diff --git a/extensions/positron-python/python_files/posit/pinned-test-requirements.txt b/extensions/positron-python/python_files/posit/pinned-test-requirements.txt index e3296784a2ee..de9a092548e1 100644 --- a/extensions/positron-python/python_files/posit/pinned-test-requirements.txt +++ b/extensions/positron-python/python_files/posit/pinned-test-requirements.txt @@ -44,6 +44,7 @@ pyarrow==21.0.0; python_version < '3.14' pytest==8.4.2 pytest-asyncio==1.2.0 pytest-mock==3.15.1 +redshift_connector==2.1.10; python_version < '3.14' syrupy==4.9.1; python_version == '3.9' syrupy==5.0.0; python_version >= '3.10' torch==2.8.0; python_version == '3.9' diff --git a/extensions/positron-python/python_files/posit/positron/connections.py b/extensions/positron-python/python_files/posit/positron/connections.py index 9931b2620468..8060c55d481e 100644 --- a/extensions/positron-python/python_files/posit/positron/connections.py +++ b/extensions/positron-python/python_files/posit/positron/connections.py @@ -301,7 +301,6 @@ def _wrap_connection(self, obj: Any) -> Connection: if not self.object_is_supported(obj): type_name = type(obj).__name__ raise UnsupportedConnectionError(f"Unsupported connection type {type_name}") - if safe_isinstance(obj, "sqlite3", "Connection"): return SQLite3Connection(obj) elif safe_isinstance(obj, "sqlalchemy", "Engine"): @@ -312,6 +311,8 @@ def _wrap_connection(self, obj: Any) -> Connection: return SnowflakeConnection(obj) elif safe_isinstance(obj, "databricks.sql.client", "Connection"): return DatabricksConnection(obj) + elif safe_isinstance(obj, "redshift_connector", "Connection"): + return RedshiftConnection(obj) else: type_name = type(obj).__name__ raise UnsupportedConnectionError(f"Unsupported connection type {type(obj)}") @@ -327,6 +328,7 @@ def object_is_supported(self, obj: Any) -> bool: or safe_isinstance(obj, "duckdb", "DuckDBPyConnection") or safe_isinstance(obj, "snowflake.connector", "SnowflakeConnection") or safe_isinstance(obj, "databricks.sql.client", "Connection") + or safe_isinstance(obj, "redshift_connector", "Connection") ) except Exception as err: logger.error(f"Error checking supported {err}") @@ -1200,3 +1202,178 @@ def _make_code(self) -> str: ")\n" "%connection_show con\n" ) + + +class RedshiftConnection(Connection): + """Support for Redshift connections to databases.""" + + def __init__(self, conn: Any): + self.conn = conn + + try: + # Unfortunatelly there's no public API to get the host, so we access the protected member. + # to at least provide some info in the connection display name. + host, _ = conn._usock.getpeername() # noqa: SLF001 + except AttributeError: + host = "" + + self.host = str(host) + + self.display_name = f"Redshift ({self.host})" + self.type = "Redshift" + self.code = self._make_code() + + self.icon = "data:image/svg+xml;base64,PD94bWwgdmVyc2lvbj0iMS4wIiBlbmNvZGluZz0iVVRGLTgiPz4KPHN2ZyB3aWR0aD0iODBweCIgaGVpZ2h0PSI4MHB4IiB2aWV3Qm94PSIwIDAgODAgODAiIHZlcnNpb249IjEuMSIgeG1sbnM9Imh0dHA6Ly93d3cudzMub3JnLzIwMDAvc3ZnIiB4bWxuczp4bGluaz0iaHR0cDovL3d3dy53My5vcmcvMTk5OS94bGluayI+CiAgICA8IS0tIEdlbmVyYXRvcjogU2tldGNoIDY0ICg5MzUzNykgLSBodHRwczovL3NrZXRjaC5jb20gLS0+CiAgICA8dGl0bGU+SWNvbi1BcmNoaXRlY3R1cmUvNjQvQXJjaF9BbWF6b24tUmVkc2hpZnRjdF82NDwvdGl0bGU+CiAgICA8ZGVzYz5DcmVhdGVkIHdpdGggU2tldGNoLjwvZGVzYz4KICAgIDxkZWZzPgogICAgICAgIDxsaW5lYXJHcmFkaWVudCB4MT0iMCUiIHkxPSIxMDAlIiB4Mj0iMTAwJSIgeTI9IjAlIiBpZD0ibGluZWFyR3JhZGllbnQtMSI+CiAgICAgICAgICAgIDxzdG9wIHN0b3AtY29sb3I9IiM0RDI3QTgiIG9mZnNldD0iMCUiPjwvc3RvcD4KICAgICAgICAgICAgPHN0b3Agc3RvcC1jb2xvcj0iI0ExNjZGRiIgb2Zmc2V0PSIxMDAlIj48L3N0b3A+CiAgICAgICAgPC9saW5lYXJHcmFkaWVudD4KICAgIDwvZGVmcz4KICAgIDxnIGlkPSJJY29uLUFyY2hpdGVjdHVyZS82NC9BcmNoX0FtYXpvbi1SZWRzaGlmdGN0XzY0IiBzdHJva2U9Im5vbmUiIHN0cm9rZS13aWR0aD0iMSIgZmlsbD0ibm9uZSIgZmlsbC1ydWxlPSJldmVub2RkIj4KICAgICAgICA8ZyBpZD0iSWNvbi1BcmNoaXRlY3R1cmUtQkcvNjQvQW5hbHl0aWNzIiBmaWxsPSJ1cmwoI2xpbmVhckdyYWRpZW50LTEpIj4KICAgICAgICAgICAgPHJlY3QgaWQ9IlJlY3RhbmdsZSIgeD0iMCIgeT0iMCIgd2lkdGg9IjgwIiBoZWlnaHQ9IjgwIj48L3JlY3Q+CiAgICAgICAgPC9nPgogICAgICAgIDxwYXRoIGQ9Ik01MC44MjUwMzU1LDM1LjE3MDQ4ODUgQzQ5Ljc2NTIwNjksMzUuMTcwNDg4NSA0OC45MDQxNTg2LDM0LjMxMTA2NjggNDguOTA0MTU4NiwzMy4yNTQyMzczIEM0OC45MDQxNTg2LDMyLjE5NzQwNzggNDkuNzY1MjA2OSwzMS4zMzY5ODkgNTAuODI1MDM1NSwzMS4zMzY5ODkgQzUxLjg4Mzg2NTIsMzEuMzM2OTg5IDUyLjc0NDkxMzUsMzIuMTk3NDA3OCA1Mi43NDQ5MTM1LDMzLjI1NDIzNzMgQzUyLjc0NDkxMzUsMzQuMzExMDY2OCA1MS44ODM4NjUyLDM1LjE3MDQ4ODUgNTAuODI1MDM1NSwzNS4xNzA0ODg1IE00NS45NTk0MTMzLDQ2LjgyNDUyNjQgQzQ0LjkwMDU4MzYsNDYuODI0NTI2NCA0NC4wMzk1MzUzLDQ1Ljk2NTEwNDcgNDQuMDM5NTM1Myw0NC45MDgyNzUyIEM0NC4wMzk1MzUzLDQzLjg1MTQ0NTcgNDQuOTAwNTgzNiw0Mi45OTIwMjM5IDQ1Ljk1OTQxMzMsNDIuOTkyMDIzOSBDNDcuMDE4MjQzLDQyLjk5MjAyMzkgNDcuODgwMjkwMiw0My44NTE0NDU3IDQ3Ljg4MDI5MDIsNDQuOTA4Mjc1MiBDNDcuODgwMjkwMiw0NS45NjUxMDQ3IDQ3LjAxODI0Myw0Ni44MjQ1MjY0IDQ1Ljk1OTQxMzMsNDYuODI0NTI2NCBNMzQuMjgyMzE5NSw0NC44ODIzNTI5IEMzMy4yMjQ0ODg3LDQ0Ljg4MjM1MjkgMzIuMzYzNDQwNCw0NC4wMjI5MzEyIDMyLjM2MzQ0MDQsNDIuOTY2MTAxNyBDMzIuMzYzNDQwNCw0MS45MDkyNzIyIDMzLjIyNDQ4ODcsNDEuMDQ4ODUzNCAzNC4yODIzMTk1LDQxLjA0ODg1MzQgQzM1LjM0MjE0ODIsNDEuMDQ4ODUzNCAzNi4yMDMxOTY1LDQxLjkwOTI3MjIgMzYuMjAzMTk2NSw0Mi45NjYxMDE3IEMzNi4yMDMxOTY1LDQ0LjAyMjkzMTIgMzUuMzQyMTQ4Miw0NC44ODIzNTI5IDM0LjI4MjMxOTUsNDQuODgyMzUyOSBNMjkuNDE3Njk2Miw1NS41NjUzMDQxIEMyOC4zNTk4NjU0LDU1LjU2NTMwNDEgMjcuNDk3ODE4Miw1NC43MDU4ODI0IDI3LjQ5NzgxODIsNTMuNjQ5MDUyOCBDMjcuNDk3ODE4Miw1Mi41OTIyMjMzIDI4LjM1OTg2NTQsNTEuNzMyODAxNiAyOS40MTc2OTYyLDUxLjczMjgwMTYgQzMwLjQ3NzUyNDgsNTEuNzMyODAxNiAzMS4zMzg1NzMyLDUyLjU5MjIyMzMgMzEuMzM4NTczMiw1My42NDkwNTI4IEMzMS4zMzg1NzMyLDU0LjcwNTg4MjQgMzAuNDc3NTI0OCw1NS41NjUzMDQxIDI5LjQxNzY5NjIsNTUuNTY1MzA0MSBNNTAuODI1MDM1NSwyOS4zNDI5NzExIEM0OC42NjQ0MjM1LDI5LjM0Mjk3MTEgNDYuOTA2MzY2NiwzMS4wOTc3MDY5IDQ2LjkwNjM2NjYsMzMuMjU0MjM3MyBDNDYuOTA2MzY2NiwzNC41NzYyNzEyIDQ3LjU3MDYzMjUsMzUuNzQxNzc0NyA0OC41ODA1MTYzLDM2LjQ1MDY0ODEgTDQ2Ljc1MjUzNjcsNDEuMDc4NzYzNyBDNDYuNDk1ODIwNCw0MS4wMjU5MjIyIDQ2LjIzMDExNDEsNDAuOTk4MDA2IDQ1Ljk1OTQxMzMsNDAuOTk4MDA2IEM0NC4yNDEzMTIyLDQwLjk5ODAwNiA0Mi43OTM5MTIsNDIuMTE0NjU2IDQyLjI2NzQ5MzgsNDMuNjU1MDM0OSBMMzguMTc3MDE0OSw0Mi43MjY4MTk1IEMzOC4wNTAxNTUxLDQwLjY4Mzk0ODIgMzYuMzYyMDIwOSwzOS4wNTQ4MzU1IDM0LjI4MjMxOTUsMzkuMDU0ODM1NSBDMzIuMTIyNzA2NSwzOS4wNTQ4MzU1IDMwLjM2NTY0ODUsNDAuODA5NTcxMyAzMC4zNjU2NDg1LDQyLjk2NjEwMTcgQzMwLjM2NTY0ODUsNDMuOTc0MDc3OCAzMC43NTkyMTM1LDQ0Ljg4NTM0NCAzMS4zODk1MTY4LDQ1LjU3OTI2MjIgTDI5LjYwNDQ4OTgsNDkuNzU3NzI2OCBDMjkuNTQwNTYwNCw0OS43NTM3Mzg4IDI5LjQ4MTYyNTYsNDkuNzM4NzgzNiAyOS40MTc2OTYyLDQ5LjczODc4MzYgQzI3LjI1ODA4MzIsNDkuNzM4NzgzNiAyNS41MDAwMjYzLDUxLjQ5MjUyMjQgMjUuNTAwMDI2Myw1My42NDkwNTI4IEMyNS41MDAwMjYzLDU1LjgwNDU4NjIgMjcuMjU4MDgzMiw1Ny41NTkzMjIgMjkuNDE3Njk2Miw1Ny41NTkzMjIgQzMxLjU3ODMwODIsNTcuNTU5MzIyIDMzLjMzNjM2NTEsNTUuODA0NTg2MiAzMy4zMzYzNjUxLDUzLjY0OTA1MjggQzMzLjMzNjM2NTEsNTIuMjY1MjA0NCAzMi42MDkxNjg4LDUxLjA1NDgzNTUgMzEuNTE5MzczMyw1MC4zNTg5MjMyIEwzMy4wOTQ2MzIyLDQ2LjY3Mjk4MTEgQzMzLjQ3MjIxNDksNDYuNzkzNjE5MSAzMy44NjU3Nzk5LDQ2Ljg3NjM3MDkgMzQuMjgyMzE5NSw0Ni44NzYzNzA5IEMzNS44MjM2MTYsNDYuODc2MzcwOSAzNy4xNDcxNTMxLDQ1Ljk3NjA3MTggMzcuNzg1NDQ3Nyw0NC42ODI5NTExIEw0Mi4xMTg2NTgzLDQ1LjY2NjAwMiBDNDIuNDczMjY2NCw0Ny40NjA2MTgxIDQ0LjA1OTUxMzIsNDguODE4NTQ0NCA0NS45NTk0MTMzLDQ4LjgxODU0NDQgQzQ4LjEyMDAyNTIsNDguODE4NTQ0NCA0OS44NzgwODIxLDQ3LjA2NDgwNTYgNDkuODc4MDgyMSw0NC45MDgyNzUyIEM0OS44NzgwODIxLDQzLjc0Mjc3MTcgNDkuMzUzNjYxNyw0Mi43MDY4Nzk0IDQ4LjU0MDU2MDQsNDEuOTg5MDMyOSBMNTAuNDYxNDM3NCwzNy4xMjc2MTcxIEM1MC41ODMzMDI3LDM3LjEzOTU4MTMgNTAuNzAwMTczNSwzNy4xNjQ1MDY1IDUwLjgyNTAzNTUsMzcuMTY0NTA2NSBDNTIuOTg0NjQ4NSwzNy4xNjQ1MDY1IDU0Ljc0MjcwNTQsMzUuNDA5NzcwNyA1NC43NDI3MDU0LDMzLjI1NDIzNzMgQzU0Ljc0MjcwNTQsMzEuMDk3NzA2OSA1Mi45ODQ2NDg1LDI5LjM0Mjk3MTEgNTAuODI1MDM1NSwyOS4zNDI5NzExIE00MCw2Ni4wMDU5ODIxIEMzMC4yNjg3NTU2LDY2LjAwNTk4MjEgMjIuOTk3NzkxOSw2My4wODQ3NDU4IDIyLjk5Nzc5MTksNjAuNDcyNTgyMyBMMjIuOTk3NzkxOSwyMy4xNTE1NDU0IEMyNi4zMDgxMzMxLDI1Ljg0MTQ3NTYgMzMuMzE1Mzg4MywyNy4yNjMyMTA0IDQwLjAyMDk3NjgsMjcuMjYzMjEwNCBDNDYuNjk2NTk4NSwyNy4yNjMyMTA0IDUzLjY3Mzg4NjgsMjUuODUzNDM5NyA1Ny4wMDIyMDgxLDIzLjE4NjQ0MDcgTDU3LjAwMjIwODEsNjAuNDcyNTgyMyBDNTcuMDAyMjA4MSw2My4wODQ3NDU4IDQ5LjczMDI0NTUsNjYuMDA1OTgyMSA0MCw2Ni4wMDU5ODIxIE00MC4wMjA5NzY4LDEzLjk5NDAxNzkgQzUwLjAyNzkxNjUsMTMuOTk0MDE3OSA1Ny4wMDIyMDgxLDE2Ljk2NTEwNDcgNTcuMDAyMjA4MSwxOS42MzIxMDM3IEM1Ny4wMDIyMDgxLDIyLjI5ODEwNTcgNTAuMDI3OTE2NSwyNS4yNjkxOTI0IDQwLjAyMDk3NjgsMjUuMjY5MTkyNCBDMzAuMDEzMDM4MiwyNS4yNjkxOTI0IDIzLjAzOTc0NTUsMjIuMjk4MTA1NyAyMy4wMzk3NDU1LDE5LjYzMjEwMzcgQzIzLjAzOTc0NTUsMTYuOTY1MTA0NyAzMC4wMTMwMzgyLDEzLjk5NDAxNzkgNDAuMDIwOTc2OCwxMy45OTQwMTc5IE01OSwxOS42MzIxMDM3IEM1OSwxNC42NzQ5NzUxIDQ5LjIyMTgwNzUsMTIgNDAuMDIwOTc2OCwxMiBDMzAuODIwMTQ2MiwxMiAyMS4wNDE5NTM2LDE0LjY3NDk3NTEgMjEuMDQxOTUzNiwxOS42MzIxMDM3IEMyMS4wNDE5NTM2LDE5LjY0MDA3OTggMjEuMDQzOTUxNCwxOS42NDkwNTI4IDIxLjA0Mzk1MTQsMTkuNjU3MDI4OSBMMjEsMTkuNjU3MDI4OSBMMjEsNjAuNDcyNTgyMyBDMjEsNjUuMzYxOTE0MyAzMC43ODkxODA0LDY4IDQwLDY4IEM0OS4yMTA4MTk2LDY4IDU5LDY1LjM2MTkxNDMgNTksNjAuNDcyNTgyMyBMNTksMTkuNjU3MDI4OSBMNTguOTk4MDAyMiwxOS42NTcwMjg5IEM1OC45OTgwMDIyLDE5LjY0OTA1MjggNTksMTkuNjQwMDc5OCA1OSwxOS42MzIxMDM3IiBpZD0iQW1hem9uLVJlZHNoaWZ0X0ljb25fNjRfU3F1aWQiIGZpbGw9IiNGRkZGRkYiPjwvcGF0aD4KICAgIDwvZz4KPC9zdmc+" + + def disconnect(self): + with contextlib.suppress(Exception): + self.conn.close() + + def list_object_types(self): + return { + "database": ConnectionObjectInfo({"contains": None, "icon": None}), + "schema": ConnectionObjectInfo({"contains": None, "icon": None}), + "table": ConnectionObjectInfo({"contains": "data", "icon": None}), + "view": ConnectionObjectInfo({"contains": "data", "icon": None}), + } + + def list_objects(self, path: list[ObjectSchema]): + if len(path) == 0: + rows = self._query("SHOW DATABASES;") + return [ + ConnectionObject({"name": row["database_name"], "kind": "database"}) for row in rows + ] + + if len(path) == 1: + database = path[0] + if database.kind != "database": + raise ValueError("Expected database on path position 0.", f"Path: {path}") + database_ident = self._qualify(database.name) + rows = self._query(f"SHOW SCHEMAS FROM DATABASE {database_ident};") + return [ + ConnectionObject( + { + "name": row["schema_name"], + "kind": "schema", + } + ) + for row in rows + ] + + if len(path) == 2: + database, schema = path + if database.kind != "database" or schema.kind != "schema": + raise ValueError( + "Expected database and schema objects at positions 0 and 1.", f"Path: {path}" + ) + location = f"{self._qualify(database.name)}.{self._qualify(schema.name)}" + tables = self._query(f"SHOW TABLES FROM SCHEMA {location};") + return [ + ConnectionObject( + { + "name": row["table_name"], + "kind": row["table_type"].lower(), + } + ) + for row in tables + ] + + raise ValueError(f"Path length must be at most 2, but got {len(path)}. Path: {path}") + + def list_fields(self, path: list[ObjectSchema]): + if len(path) != 3: + raise ValueError(f"Path length must be 3, but got {len(path)}. Path: {path}") + + database, schema, table = path + if ( + database.kind != "database" + or schema.kind != "schema" + or table.kind not in ("table", "view") + ): + raise ValueError( + "Expected database, schema, and table/view kinds in the path.", + f"Path: {path}", + ) + + identifier = ".".join( + [self._qualify(database.name), self._qualify(schema.name), self._qualify(table.name)] + ) + rows = self._query(f"SHOW COLUMNS FROM TABLE {identifier};") + return [ + ConnectionObjectFields( + { + "name": row["column_name"], + "dtype": row["data_type"], + } + ) + for row in rows + ] + + def preview_object(self, path: list[ObjectSchema], var_name: str | None = None): + if len(path) != 3: + raise ValueError(f"Path length must be 3, but got {len(path)}. Path: {path}") + + database, schema, table = path + if ( + database.kind != "database" + or schema.kind != "schema" + or table.kind not in ("table", "view") + ): + raise ValueError( + "Expected database, schema, and table/view kinds in the path.", + f"Path: {path}", + ) + + identifier = ".".join( + [self._qualify(database.name), self._qualify(schema.name), self._qualify(table.name)] + ) + sql = f"SELECT * FROM {identifier} LIMIT 1000;" + + with self.conn.cursor() as cursor: + try: + cursor.execute(sql) + frame = cursor.fetch_dataframe() + except Exception as e: + # Rollback on error to avoid transaction issues + # for subsequent queries + self.conn.rollback() + raise e + + var_name = var_name or "conn" + return frame, ( + f"with {var_name}.cursor() as cursor:\n" + f" cursor.execute({sql!r})\n" + f" {table.name} = cursor.fetch_dataframe()" + ) + + def _query(self, sql: str) -> list[dict[str, Any]]: + cursor = self.conn.cursor() + try: + cursor.execute(sql) + rows = cursor.fetchall() + description = cursor.description or [] + columns = [col[0] for col in description] + return [dict(zip(columns, row)) for row in rows] + except Exception as e: + # Rollback on error to avoid transaction issues + # for subsequent queries + self.conn.rollback() + raise e + finally: + cursor.close() + + def _qualify(self, identifier: str) -> str: + escaped = identifier.replace('"', '""') + return f'"{escaped}"' + + def _make_code(self) -> str: + return ( + "# Requires redshift-connector package\n" + "# Authentication steps may be incomplete, adjust as needed.\n" + "import redshift_connector\n" + "con = redshift_connector.connect(\n" + f" iam = True,\n" + f" host = '{self.host}',\n" + ")\n" + "%connection_show con\n" + ) diff --git a/extensions/positron-python/python_files/posit/positron/tests/test_connections.py b/extensions/positron-python/python_files/posit/positron/tests/test_connections.py index 37ab01e5975e..e2ec028c0de9 100644 --- a/extensions/positron-python/python_files/posit/positron/tests/test_connections.py +++ b/extensions/positron-python/python_files/posit/positron/tests/test_connections.py @@ -43,6 +43,14 @@ except ImportError: HAS_DATABRICKS = False + +try: + import redshift_connector + + HAS_REDSHIFT = "REDSHIFT_HOST" in os.environ +except ImportError: + HAS_REDSHIFT = False + from positron.access_keys import encode_access_key from positron.connections import ConnectionsService @@ -706,3 +714,135 @@ def _view_in_connections_pane(self, variables_comm: DummyComm, path): assert variables_comm.messages == [json_rpc_response({})] variables_comm.messages.clear() return tuple(encoded_paths) + + +@pytest.mark.skipif(not HAS_REDSHIFT, reason="Redshift not available") +class TestRedshiftConnectionsService: + REDSHIFT_HOST = os.environ.get("REDSHIFT_HOST") + REDSHIFT_PROFILE = os.environ.get("REDSHIFT_PROFILE", "default") + REDSHIFT_DATABASE = "dev" + REDSHIFT_SCHEMA = "public" + REDSHIFT_TABLE = "airlines" + + def _connect(self): + return redshift_connector.connect( + iam=True, + host=self.REDSHIFT_HOST, + database=self.REDSHIFT_DATABASE, + profile=self.REDSHIFT_PROFILE, + ) + + def _open_comm(self, connections_service: ConnectionsService): + con = self._connect() + comm_id = connections_service.register_connection(con) + dummy_comm = DummyComm(TARGET_NAME, comm_id=comm_id) + connections_service.on_comm_open(dummy_comm) + dummy_comm.messages.clear() + return dummy_comm, comm_id + + def _database_path(self): + return [{"kind": "database", "name": self.REDSHIFT_DATABASE}] + + def _schema_path(self): + return [*self._database_path(), {"kind": "schema", "name": self.REDSHIFT_SCHEMA}] + + def _table_path(self): + return [*self._schema_path(), {"kind": "table", "name": self.REDSHIFT_TABLE}] + + def _resolve_path(self, kind: str): + if kind == "root": + return [] + if kind == "database": + return self._database_path() + if kind == "schema": + return self._schema_path() + if kind == "table": + return self._table_path() + raise ValueError(f"Unknown path kind: {kind}") + + def test_register_connection(self, connections_service: ConnectionsService): + con = self._connect() + comm_id = connections_service.register_connection(con) + assert comm_id in connections_service.comms + + @pytest.mark.parametrize( + ("path_kind"), + [ + pytest.param("root", id="root"), + pytest.param("database", id="database"), + pytest.param("schema", id="schema"), + pytest.param("table", id="table"), + ], + ) + def test_contains_data(self, connections_service: ConnectionsService, path_kind: str): + dummy_comm, comm_id = self._open_comm(connections_service) + path = self._resolve_path(path_kind) + + msg = _make_msg(params={"path": path}, method="contains_data", comm_id=comm_id) + dummy_comm.handle_msg(msg) + result = dummy_comm.messages[0]["data"]["result"] + assert result == (path_kind == "table") + + @pytest.mark.parametrize( + ("path_kind", "expected"), + [ + pytest.param("root", "data:image", id="root"), + pytest.param("database", "", id="database"), + pytest.param("schema", "", id="schema"), + pytest.param("table", "", id="table"), + ], + ) + def test_get_icon(self, connections_service: ConnectionsService, path_kind: str, expected: str): + dummy_comm, comm_id = self._open_comm(connections_service) + path = self._resolve_path(path_kind) + + msg = _make_msg(params={"path": path}, method="get_icon", comm_id=comm_id) + dummy_comm.handle_msg(msg) + result = dummy_comm.messages[0]["data"]["result"] + if expected: + assert expected in result + else: + assert result == "" + + @pytest.mark.parametrize( + "path_kind", + [ + pytest.param("root", id="databases"), + pytest.param("database", id="schemas"), + pytest.param("schema", id="tables"), + ], + ) + def test_list_objects(self, connections_service: ConnectionsService, path_kind: str): + dummy_comm, comm_id = self._open_comm(connections_service) + path = self._resolve_path(path_kind) + expected = { + "root": self.REDSHIFT_DATABASE, + "database": self.REDSHIFT_SCHEMA, + "schema": self.REDSHIFT_TABLE, + }[path_kind] + + msg = _make_msg(params={"path": path}, method="list_objects", comm_id=comm_id) + dummy_comm.handle_msg(msg) + result = dummy_comm.messages[0]["data"]["result"] + names = [item["name"] for item in result] + assert expected in names + + def test_list_fields(self, connections_service: ConnectionsService): + dummy_comm, comm_id = self._open_comm(connections_service) + path = self._table_path() + + msg = _make_msg(params={"path": path}, method="list_fields", comm_id=comm_id) + dummy_comm.handle_msg(msg) + result = dummy_comm.messages[0]["data"]["result"] + field_names = {field["name"].lower() for field in result} + assert {"carrier", "name"}.issubset(field_names) + + def test_preview_object(self, connections_service: ConnectionsService): + dummy_comm, comm_id = self._open_comm(connections_service) + path = self._table_path() + + msg = _make_msg(params={"path": path}, method="preview_object", comm_id=comm_id) + dummy_comm.handle_msg(msg) + connections_service._kernel.data_explorer_service.shutdown() # noqa: SLF001 + result = dummy_comm.messages[0]["data"]["result"] + assert result is None diff --git a/extensions/positron-python/python_files/posit/test-requirements.txt b/extensions/positron-python/python_files/posit/test-requirements.txt index ce3a6a124995..82d225c0c2ed 100644 --- a/extensions/positron-python/python_files/posit/test-requirements.txt +++ b/extensions/positron-python/python_files/posit/test-requirements.txt @@ -23,6 +23,7 @@ pytest pytest-asyncio pytest-mock syrupy +redshift_connector torch scipy snowflake-connector-python; python_version < '3.14'