pymnemon 0.1.1__py3-none-any.whl

mnemon/utils/databricks_engine.py

@@ -0,0 +1,165 @@
+"""
+Minimal Databricks engine shim to mimic SQLAlchemy for metadata + queries.
+
+The native Databricks SQL driver is not a SQLAlchemy dialect, so we wrap it
+to provide:
+- `connection.execute(...)` that behaves like a Result (returns_rows, keys, fetchall)
+- `engine.dialect.get_schema_names/get_table_names/get_columns` for reflection-like use
+"""
+from dataclasses import dataclass
+from typing import Any, Iterable, List, Optional
+
+from databricks import sql
+from sqlalchemy.sql import ClauseElement
+
+
+@dataclass
+class _ResultAdapter:
+    """Lightweight adapter to mimic SQLAlchemy Result for mnemon usage."""
+    cursor: Any
+    returns_rows: bool = True
+
+    def keys(self) -> List[str]:
+        if not getattr(self.cursor, "description", None):
+            return []
+        return [col[0] for col in self.cursor.description]
+
+    def fetchall(self) -> Iterable:
+        return self.cursor.fetchall()
+
+
+class _DatabricksDialect:
+    """
+    Tiny dialect shim providing inspector-style helpers for tables/columns.
+    Uses SHOW/DESCRIBE to avoid relying on information_schema availability.
+    """
+
+    def get_inspector(self, engine: "DatabricksEngine") -> "DatabricksInspector":
+        return DatabricksInspector(engine)
+
+    def get_schema_names(self, connection: "DatabricksConnection") -> List[str]:
+        cur = connection.execute("SHOW DATABASES")
+        return [row[0] for row in cur.fetchall()]
+
+    def get_table_names(
+        self,
+        connection: "DatabricksConnection",
+        schema: Optional[str] = None,
+    ) -> List[str]:
+        stmt = f"SHOW TABLES IN `{schema}`" if schema else "SHOW TABLES"
+        cur = connection.execute(stmt)
+        rows = cur.fetchall()
+        # SHOW TABLES columns: database, tableName, isTemporary
+        return [row[1] if len(row) > 1 else row[0] for row in rows]
+
+    def get_columns(
+        self,
+        connection: "DatabricksConnection",
+        table_name: str,
+        schema: Optional[str] = None,
+    ) -> List[dict]:
+        qualified = f"`{schema}`.`{table_name}`" if schema else f"`{table_name}`"
+        cur = connection.execute(f"DESCRIBE TABLE {qualified}")
+        cols: List[dict] = []
+        for row in cur.fetchall():
+            # DESCRIBE TABLE returns: col_name, data_type, comment, ...
+            if not row or len(row) < 2:
+                continue
+            name, dtype = row[0], row[1]
+            if name is None or str(name).startswith("#"):
+                continue  # skip footer rows
+            cols.append(
+                {
+                    "name": name,
+                    "type": dtype,
+                    "nullable": True,  # DESCRIBE does not expose nullability
+                    "default": None,
+                }
+            )
+        return cols
+
+
+class DatabricksInspector:
+    """Inspector-style helper to mirror SQLAlchemy's reflection surface."""
+
+    def __init__(self, engine: "DatabricksEngine"):
+        self.engine = engine
+
+    def get_schema_names(self) -> List[str]:
+        with self.engine.connect() as conn:
+            return self.engine.dialect.get_schema_names(conn)
+
+    def get_table_names(self, schema: Optional[str] = None) -> List[str]:
+        with self.engine.connect() as conn:
+            return self.engine.dialect.get_table_names(conn, schema=schema)
+
+    def get_columns(self, table_name: str, schema: Optional[str] = None) -> List[dict]:
+        with self.engine.connect() as conn:
+            return self.engine.dialect.get_columns(conn, table_name, schema=schema)
+
+
+class DatabricksEngine:
+    def __init__(self, conn):
+        self._conn = conn
+        self.dialect = _DatabricksDialect()
+
+    def connect(self):
+        return DatabricksConnection(self._conn, self.dialect, engine=self)
+
+    def dispose(self):
+        try:
+            self._conn.close()
+        except Exception:
+            pass
+
+
+class DatabricksConnection:
+    def __init__(self, conn, dialect: _DatabricksDialect, engine: Optional["DatabricksEngine"] = None):
+        self._conn = conn
+        self.dialect = dialect
+        self.engine = engine
+
+    def execute(self, statement, parameters=None):
+        # Accept either raw SQL string or a compiled ClauseElement
+        if isinstance(statement, ClauseElement):
+            compiled = statement.compile(compile_kwargs={"literal_binds": True})
+            sql_stmt = str(compiled)
+        else:
+            sql_stmt = str(statement)
+
+        cursor = self._conn.cursor()
+        cursor.execute(sql_stmt, parameters or {})
+        return _ResultAdapter(cursor)
+
+    # Alias to mirror SQLAlchemy Connection API
+    def exec_driver_sql(self, statement: str, parameters=None):
+        return self.execute(statement, parameters)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc, tb):
+        self.close()
+
+    def close(self):
+        pass
+
+
+def build_databricks_connection(auth_method: str, config: dict):
+    common = {
+        "server_hostname": config["host"],
+        "http_path": config["http_path"],
+    }
+
+    if auth_method == "oauth_m2m":
+        return sql.connect(
+            **common,
+            auth_type="oauth-m2m",
+            client_id=config["client_id"],
+            client_secret=config["client_secret"],
+        )
+
+    return sql.connect(
+        **common,
+        access_token=config["access_token"],
+    )

mnemon/utils/dialect_queries.py

@@ -0,0 +1,374 @@
+"""
+DIALECT_QUERIES - Database-specific SQL for introspection methods.
+Used by SchemaDBConnection for cross-database compatibility.
+"""
+
+DIALECT_QUERIES = {
+    "postgresql": {
+        "get_schema_tables": """
+            SELECT table_name
+            FROM information_schema.tables
+            WHERE table_schema = :schema
+              AND table_type = 'BASE TABLE'
+            ORDER BY table_name
+        """,
+        "get_table_columns": """
+            SELECT column_name, data_type, is_nullable, column_default
+            FROM information_schema.columns
+            WHERE table_schema = :schema AND table_name = :table
+            ORDER BY ordinal_position
+        """,
+        "get_table_sample": 'SELECT * FROM "{schema}"."{table}" LIMIT :limit',
+        "get_table_count": 'SELECT COUNT(*) AS count FROM "{schema}"."{table}"',
+        "get_schemas": """
+            SELECT schema_name
+            FROM information_schema.schemata
+            WHERE schema_name NOT IN ('pg_catalog', 'information_schema', 'pg_toast')
+            ORDER BY schema_name
+        """,
+    },
+
+    "snowflake": {
+        "get_schema_tables": """
+            SELECT table_name
+            FROM information_schema.tables
+            WHERE table_schema = :schema
+              AND table_type = 'BASE TABLE'
+            ORDER BY table_name
+        """,
+        "get_table_columns": """
+            SELECT column_name, data_type, is_nullable, column_default
+            FROM information_schema.columns
+            WHERE table_schema = :schema AND table_name = :table
+            ORDER BY ordinal_position
+        """,
+        "get_table_sample": 'SELECT * FROM "{schema}"."{table}" LIMIT :limit',
+        "get_table_count": 'SELECT COUNT(*) AS count FROM "{schema}"."{table}"',
+        "get_schemas": """
+            SELECT schema_name
+            FROM information_schema.schemata
+            ORDER BY schema_name
+        """,
+    },
+
+    "bigquery": {
+        "get_schema_tables": """
+            SELECT table_name
+            FROM `{project}.{schema}.INFORMATION_SCHEMA.TABLES`
+            WHERE table_type = 'BASE TABLE'
+            ORDER BY table_name
+        """,
+        "get_table_columns": """
+            SELECT column_name, data_type, is_nullable, column_default
+            FROM `{project}.{schema}.INFORMATION_SCHEMA.COLUMNS`
+            WHERE table_name = :table
+            ORDER BY ordinal_position
+        """,
+        "get_table_sample": "SELECT * FROM `{project}.{schema}.{table}` LIMIT :limit",
+        "get_table_count": "SELECT COUNT(*) AS count FROM `{project}.{schema}.{table}`",
+        "get_schemas": """
+            SELECT schema_name
+            FROM `{project}.INFORMATION_SCHEMA.SCHEMATA`
+            ORDER BY schema_name
+        """,
+    },
+
+    "redshift": {
+        "get_schema_tables": """
+            SELECT tablename AS table_name
+            FROM pg_tables
+            WHERE schemaname = :schema
+            ORDER BY tablename
+        """,
+        "get_table_columns": """
+            SELECT column_name, data_type, is_nullable, column_default
+            FROM information_schema.columns
+            WHERE table_schema = :schema AND table_name = :table
+            ORDER BY ordinal_position
+        """,
+        "get_table_sample": 'SELECT * FROM "{schema}"."{table}" LIMIT :limit',
+        "get_table_count": 'SELECT COUNT(*) AS count FROM "{schema}"."{table}"',
+        "get_schemas": """
+            SELECT schema_name
+            FROM information_schema.schemata
+            WHERE schema_name NOT IN ('pg_catalog', 'information_schema', 'pg_internal')
+            ORDER BY schema_name
+        """,
+    },
+
+    "clickhouse": {
+        "get_schema_tables": """
+            SELECT name AS table_name
+            FROM system.tables
+            WHERE database = :schema
+              AND engine NOT IN ('View', 'MaterializedView')
+            ORDER BY name
+        """,
+        "get_table_columns": """
+            SELECT name AS column_name, type AS data_type, 
+                   if(type LIKE 'Nullable%', 'YES', 'NO') AS is_nullable,
+                   default_expression AS column_default
+            FROM system.columns
+            WHERE database = :schema AND table = :table
+            ORDER BY position
+        """,
+        "get_table_sample": "SELECT * FROM {schema}.{table} LIMIT :limit",
+        "get_table_count": "SELECT COUNT(*) AS count FROM {schema}.{table}",
+        "get_schemas": """
+            SELECT name AS schema_name
+            FROM system.databases
+            WHERE name NOT IN ('system', 'INFORMATION_SCHEMA', 'information_schema')
+            ORDER BY name
+        """,
+    },
+
+    "duckdb": {
+        "get_schema_tables": """
+            SELECT table_name
+            FROM information_schema.tables
+            WHERE table_schema = :schema
+              AND table_type = 'BASE TABLE'
+            ORDER BY table_name
+        """,
+        "get_table_columns": """
+            SELECT column_name, data_type, is_nullable, column_default
+            FROM information_schema.columns
+            WHERE table_schema = :schema AND table_name = :table
+            ORDER BY ordinal_position
+        """,
+        "get_table_sample": 'SELECT * FROM "{schema}"."{table}" LIMIT :limit',
+        "get_table_count": 'SELECT COUNT(*) AS count FROM "{schema}"."{table}"',
+        "get_schemas": """
+            SELECT schema_name
+            FROM information_schema.schemata
+            ORDER BY schema_name
+        """,
+    },
+
+    "mysql": {
+        "get_schema_tables": """
+            SELECT table_name
+            FROM information_schema.tables
+            WHERE table_schema = :schema
+              AND table_type = 'BASE TABLE'
+            ORDER BY table_name
+        """,
+        "get_table_columns": """
+            SELECT column_name, data_type, is_nullable, column_default
+            FROM information_schema.columns
+            WHERE table_schema = :schema AND table_name = :table
+            ORDER BY ordinal_position
+        """,
+        "get_table_sample": "SELECT * FROM `{schema}`.`{table}` LIMIT :limit",
+        "get_table_count": "SELECT COUNT(*) AS count FROM `{schema}`.`{table}`",
+        "get_schemas": """
+            SELECT schema_name
+            FROM information_schema.schemata
+            WHERE schema_name NOT IN ('mysql', 'information_schema', 'performance_schema', 'sys')
+            ORDER BY schema_name
+        """,
+    },
+
+    "sqlserver": {
+        "get_schema_tables": """
+            SELECT table_name
+            FROM information_schema.tables
+            WHERE table_schema = :schema
+              AND table_type = 'BASE TABLE'
+            ORDER BY table_name
+        """,
+        "get_table_columns": """
+            SELECT column_name, data_type, is_nullable, column_default
+            FROM information_schema.columns
+            WHERE table_schema = :schema AND table_name = :table
+            ORDER BY ordinal_position
+        """,
+        "get_table_sample": "SELECT TOP (:limit) * FROM [{schema}].[{table}]",
+        "get_table_count": "SELECT COUNT(*) AS count FROM [{schema}].[{table}]",
+        "get_schemas": """
+            SELECT schema_name
+            FROM information_schema.schemata
+            WHERE schema_name NOT IN ('sys', 'guest', 'INFORMATION_SCHEMA')
+            ORDER BY schema_name
+        """,
+    },
+
+    "databricks": {
+        "get_schema_tables": """
+            SELECT table_name
+            FROM information_schema.tables
+            WHERE table_schema = :schema
+              AND table_type = 'BASE TABLE'
+            ORDER BY table_name
+        """,
+        "get_table_columns": """
+            SELECT column_name, data_type, is_nullable, column_default
+            FROM information_schema.columns
+            WHERE table_schema = :schema AND table_name = :table
+            ORDER BY ordinal_position
+        """,
+        "get_table_sample": "SELECT * FROM `{schema}`.`{table}` LIMIT :limit",
+        "get_table_count": "SELECT COUNT(*) AS count FROM `{schema}`.`{table}`",
+        "get_schemas": """
+            SELECT schema_name
+            FROM information_schema.schemata
+            ORDER BY schema_name
+        """,
+    },
+
+    "athena": {
+        "get_schema_tables": """
+            SELECT table_name
+            FROM information_schema.tables
+            WHERE table_schema = :schema
+              AND table_type = 'BASE TABLE'
+            ORDER BY table_name
+        """,
+        "get_table_columns": """
+            SELECT column_name, data_type, is_nullable, '' AS column_default
+            FROM information_schema.columns
+            WHERE table_schema = :schema AND table_name = :table
+            ORDER BY ordinal_position
+        """,
+        "get_table_sample": 'SELECT * FROM "{schema}"."{table}" LIMIT :limit',
+        "get_table_count": 'SELECT COUNT(*) AS count FROM "{schema}"."{table}"',
+        "get_schemas": """
+            SELECT schema_name
+            FROM information_schema.schemata
+            ORDER BY schema_name
+        """,
+    },
+
+    "trino": {
+        "get_schema_tables": """
+            SELECT table_name
+            FROM information_schema.tables
+            WHERE table_schema = :schema
+              AND table_type = 'BASE TABLE'
+            ORDER BY table_name
+        """,
+        "get_table_columns": """
+            SELECT column_name, data_type, is_nullable, '' AS column_default
+            FROM information_schema.columns
+            WHERE table_schema = :schema AND table_name = :table
+            ORDER BY ordinal_position
+        """,
+        "get_table_sample": 'SELECT * FROM "{schema}"."{table}" LIMIT :limit',
+        "get_table_count": 'SELECT COUNT(*) AS count FROM "{schema}"."{table}"',
+        "get_schemas": """
+            SELECT schema_name
+            FROM information_schema.schemata
+            ORDER BY schema_name
+        """,
+    },
+
+    "sparksql": {
+        "get_schema_tables": """
+            SHOW TABLES IN {schema}
+        """,
+        "get_table_columns": """
+            DESCRIBE {schema}.{table}
+        """,
+        "get_table_sample": "SELECT * FROM {schema}.{table} LIMIT :limit",
+        "get_table_count": "SELECT COUNT(*) AS count FROM {schema}.{table}",
+        "get_schemas": """
+            SHOW DATABASES
+        """,
+    },
+
+    "vertica": {
+        "get_schema_tables": """
+            SELECT table_name
+            FROM v_catalog.tables
+            WHERE table_schema = :schema
+            ORDER BY table_name
+        """,
+        "get_table_columns": """
+            SELECT column_name, data_type, is_nullable, column_default
+            FROM v_catalog.columns
+            WHERE table_schema = :schema AND table_name = :table
+            ORDER BY ordinal_position
+        """,
+        "get_table_sample": 'SELECT * FROM "{schema}"."{table}" LIMIT :limit',
+        "get_table_count": 'SELECT COUNT(*) AS count FROM "{schema}"."{table}"',
+        "get_schemas": """
+            SELECT schema_name
+            FROM v_catalog.schemata
+            WHERE NOT is_system_schema
+            ORDER BY schema_name
+        """,
+    },
+
+    "oracle": {
+        "get_schema_tables": """
+            SELECT table_name
+            FROM all_tables
+            WHERE owner = UPPER(:schema)
+            ORDER BY table_name
+        """,
+        "get_table_columns": """
+            SELECT column_name, data_type, nullable AS is_nullable, data_default AS column_default
+            FROM all_tab_columns
+            WHERE owner = UPPER(:schema) AND table_name = UPPER(:table)
+            ORDER BY column_id
+        """,
+        "get_table_sample": 'SELECT * FROM "{schema}"."{table}" WHERE ROWNUM <= :limit',
+        "get_table_count": 'SELECT COUNT(*) AS count FROM "{schema}"."{table}"',
+        "get_schemas": """
+            SELECT username AS schema_name
+            FROM all_users
+            ORDER BY username
+        """,
+    },
+
+    "teradata": {
+        "get_schema_tables": """
+            SELECT TableName AS table_name
+            FROM DBC.TablesV
+            WHERE DatabaseName = :schema
+              AND TableKind = 'T'
+            ORDER BY TableName
+        """,
+        "get_table_columns": """
+            SELECT ColumnName AS column_name, ColumnType AS data_type,
+                   CASE WHEN Nullable = 'Y' THEN 'YES' ELSE 'NO' END AS is_nullable,
+                   DefaultValue AS column_default
+            FROM DBC.ColumnsV
+            WHERE DatabaseName = :schema AND TableName = :table
+            ORDER BY ColumnId
+        """,
+        "get_table_sample": 'SELECT TOP :limit * FROM "{schema}"."{table}"',
+        "get_table_count": 'SELECT COUNT(*) AS count FROM "{schema}"."{table}"',
+        "get_schemas": """
+            SELECT DatabaseName AS schema_name
+            FROM DBC.DatabasesV
+            ORDER BY DatabaseName
+        """,
+    },
+
+    "db2": {
+        "get_schema_tables": """
+            SELECT tabname AS table_name
+            FROM syscat.tables
+            WHERE tabschema = UPPER(:schema)
+              AND type = 'T'
+            ORDER BY tabname
+        """,
+        "get_table_columns": """
+            SELECT colname AS column_name, typename AS data_type,
+                   CASE WHEN nulls = 'Y' THEN 'YES' ELSE 'NO' END AS is_nullable,
+                   default AS column_default
+            FROM syscat.columns
+            WHERE tabschema = UPPER(:schema) AND tabname = UPPER(:table)
+            ORDER BY colno
+        """,
+        "get_table_sample": 'SELECT * FROM "{schema}"."{table}" FETCH FIRST :limit ROWS ONLY',
+        "get_table_count": 'SELECT COUNT(*) AS count FROM "{schema}"."{table}"',
+        "get_schemas": """
+            SELECT schemaname AS schema_name
+            FROM syscat.schemata
+            WHERE schemaname NOT LIKE 'SYS%'
+            ORDER BY schemaname
+        """,
+    },
+}