feat: add data catalog support for oracledb

src/harlequin_odbc/adapter.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
 from contextlib import suppress
+from decimal import Decimal
+from datetime import date, datetime, time
 from typing import TYPE_CHECKING, Any, Sequence
 
 import pyodbc
@@ -29,14 +31,12 @@ if TYPE_CHECKING:
     pass
 
 
-class HarlequinOdbcCursor(HarlequinCursor):
+class OdbcTypeLabels:
-    def __init__(self, cur: pyodbc.Cursor) -> None:
+    """Shared Harlequin type-label mapper for result sets and catalog metadata."""
-        self.cur = cur
-        self._limit: int | None = None
 
-    def columns(self) -> list[tuple[str, str]]:
+    UNKNOWN = "?"
-        # todo: use getTypeInfo
+
-        type_mapping = {
+    PYTHON_TYPE_NAMES: dict[str, str] = {
         "bool": "t/f",
         "int": "##",
         "float": "#.#",
@@ -48,10 +48,75 @@ class HarlequinOdbcCursor(HarlequinCursor):
         "datetime": "dt",
         "UUID": "uid",
     }
+
+    PYTHON_TYPES: dict[type[Any], str] = {
+        bool: "t/f",
+        int: "##",
+        float: "#.#",
+        Decimal: "#.#",
+        str: "s",
+        bytes: "0b",
+        date: "d",
+        time: "t",
+        datetime: "dt",
+    }
+
+    @classmethod
+    def from_python_type(cls, col_type: type[Any] | None) -> str:
+        if col_type is None:
+            return cls.UNKNOWN
+        return cls.PYTHON_TYPES.get(
+            col_type, cls.PYTHON_TYPE_NAMES.get(col_type.__name__, cls.UNKNOWN)
+        )
+
+    @classmethod
+    def from_oracle_metadata(
+        cls,
+        data_type: str | None,
+        data_precision: int | None = None,
+        data_scale: int | None = None,
+    ) -> str:
+        """Map Oracle dictionary metadata to the same labels used for query results."""
+        if data_type is None:
+            return cls.UNKNOWN
+
+        normalized_type = data_type.upper()
+
+        if normalized_type in {"CHAR", "NCHAR", "VARCHAR2", "NVARCHAR2", "CLOB", "NCLOB", "LONG"}:
+            return cls.PYTHON_TYPE_NAMES["str"]
+
+        if normalized_type == "NUMBER":
+            if data_precision == 1 and (data_scale is None or data_scale == 0):
+                return cls.PYTHON_TYPE_NAMES["bool"]
+            if data_scale is None or data_scale == 0:
+                return cls.PYTHON_TYPE_NAMES["int"]
+            return cls.PYTHON_TYPE_NAMES["Decimal"]
+
+        if normalized_type in {"BINARY_FLOAT", "BINARY_DOUBLE", "FLOAT"}:
+            return cls.PYTHON_TYPE_NAMES["float"]
+
+        if normalized_type == "DATE" or normalized_type.startswith("TIMESTAMP"):
+            return cls.PYTHON_TYPE_NAMES["datetime"]
+
+        if normalized_type.startswith("INTERVAL"):
+            return "td"
+
+        if normalized_type in {"RAW", "LONG RAW", "BLOB", "BFILE"}:
+            return cls.PYTHON_TYPE_NAMES["bytes"]
+
+        return cls.UNKNOWN
+
+
+class HarlequinOdbcCursor(HarlequinCursor):
+    def __init__(self, cur: pyodbc.Cursor) -> None:
+        self.cur = cur
+        self._limit: int | None = None
+
+    def columns(self) -> list[tuple[str, str]]:
         return [
             (
                 col_name if col_name else "(No column name)",
-                type_mapping.get(col_type.__name__, "?"),
+                OdbcTypeLabels.from_python_type(col_type),
             )
             for col_name, col_type, *_ in self.cur.description
         ]
@@ -80,7 +145,9 @@ class HarlequinOdbcConnection(HarlequinConnection):
         init_message: str = "",
     ) -> None:
         assert len(conn_str) == 1
+        self.conn_str = conn_str
         self.init_message = init_message
+        self._is_oracle_db: bool | None = None
         try:
             self.conn = pyodbc.connect(conn_str[0], autocommit=True)
             self.aux_conn = pyodbc.connect(conn_str[0], autocommit=True)
@@ -106,6 +173,7 @@ class HarlequinOdbcConnection(HarlequinConnection):
 
     def get_catalog(self) -> Catalog:
         raw_catalog = self._list_tables()
+        include_db_in_identifier = not self._is_oracle()
         db_items: list[CatalogItem] = []
         for db, schemas in raw_catalog.items():
             schema_items: list[CatalogItem] = []
@@ -119,6 +187,7 @@ class HarlequinOdbcConnection(HarlequinConnection):
                             db_label=db,
                             rel_type=rel_type,
                             connection=self,
+                            include_db_in_identifier=include_db_in_identifier,
                         )
                     )
                 schema_items.append(
@@ -127,6 +196,7 @@ class HarlequinOdbcConnection(HarlequinConnection):
                         db_label=db,
                         connection=self,
                         children=rel_items,
+                        include_db_in_identifier=include_db_in_identifier,
                     )
                 )
             db_items.append(
@@ -145,6 +215,9 @@ class HarlequinOdbcConnection(HarlequinConnection):
             self.aux_conn.close()
 
     def _list_tables(self) -> dict[str, dict[str, list[tuple[str, str]]]]:
+        if self._is_oracle():
+            return self._list_oracle_tables()
+
         cur = self.aux_conn.cursor()
         catalog: dict[str, dict[str, list[tuple[str, str]]]] = {}
         for db_name, schema_name, rel_name, rel_type, *_ in cur.tables(catalog="%"):
@@ -166,10 +239,107 @@ class HarlequinOdbcConnection(HarlequinConnection):
     def _list_columns_in_relation(
         self, catalog_name: str, schema_name: str, rel_name: str
     ) -> list[tuple[str, str]]:
+        if self._is_oracle():
+            return self._list_oracle_columns_in_relation(schema_name, rel_name)
+
         cur = self.aux_conn.cursor()
         raw_cols = cur.columns(table=rel_name, catalog=catalog_name, schema=schema_name)
         return [(col[3], col[5]) for col in raw_cols]
 
+    def _is_oracle(self) -> bool:
+        if self._is_oracle_db is None:
+            dbms_name = ""
+            with suppress(Exception):
+                dbms_name_info = self.conn.getinfo(pyodbc.SQL_DBMS_NAME)
+                if dbms_name_info:
+                    dbms_name = str(dbms_name_info)
+            if not dbms_name:
+                dbms_name = self.conn_str[0]
+            self._is_oracle_db = "oracle" in dbms_name.lower()
+        return self._is_oracle_db is True
+
+    def _list_oracle_tables(self) -> dict[str, dict[str, list[tuple[str, str]]]]:
+        cur = self.aux_conn.cursor()
+        database_name = self._get_oracle_database_name(cur)
+        cur.execute(
+            """
+            select owner, table_name, rel_type
+            from (
+                select
+                    owner,
+                    table_name,
+                    case temporary
+                        when 'Y' then 'GLOBAL TEMPORARY'
+                        else 'TABLE'
+                    end as rel_type
+                from all_tables
+                where owner not in (
+                    'SYS', 'SYSTEM', 'XDB', 'CTXSYS', 'MDSYS', 'ORDSYS', 'WMSYS'
+                )
+                union all
+                select owner, view_name as table_name, 'VIEW' as rel_type
+                from all_views
+                where owner not in (
+                    'SYS', 'SYSTEM', 'XDB', 'CTXSYS', 'MDSYS', 'ORDSYS', 'WMSYS'
+                )
+            )
+            order by owner, table_name
+            """
+        )
+
+        catalog: dict[str, dict[str, list[tuple[str, str]]]] = {database_name: {}}
+        for schema_name, rel_name, rel_type in cur:
+            if schema_name not in catalog[database_name]:
+                catalog[database_name][schema_name] = []
+            catalog[database_name][schema_name].append((rel_name, rel_type))
+
+        return catalog
+
+    def _get_oracle_database_name(self, cur: pyodbc.Cursor) -> str:
+        with suppress(Exception):
+            cur.execute("select sys_context('USERENV', 'DB_NAME') from dual")
+            database_name = cur.fetchval()
+            if database_name:
+                return str(database_name)
+
+        with suppress(Exception):
+            database_name = self.conn.getinfo(pyodbc.SQL_DATABASE_NAME)
+            if database_name:
+                return str(database_name)
+
+        return "Oracle"
+
+    def _list_oracle_columns_in_relation(
+        self, schema_name: str, rel_name: str
+    ) -> list[tuple[str, str]]:
+        cur = self.aux_conn.cursor()
+        cur.execute(
+            """
+            select
+                column_name,
+                data_type,
+                data_precision,
+                data_scale
+            from all_tab_columns
+            where owner = ?
+                and table_name = ?
+            order by column_id
+            """,
+            schema_name,
+            rel_name,
+        )
+        return [
+            (
+                column_name,
+                OdbcTypeLabels.from_oracle_metadata(
+                    data_type=data_type,
+                    data_precision=data_precision,
+                    data_scale=data_scale,
+                ),
+            )
+            for column_name, data_type, data_precision, data_scale in cur
+        ]
+
     def get_completions(self) -> list[HarlequinCompletion]:
         return []
 
@@ -191,3 +361,4 @@ class HarlequinOdbcAdapter(HarlequinAdapter):
     def connect(self) -> HarlequinOdbcConnection:
         conn = HarlequinOdbcConnection(self.conn_str)
         return conn
+

src/harlequin_odbc/catalog.py

@@ -60,6 +60,7 @@ class RelationCatalogItem(InteractiveCatalogItem["HarlequinOdbcConnection"]):
         db_label: str,
         rel_type: str,
         connection: "HarlequinOdbcConnection",
+        include_db_in_identifier: bool = True,
     ) -> "RelationCatalogItem":
         rel_type_map: dict[str, type[RelationCatalogItem]] = {
             "TABLE": TableCatalogItem,
@@ -70,8 +71,13 @@ class RelationCatalogItem(InteractiveCatalogItem["HarlequinOdbcConnection"]):
         }
 
         item_class = rel_type_map.get(rel_type, TableCatalogItem)
+        if include_db_in_identifier:
+            qualified_identifier = f'"{db_label}"."{schema_label}"."{label}"'
+        else:
+            qualified_identifier = f'"{schema_label}"."{label}"'
+
         return item_class(
-            qualified_identifier=f'"{db_label}"."{schema_label}"."{label}"',
+            qualified_identifier=qualified_identifier,
             query_name=f'"{schema_label}"."{label}"',
             label=label,
             schema_label=schema_label,
@@ -137,12 +143,17 @@ class SchemaCatalogItem(InteractiveCatalogItem["HarlequinOdbcConnection"]):
         db_label: str,
         connection: "HarlequinOdbcConnection",
         children: list[CatalogItem] | None = None,
+        include_db_in_identifier: bool = True,
     ) -> "SchemaCatalogItem":
         schema_identifier = f'"{label}"'
         if children is None:
             children = []
+        qualified_identifier = schema_identifier
+        if include_db_in_identifier:
+            qualified_identifier = f'"{db_label}".{schema_identifier}'
+
         return cls(
-            qualified_identifier=f'"{db_label}".{schema_identifier}',
+            qualified_identifier=qualified_identifier,
             query_name=schema_identifier,
             label=label,
             db_label=db_label,

tests/test_oracle_catalog.py

@@ -0,0 +1,100 @@
+from __future__ import annotations
+
+from collections.abc import Iterator
+
+import pyodbc
+import pytest
+
+from harlequin_odbc.adapter import HarlequinOdbcConnection
+from harlequin_odbc.catalog import (
+    DatabaseCatalogItem,
+    SchemaCatalogItem,
+    TableCatalogItem,
+    ViewCatalogItem,
+)
+
+
+class FakeCursor:
+    def __init__(self) -> None:
+        self.executed: list[tuple[str, tuple[object, ...]]] = []
+        self.rows: list[tuple[str, ...]] = []
+
+    def execute(self, sql: str, *params: object) -> FakeCursor:
+        self.executed.append((sql, params))
+        normalized_sql = sql.lower()
+        if "sys_context" in normalized_sql:
+            self.rows = [("ORCL",)]
+        elif "all_tab_columns" in normalized_sql:
+            self.rows = [("EMPLOYEE_ID", "NUMBER(6,0)"), ("FIRST_NAME", "VARCHAR2(20)")]
+        elif "all_tables" in normalized_sql and "all_views" in normalized_sql:
+            self.rows = [("HR", "EMPLOYEES", "TABLE"), ("HR", "EMPLOYEE_VIEW", "VIEW")]
+        else:
+            self.rows = []
+        return self
+
+    def fetchval(self) -> str | None:
+        if not self.rows:
+            return None
+        return self.rows[0][0]
+
+    def tables(self, *_args: object, **_kwargs: object) -> object:
+        raise AssertionError("Oracle catalog should not call SQLTables")
+
+    def columns(self, *_args: object, **_kwargs: object) -> object:
+        raise AssertionError("Oracle catalog should not call SQLColumns")
+
+    def __iter__(self) -> Iterator[tuple[str, ...]]:
+        return iter(self.rows)
+
+
+class FakeConnection:
+    def __init__(self) -> None:
+        self.cursor_instance = FakeCursor()
+
+    def cursor(self) -> FakeCursor:
+        return self.cursor_instance
+
+    def getinfo(self, info_type: int) -> str:
+        if info_type == pyodbc.SQL_DBMS_NAME:
+            return "Oracle"
+        if info_type == pyodbc.SQL_DATABASE_NAME:
+            return "ORCL"
+        return ""
+
+    def close(self) -> None:
+        pass
+
+
+def test_oracle_catalog_uses_data_dictionary(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    fake_connection = FakeConnection()
+
+    def connect(*_args: object, **_kwargs: object) -> FakeConnection:
+        return fake_connection
+
+    monkeypatch.setattr(pyodbc, "connect", connect)
+    connection = HarlequinOdbcConnection(("Driver={Oracle};",))
+
+    catalog = connection.get_catalog()
+
+    [database_item] = catalog.items
+    assert isinstance(database_item, DatabaseCatalogItem)
+    assert database_item.label == "ORCL"
+    [schema_item] = database_item.children
+    assert isinstance(schema_item, SchemaCatalogItem)
+    assert schema_item.label == "HR"
+
+    [table_item, view_item] = schema_item.children
+    assert isinstance(table_item, TableCatalogItem)
+    assert table_item.label == "EMPLOYEES"
+    assert table_item.query_name == '"HR"."EMPLOYEES"'
+    assert isinstance(view_item, ViewCatalogItem)
+    assert view_item.label == "EMPLOYEE_VIEW"
+
+    column_items = table_item.fetch_children()
+    assert [(item.label, item.type_label) for item in column_items] == [
+        ("EMPLOYEE_ID", "NUMBER(6,0)"),
+        ("FIRST_NAME", "VARCHAR2(20)"),
+    ]
+    assert fake_connection.cursor_instance.executed[-1][1] == ("HR", "EMPLOYEES")