Merging upstream version 0.3.0.

Signed-off-by: Daniel Baumann <daniel@debian.org>

.gitignore

@@ -1,3 +1,6 @@
+# test data
+sqlserver_data/
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

CHANGELOG.md

@@ -4,6 +4,12 @@ All notable changes to this project will be documented in this file.
 
 ## [Unreleased]
 
+## [0.3.0] - 2025-02-25
+
+-   The Data Catalog now displays all databases on the connected server, not just the currently-connected database ([tconbeer/harlequin#415](https://github.com/tconbeer/harlequin/discussions/415)).
+-   Columns in the Data Catalog are now fetched lazily ([#12](https://github.com/tconbeer/harlequin-odbc/issues/12), [#13](https://github.com/tconbeer/harlequin-odbc/issues/13)).
+-   Data Catalog items now support basic interactions ([#14](https://github.com/tconbeer/harlequin-odbc/issues/14)).
+
 ## [0.2.0] - 2025-01-08
 
 -   Drops support for Python 3.8
@@ -22,7 +28,9 @@ All notable changes to this project will be documented in this file.
 
 -   Adds a basic ODBC adapter.
 
-[Unreleased]: https://github.com/tconbeer/harlequin-odbc/compare/0.2.0...HEAD
+[Unreleased]: https://github.com/tconbeer/harlequin-odbc/compare/0.3.0...HEAD
+
+[0.3.0]: https://github.com/tconbeer/harlequin-odbc/compare/0.2.0...0.3.0
 
 [0.2.0]: https://github.com/tconbeer/harlequin-odbc/compare/0.1.1...0.2.0
 

docker-compose.yml

@@ -7,9 +7,8 @@ services:
       ACCEPT_EULA: Y
       MSSQL_SA_PASSWORD: for-testing
     volumes:
-      - sqlserver_data:/var/opt/mssql
+      - ./sqlserver_data/data:/var/opt/mssql/data
+      - ./sqlserver_data/log:/var/opt/mssql/log
+      - ./sqlserver_data/secrets:/var/opt/mssql/secrets
     ports:
         - 1433:1433
- 
-volumes:
-  sqlserver_data:

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "harlequin-odbc"
-version = "0.2.0"
+version = "0.3.0"
 description = "A Harlequin adapter for ODBC drivers."
 authors = ["Ted Conbeer <tconbeer@users.noreply.github.com>"]
 license = "MIT"
@@ -14,7 +14,7 @@ odbc = "harlequin_odbc:HarlequinOdbcAdapter"
 
 [tool.poetry.dependencies]
 python = ">=3.9,<3.14"
-harlequin = ">=1.9.1,<3"
+harlequin = ">=1.25,<3"
 pyodbc = "^5.0"
 
 [tool.poetry.group.dev.dependencies]
@@ -30,13 +30,13 @@ build-backend = "poetry.core.masonry.api"
 
 
 [tool.ruff]
-target-version = "py38"
+target-version = "py39"
 
 [tool.ruff.lint]
 select = ["A", "B", "E", "F", "I"]
 
 [tool.mypy]
-python_version = "3.8"
+python_version = "3.9"
 files = [
     "src/**/*.py",
     "tests/**/*.py",

src/harlequin_odbc/adapter.py

@@ -1,8 +1,9 @@
 from __future__ import annotations
 
-from typing import Any, Sequence
+from contextlib import suppress
+from typing import TYPE_CHECKING, Any, Sequence
 
-import pyodbc  # type: ignore
+import pyodbc
 from harlequin import (
     HarlequinAdapter,
     HarlequinConnection,
@@ -17,8 +18,16 @@ from harlequin.exception import (
 )
 from textual_fastdatatable.backend import AutoBackendType
 
+from harlequin_odbc.catalog import (
+    DatabaseCatalogItem,
+    RelationCatalogItem,
+    SchemaCatalogItem,
+)
 from harlequin_odbc.cli_options import ODBC_OPTIONS
 
+if TYPE_CHECKING:
+    pass
+
 
 class HarlequinOdbcCursor(HarlequinCursor):
     def __init__(self, cur: pyodbc.Cursor) -> None:
@@ -86,7 +95,7 @@ class HarlequinOdbcConnection(HarlequinConnection):
             cur.execute(query)
         except Exception as e:
             raise HarlequinQueryError(
-                msg=str(e),
+                msg=f"{e.__class__.__name__}: {e}",
                 title="Harlequin encountered an error while executing your query.",
             ) from e
         else:
@@ -103,72 +112,48 @@ class HarlequinOdbcConnection(HarlequinConnection):
             for schema, relations in schemas.items():
                 rel_items: list[CatalogItem] = []
                 for rel, rel_type in relations:
-                    cols = self._list_columns_in_relation(
-                        catalog_name=db, schema_name=schema, rel_name=rel
-                    )
-                    col_items = [
-                        CatalogItem(
-                            qualified_identifier=f'"{db}"."{schema}"."{rel}"."{col}"',
-                            query_name=f'"{col}"',
-                            label=col,
-                            type_label=col_type,
-                        )
-                        for col, col_type in cols
-                    ]
                     rel_items.append(
-                        CatalogItem(
+                        RelationCatalogItem.from_label(
-                            qualified_identifier=f'"{db}"."{schema}"."{rel}"',
-                            query_name=f'"{db}"."{schema}"."{rel}"',
                             label=rel,
-                            type_label=rel_type,
+                            schema_label=schema,
-                            children=col_items,
+                            db_label=db,
+                            rel_type=rel_type,
+                            connection=self,
                         )
                     )
                 schema_items.append(
-                    CatalogItem(
+                    SchemaCatalogItem.from_label(
-                        qualified_identifier=f'"{db}"."{schema}"',
-                        query_name=f'"{db}"."{schema}"',
                         label=schema,
-                        type_label="s",
+                        db_label=db,
+                        connection=self,
                         children=rel_items,
                     )
                 )
             db_items.append(
-                CatalogItem(
+                DatabaseCatalogItem.from_label(
-                    qualified_identifier=f'"{db}"',
-                    query_name=f'"{db}"',
                     label=db,
-                    type_label="db",
+                    connection=self,
                     children=schema_items,
                 )
             )
         return Catalog(items=db_items)
 
+    def close(self) -> None:
+        with suppress(Exception):
+            self.conn.close()
+        with suppress(Exception):
+            self.aux_conn.close()
+
     def _list_tables(self) -> dict[str, dict[str, list[tuple[str, str]]]]:
         cur = self.aux_conn.cursor()
-        rel_type_map = {
-            "TABLE": "t",
-            "VIEW": "v",
-            "SYSTEM TABLE": "st",
-            "GLOBAL TEMPORARY": "tmp",
-            "LOCAL TEMPORARY": "tmp",
-        }
         catalog: dict[str, dict[str, list[tuple[str, str]]]] = {}
-        for db_name, schema_name, rel_name, rel_type, *_ in cur.tables():
+        for db_name, schema_name, rel_name, rel_type, *_ in cur.tables(catalog="%"):
             if db_name not in catalog:
-                catalog[db_name] = {
+                catalog[db_name] = {schema_name: [(rel_name, rel_type)]}
-                    schema_name: [
-                        (rel_name, rel_type_map.get(rel_type, str(rel_type).lower()))
-                    ]
-                }
             elif schema_name not in catalog[db_name]:
-                catalog[db_name][schema_name] = [
+                catalog[db_name][schema_name] = [(rel_name, rel_type)]
-                    (rel_name, rel_type_map.get(rel_type, rel_type))
-                ]
             else:
-                catalog[db_name][schema_name].append(
+                catalog[db_name][schema_name].append((rel_name, rel_type))
-                    (rel_name, rel_type_map.get(rel_type, rel_type))
-                )
         return catalog
 
     def _list_columns_in_relation(

src/harlequin_odbc/catalog.py

@@ -0,0 +1,180 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, ClassVar
+
+from harlequin.catalog import CatalogItem, InteractiveCatalogItem
+
+from harlequin_odbc.interactions import (
+    execute_drop_database_statement,
+    execute_drop_table_statement,
+    execute_drop_view_statement,
+    execute_use_statement,
+    insert_columns_at_cursor,
+    show_select_star,
+)
+
+if TYPE_CHECKING:
+    from harlequin_odbc.adapter import HarlequinOdbcConnection
+
+
+@dataclass
+class ColumnCatalogItem(InteractiveCatalogItem["HarlequinOdbcConnection"]):
+    parent: "RelationCatalogItem" | None = None
+
+    @classmethod
+    def from_parent(
+        cls,
+        parent: "RelationCatalogItem",
+        label: str,
+        type_label: str,
+    ) -> "ColumnCatalogItem":
+        column_qualified_identifier = f'{parent.qualified_identifier}."{label}"'
+        column_query_name = f'"{label}"'
+        return cls(
+            qualified_identifier=column_qualified_identifier,
+            query_name=column_query_name,
+            label=label,
+            type_label=type_label,
+            connection=parent.connection,
+            parent=parent,
+            loaded=True,
+        )
+
+
+@dataclass
+class RelationCatalogItem(InteractiveCatalogItem["HarlequinOdbcConnection"]):
+    INTERACTIONS = [
+        ("Insert Columns at Cursor", insert_columns_at_cursor),
+        ("Preview Data", show_select_star),
+    ]
+    TYPE_LABEL: ClassVar[str] = ""
+    schema_label: str = ""
+    db_label: str = ""
+
+    @classmethod
+    def from_label(
+        cls,
+        label: str,
+        schema_label: str,
+        db_label: str,
+        rel_type: str,
+        connection: "HarlequinOdbcConnection",
+    ) -> "RelationCatalogItem":
+        rel_type_map: dict[str, type[RelationCatalogItem]] = {
+            "TABLE": TableCatalogItem,
+            "VIEW": ViewCatalogItem,
+            "SYSTEM TABLE": SystemTableCatalogItem,
+            "GLOBAL TEMPORARY": GlobalTempTableCatalogItem,
+            "LOCAL TEMPORARY": LocalTempTableCatalogItem,
+        }
+
+        item_class = rel_type_map.get(rel_type, TableCatalogItem)
+        return item_class(
+            qualified_identifier=f'"{db_label}"."{schema_label}"."{label}"',
+            query_name=f'"{schema_label}"."{label}"',
+            label=label,
+            schema_label=schema_label,
+            db_label=db_label,
+            type_label=item_class.TYPE_LABEL,
+            connection=connection,
+        )
+
+    def fetch_children(self) -> list[ColumnCatalogItem]:
+        if self.connection is None:
+            return []
+        cols = self.connection._list_columns_in_relation(
+            catalog_name=self.db_label,
+            schema_name=self.schema_label,
+            rel_name=self.label,
+        )
+        return [
+            ColumnCatalogItem.from_parent(
+                parent=self, label=col_label, type_label=col_type_label
+            )
+            for col_label, col_type_label in cols
+        ]
+
+
+class ViewCatalogItem(RelationCatalogItem):
+    INTERACTIONS = RelationCatalogItem.INTERACTIONS + [
+        ("Drop View", execute_drop_view_statement),
+    ]
+    TYPE_LABEL: ClassVar[str] = "v"
+
+
+class TableCatalogItem(RelationCatalogItem):
+    INTERACTIONS = RelationCatalogItem.INTERACTIONS + [
+        ("Drop Table", execute_drop_table_statement),
+    ]
+    TYPE_LABEL: ClassVar[str] = "t"
+
+
+class SystemTableCatalogItem(RelationCatalogItem):
+    TYPE_LABEL: ClassVar[str] = "st"
+
+
+class TempTableCatalogItem(TableCatalogItem):
+    TYPE_LABEL: ClassVar[str] = "tmp"
+
+
+class GlobalTempTableCatalogItem(TempTableCatalogItem):
+    pass
+
+
+class LocalTempTableCatalogItem(TempTableCatalogItem):
+    pass
+
+
+@dataclass
+class SchemaCatalogItem(InteractiveCatalogItem["HarlequinOdbcConnection"]):
+    db_label: str = ""
+
+    @classmethod
+    def from_label(
+        cls,
+        label: str,
+        db_label: str,
+        connection: "HarlequinOdbcConnection",
+        children: list[CatalogItem] | None = None,
+    ) -> "SchemaCatalogItem":
+        schema_identifier = f'"{label}"'
+        if children is None:
+            children = []
+        return cls(
+            qualified_identifier=f'"{db_label}".{schema_identifier}',
+            query_name=schema_identifier,
+            label=label,
+            db_label=db_label,
+            type_label="sch",
+            connection=connection,
+            children=children,
+            loaded=True,
+        )
+
+
+class DatabaseCatalogItem(InteractiveCatalogItem["HarlequinOdbcConnection"]):
+    INTERACTIONS = [
+        ("Use Database", execute_use_statement),
+        ("Drop Database", execute_drop_database_statement),
+    ]
+
+    @classmethod
+    def from_label(
+        cls,
+        label: str,
+        connection: "HarlequinOdbcConnection",
+        children: list[CatalogItem] | None = None,
+    ) -> "DatabaseCatalogItem":
+        database_identifier = f'"{label}"'
+        if children is None:
+            children = []
+        return cls(
+            qualified_identifier=database_identifier,
+            query_name=database_identifier,
+            label=label,
+            type_label="db",
+            connection=connection,
+            children=children,
+            loaded=True,
+        )

src/harlequin_odbc/interactions.py

@@ -0,0 +1,120 @@
+from __future__ import annotations
+
+from textwrap import dedent
+from typing import TYPE_CHECKING, Literal, Sequence
+
+from harlequin.catalog import CatalogItem
+from harlequin.exception import HarlequinQueryError
+
+if TYPE_CHECKING:
+    from harlequin.driver import HarlequinDriver
+
+    from harlequin_odbc.catalog import (
+        ColumnCatalogItem,
+        DatabaseCatalogItem,
+        RelationCatalogItem,
+    )
+
+
+def execute_use_statement(
+    item: "DatabaseCatalogItem",
+    driver: "HarlequinDriver",
+) -> None:
+    if item.connection is None:
+        return
+    try:
+        item.connection.execute(f"use {item.query_name}")
+    except HarlequinQueryError:
+        driver.notify("Could not switch context", severity="error")
+        raise
+    else:
+        driver.notify(f"Editor context switched to {item.label}")
+
+
+def execute_drop_database_statement(
+    item: "DatabaseCatalogItem",
+    driver: "HarlequinDriver",
+) -> None:
+    def _drop_database() -> None:
+        if item.connection is None:
+            return
+        try:
+            item.connection.execute(f"drop database {item.query_name}")
+        except HarlequinQueryError:
+            driver.notify(f"Could not drop database {item.label}", severity="error")
+            raise
+        else:
+            driver.notify(f"Dropped database {item.label}")
+            driver.refresh_catalog()
+
+    if item.children or item.fetch_children():
+        driver.confirm_and_execute(callback=_drop_database)
+    else:
+        _drop_database()
+
+
+def execute_drop_relation_statement(
+    item: "RelationCatalogItem",
+    driver: "HarlequinDriver",
+    relation_type: Literal["view", "table", "foreign table"],
+) -> None:
+    def _drop_relation() -> None:
+        if item.connection is None:
+            return
+        try:
+            item.connection.execute(f"drop {relation_type} {item.query_name}")
+        except HarlequinQueryError:
+            driver.notify(
+                f"Could not drop {relation_type} {item.label}", severity="error"
+            )
+            raise
+        else:
+            driver.notify(f"Dropped {relation_type} {item.label}")
+            driver.refresh_catalog()
+
+    driver.confirm_and_execute(callback=_drop_relation)
+
+
+def execute_drop_table_statement(
+    item: "RelationCatalogItem", driver: "HarlequinDriver"
+) -> None:
+    execute_drop_relation_statement(item=item, driver=driver, relation_type="table")
+
+
+def execute_drop_foreign_table_statement(
+    item: "RelationCatalogItem", driver: "HarlequinDriver"
+) -> None:
+    execute_drop_relation_statement(
+        item=item, driver=driver, relation_type="foreign table"
+    )
+
+
+def execute_drop_view_statement(
+    item: "RelationCatalogItem", driver: "HarlequinDriver"
+) -> None:
+    execute_drop_relation_statement(item=item, driver=driver, relation_type="view")
+
+
+def show_select_star(
+    item: "RelationCatalogItem",
+    driver: "HarlequinDriver",
+) -> None:
+    driver.insert_text_in_new_buffer(
+        dedent(
+            f"""
+            select *
+            from {item.query_name}
+            """.strip("\n")
+        )
+    )
+
+
+def insert_columns_at_cursor(
+    item: "RelationCatalogItem",
+    driver: "HarlequinDriver",
+) -> None:
+    if item.loaded:
+        cols: Sequence["CatalogItem" | "ColumnCatalogItem"] = item.children
+    else:
+        cols = item.fetch_children()
+    driver.insert_text_at_selection(text=",\n".join(c.query_name for c in cols))

tests/conftest.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+from typing import Generator
+
+import pyodbc
+import pytest
+
+from harlequin_odbc.adapter import (
+    HarlequinOdbcAdapter,
+    HarlequinOdbcConnection,
+)
+
+MASTER_DB_CONN = "Driver={ODBC Driver 18 for SQL Server};Server=tcp:localhost,1433;Database=master;Uid=sa;Pwd={for-testing};Encrypt=yes;TrustServerCertificate=yes;Connection Timeout=5;"  # noqa: E501
+TEST_DB_CONN = "Driver={ODBC Driver 18 for SQL Server};Server=tcp:localhost,1433;Database=test;Uid=sa;Pwd={for-testing};Encrypt=yes;TrustServerCertificate=yes;Connection Timeout=5;"  # noqa: E501
+
+
+@pytest.fixture
+def connection() -> Generator[HarlequinOdbcConnection, None, None]:
+    master_conn = pyodbc.connect(MASTER_DB_CONN, autocommit=True)
+    cur = master_conn.cursor()
+    cur.execute("drop database if exists test;")
+    cur.execute("create database test;")
+    cur.close()
+    master_conn.close()
+    conn = HarlequinOdbcAdapter(conn_str=(TEST_DB_CONN,)).connect()
+    yield conn
+    conn.close()

tests/test_catalog.py

@@ -0,0 +1,99 @@
+from typing import Generator
+
+import pytest
+from harlequin.catalog import InteractiveCatalogItem
+
+from harlequin_odbc.adapter import HarlequinOdbcConnection
+from harlequin_odbc.catalog import (
+    ColumnCatalogItem,
+    DatabaseCatalogItem,
+    RelationCatalogItem,
+    SchemaCatalogItem,
+    TableCatalogItem,
+    ViewCatalogItem,
+)
+
+
+@pytest.fixture
+def connection_with_objects(
+    connection: HarlequinOdbcConnection,
+) -> Generator[HarlequinOdbcConnection, None, None]:
+    connection.execute("create schema one")
+    connection.execute("select 1 as a, '2' as b into one.foo")
+    connection.execute("select 1 as a, '2' as b into one.bar")
+    connection.execute("select 1 as a, '2' as b into one.baz")
+    connection.execute("create schema two")
+    connection.execute("create view two.qux as select * from one.foo")
+    connection.execute("create schema three")
+
+    yield connection
+
+    connection.execute("drop table one.foo")
+    connection.execute("drop table one.bar")
+    connection.execute("drop table one.baz")
+    connection.execute("drop schema one")
+    connection.execute("drop view two.qux")
+    connection.execute("drop schema two")
+    connection.execute("drop schema three")
+
+
+def test_catalog(connection_with_objects: HarlequinOdbcConnection) -> None:
+    conn = connection_with_objects
+
+    catalog = conn.get_catalog()
+
+    # at least two databases, postgres and test
+    assert len(catalog.items) >= 2
+
+    [test_db_item] = filter(lambda item: item.label == "test", catalog.items)
+    assert isinstance(test_db_item, InteractiveCatalogItem)
+    assert isinstance(test_db_item, DatabaseCatalogItem)
+    assert test_db_item.children
+    assert test_db_item.loaded
+
+    schema_items = test_db_item.children
+    assert all(isinstance(item, SchemaCatalogItem) for item in schema_items)
+
+    [schema_one_item] = filter(lambda item: item.label == "one", schema_items)
+    assert isinstance(schema_one_item, SchemaCatalogItem)
+    assert schema_one_item.children
+    assert schema_one_item.loaded
+
+    table_items = schema_one_item.children
+    assert all(isinstance(item, RelationCatalogItem) for item in table_items)
+
+    [foo_item] = filter(lambda item: item.label == "foo", table_items)
+    assert isinstance(foo_item, TableCatalogItem)
+    assert not foo_item.children
+    assert not foo_item.loaded
+
+    foo_column_items = foo_item.fetch_children()
+    assert all(isinstance(item, ColumnCatalogItem) for item in foo_column_items)
+
+    [schema_two_item] = filter(lambda item: item.label == "two", schema_items)
+    assert isinstance(schema_two_item, SchemaCatalogItem)
+    assert schema_two_item.children
+    assert schema_two_item.loaded
+
+    view_items = schema_two_item.children
+    assert all(isinstance(item, ViewCatalogItem) for item in view_items)
+
+    [qux_item] = filter(lambda item: item.label == "qux", view_items)
+    assert isinstance(qux_item, ViewCatalogItem)
+    assert not qux_item.children
+    assert not qux_item.loaded
+
+    qux_column_items = qux_item.fetch_children()
+    assert all(isinstance(item, ColumnCatalogItem) for item in qux_column_items)
+
+    assert [item.label for item in foo_column_items] == [
+        item.label for item in qux_column_items
+    ]
+
+    # ensure calling fetch_children on cols doesn't raise
+    children_items = foo_column_items[0].fetch_children()
+    assert not children_items
+
+    # empty schemas don't appear in the catalog
+    schema_three_items = list(filter(lambda item: item.label == "three", schema_items))
+    assert not schema_three_items