databao-context-engine 0.1.1__py3-none-any.whl

databao_context_engine/plugins/databases/clickhouse_introspector.py

@@ -0,0 +1,162 @@
+from __future__ import annotations
+
+from typing import Any, Mapping
+
+import clickhouse_connect
+from pydantic import Field
+
+from databao_context_engine.plugins.base_db_plugin import BaseDatabaseConfigFile
+from databao_context_engine.plugins.databases.base_introspector import BaseIntrospector, SQLQuery
+from databao_context_engine.plugins.databases.databases_types import DatabaseSchema
+from databao_context_engine.plugins.databases.introspection_model_builder import IntrospectionModelBuilder
+
+
+class ClickhouseConfigFile(BaseDatabaseConfigFile):
+    type: str = Field(default="databases/clickhouse")
+    connection: dict[str, Any] = Field(
+        description="Connection parameters for the Clickhouse database. It can contain any of the keys supported by the Clickhouse connection library (see https://clickhouse.com/docs/integrations/language-clients/python/driver-api#connection-arguments)"
+    )
+
+
+class ClickhouseIntrospector(BaseIntrospector[ClickhouseConfigFile]):
+    _IGNORED_SCHEMAS = {"information_schema", "system", "INFORMATION_SCHEMA"}
+
+    supports_catalogs = True
+
+    def _connect(self, file_config: ClickhouseConfigFile):
+        connection = file_config.connection
+        if not isinstance(connection, Mapping):
+            raise ValueError("Invalid YAML config: 'connection' must be a mapping of connection parameters")
+
+        return clickhouse_connect.get_client(**connection)
+
+    def _connect_to_catalog(self, file_config: ClickhouseConfigFile, catalog: str):
+        return self._connect(file_config)
+
+    def _get_catalogs(self, connection, file_config: ClickhouseConfigFile) -> list[str]:
+        return ["clickhouse"]
+
+    def _sql_list_schemas(self, catalogs: list[str] | None) -> SQLQuery:
+        return SQLQuery(
+            """
+                SELECT 
+                    name AS schema_name, 
+                    'clickhouse' AS catalog_name 
+                FROM 
+                    system.databases
+            """,
+            None,
+        )
+
+    def collect_catalog_model(self, connection, catalog: str, schemas: list[str]) -> list[DatabaseSchema] | None:
+        if not schemas:
+            return []
+
+        schemas_sql = ", ".join(self._quote_literal(s) for s in schemas)
+
+        comps = self._component_queries()
+        results: dict[str, list[dict]] = {cq: [] for cq in comps}
+        for cq, template_sql in comps.items():
+            sql = template_sql.replace("{SCHEMAS}", schemas_sql)
+            results[cq] = self._fetchall_dicts(connection, sql, None)
+
+        return IntrospectionModelBuilder.build_schemas_from_components(
+            schemas=schemas,
+            rels=results.get("relations", []),
+            cols=results.get("columns", []),
+            pk_cols=[],
+            uq_cols=[],
+            checks=[],
+            fk_cols=[],
+            idx_cols=results.get("idx", []),
+        )
+
+    def _component_queries(self) -> dict[str, str]:
+        return {"relations": self._sql_relations(), "columns": self._sql_columns(), "idx": self._sql_indexes()}
+
+    def _sql_relations(self) -> str:
+        return r"""
+            SELECT
+                t.database AS schema_name,
+                t.name AS table_name,
+                multiIf(
+                    t.engine = 'View', 'view',
+                    t.engine = 'MaterializedView', 'materialized_view',
+                    t.engine IN (
+                    'File', 'URL',
+                    'S3', 'S3Queue',
+                    'AzureBlobStorage', 'AzureQueue',
+                    'HDFS',
+                    'MySQL', 'PostgreSQL', 'MongoDB', 'Redis',
+                    'JDBC', 'ODBC', 'SQLite',
+                    'Kafka', 'RabbitMQ', 'NATS',
+                    'ExternalDistributed',
+                    'DeltaLake', 'Iceberg', 'Hudi', 'Hive',
+                    'MaterializedPostgreSQL',
+                    'YTsaurus', 'ArrowFlight'
+                    ), 'external_table',
+                    'table'
+                ) AS kind,
+                t.comment AS description
+            FROM 
+                system.tables t
+            WHERE 
+                t.database IN ({SCHEMAS})
+            ORDER BY 
+                t.name
+        """
+
+    def _sql_columns(self) -> str:
+        return r"""
+            SELECT
+                c.database AS schema_name,
+                c.table AS table_name,
+                c.name AS column_name,
+                c.position AS ordinal_position,
+                c.type AS data_type,
+                (c.type LIKE 'Nullable(%') AS is_nullable,
+                c.default_expression AS default_expression,
+                CASE 
+                    WHEN c.default_kind IN ('MATERIALIZED','ALIAS') THEN 'computed' 
+                END AS generated,
+                c.comment AS description
+            FROM 
+                system.columns c
+            WHERE 
+                c.database IN ({SCHEMAS})
+            ORDER BY 
+                c.table, 
+                c.position
+        """
+
+    def _sql_indexes(self) -> str:
+        return r"""
+            SELECT
+                i.database AS schema_name,
+                i.table AS table_name,
+                i.name AS index_name,
+                1 AS position,
+                i.expr AS expr,
+                0 AS is_unique,
+                i.type AS method,
+                NULL AS predicate
+            FROM 
+                system.data_skipping_indices i
+            WHERE 
+                i.database IN ({SCHEMAS})
+            ORDER BY 
+                i.table, 
+                i.name
+        """
+
+    def _sql_sample_rows(self, catalog: str, schema: str, table: str, limit: int) -> SQLQuery:
+        sql = f'SELECT * FROM "{schema}"."{table}" LIMIT %s'
+        return SQLQuery(sql, (limit,))
+
+    def _fetchall_dicts(self, connection, sql: str, params) -> list[dict]:
+        res = connection.query(sql, parameters=params) if params else connection.query(sql)
+        cols = [c.lower() for c in res.column_names]
+        return [dict(zip(cols, row)) for row in res.result_rows]
+
+    def _quote_literal(self, value: str) -> str:
+        return "'" + str(value).replace("'", "\\'") + "'"

databao_context_engine/plugins/databases/database_chunker.py

@@ -0,0 +1,69 @@
+from dataclasses import dataclass
+
+from databao_context_engine.pluginlib.build_plugin import EmbeddableChunk
+from databao_context_engine.plugins.databases.databases_types import (
+    DatabaseColumn,
+    DatabaseIntrospectionResult,
+    DatabaseTable,
+)
+
+
+@dataclass
+class DatabaseTableChunkContent:
+    catalog_name: str
+    schema_name: str
+    table: DatabaseTable
+
+
+@dataclass
+class DatabaseColumnChunkContent:
+    catalog_name: str
+    schema_name: str
+    table_name: str
+    column: DatabaseColumn
+
+
+def build_database_chunks(result: DatabaseIntrospectionResult) -> list[EmbeddableChunk]:
+    chunks = []
+    for catalog in result.catalogs:
+        for schema in catalog.schemas:
+            for table in schema.tables:
+                chunks.append(_create_table_chunk(catalog.name, schema.name, table))
+
+                for column in table.columns:
+                    chunks.append(_create_column_chunk(catalog.name, schema.name, table.name, column))
+
+    return chunks
+
+
+def _create_table_chunk(catalog_name: str, schema_name: str, table: DatabaseTable) -> EmbeddableChunk:
+    return EmbeddableChunk(
+        embeddable_text=_build_table_chunk_text(table),
+        content=DatabaseTableChunkContent(
+            catalog_name=catalog_name,
+            schema_name=schema_name,
+            table=table,
+        ),
+    )
+
+
+def _create_column_chunk(
+    catalog_name: str, schema_name: str, table_name: str, column: DatabaseColumn
+) -> EmbeddableChunk:
+    return EmbeddableChunk(
+        embeddable_text=_build_column_chunk_text(table_name, column),
+        content=DatabaseColumnChunkContent(
+            catalog_name=catalog_name,
+            schema_name=schema_name,
+            table_name=table_name,
+            column=column,
+        ),
+    )
+
+
+def _build_table_chunk_text(database_table: DatabaseTable) -> str:
+    return f"Table {database_table.name} with columns {','.join([column.name for column in database_table.columns])}"
+
+
+def _build_column_chunk_text(table_name: str, database_object: DatabaseColumn) -> str:
+    return f"Column {database_object.name} in table {table_name}"

databao_context_engine/plugins/databases/databases_types.py

@@ -0,0 +1,114 @@
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Literal
+
+
+class DatasetKind(str, Enum):
+    TABLE = "table"
+    VIEW = "view"
+    MATERIALIZED_VIEW = "materialized_view"
+    EXTERNAL_TABLE = "external_table"
+
+    @classmethod
+    def from_raw(cls, raw: str | None) -> "DatasetKind":
+        default = cls.TABLE
+        if raw is None:
+            return default
+        if isinstance(raw, str):
+            raw = raw.lower()
+        try:
+            return cls(raw)
+        except Exception:
+            return default
+
+
+@dataclass
+class CheckConstraint:
+    name: str | None
+    expression: str
+    validated: bool | None
+
+
+@dataclass
+class KeyConstraint:
+    name: str | None
+    columns: list[str]
+    validated: bool | None
+
+
+@dataclass
+class ForeignKeyColumnMap:
+    from_column: str
+    to_column: str
+
+
+@dataclass
+class ForeignKey:
+    name: str | None
+    mapping: list[ForeignKeyColumnMap]
+    referenced_table: str
+    enforced: bool | None = None
+    validated: bool | None = None
+    on_update: str | None = None
+    on_delete: str | None = None
+    cardinality_inferred: Literal["one_to_one", "many_to_one"] | None = None
+
+
+@dataclass
+class Index:
+    name: str
+    columns: list[str]
+    unique: bool = False
+    method: str | None = None
+    predicate: str | None = None
+
+
+@dataclass
+class DatabaseColumn:
+    name: str
+    type: str
+    nullable: bool
+    description: str | None = None
+    default_expression: str | None = None
+    generated: Literal["identity", "computed"] | None = None
+    checks: list[CheckConstraint] = field(default_factory=list)
+
+
+@dataclass
+class DatabasePartitionInfo:
+    meta: dict[str, Any]
+    partition_tables: list[str]
+
+
+@dataclass
+class DatabaseTable:
+    name: str
+    columns: list[DatabaseColumn]
+    samples: list[dict[str, Any]]
+    partition_info: DatabasePartitionInfo | None = None
+    description: str | None = None
+    kind: DatasetKind = DatasetKind.TABLE
+    primary_key: KeyConstraint | None = None
+    unique_constraints: list[KeyConstraint] = field(default_factory=list)
+    checks: list[CheckConstraint] = field(default_factory=list)
+    indexes: list[Index] = field(default_factory=list)
+    foreign_keys: list[ForeignKey] = field(default_factory=list)
+
+
+@dataclass
+class DatabaseSchema:
+    name: str
+    tables: list[DatabaseTable]
+    description: str | None = None
+
+
+@dataclass
+class DatabaseCatalog:
+    name: str
+    schemas: list[DatabaseSchema]
+    description: str | None = None
+
+
+@dataclass
+class DatabaseIntrospectionResult:
+    catalogs: list[DatabaseCatalog]

databao_context_engine/plugins/databases/duckdb_introspector.py

@@ -0,0 +1,325 @@
+from __future__ import annotations
+
+import duckdb
+from pydantic import BaseModel, Field
+
+from databao_context_engine.plugins.base_db_plugin import BaseDatabaseConfigFile
+from databao_context_engine.plugins.databases.base_introspector import BaseIntrospector, SQLQuery
+from databao_context_engine.plugins.databases.databases_types import DatabaseSchema
+from databao_context_engine.plugins.databases.introspection_model_builder import IntrospectionModelBuilder
+
+
+class DuckDBConfigFile(BaseDatabaseConfigFile):
+    type: str = Field(default="databases/duckdb")
+    connection: DuckDBConnectionConfig
+
+
+class DuckDBConnectionConfig(BaseModel):
+    database: str = Field(description="Path to the DuckDB database file")
+
+
+class DuckDBIntrospector(BaseIntrospector[DuckDBConfigFile]):
+    _IGNORED_CATALOGS = {"system", "temp"}
+    _IGNORED_SCHEMAS = {"information_schema", "pg_catalog"}
+    supports_catalogs = True
+
+    def _connect(self, file_config: DuckDBConfigFile):
+        database_path = str(file_config.connection.database)
+        return duckdb.connect(database=database_path)
+
+    def _connect_to_catalog(self, file_config: DuckDBConfigFile, catalog: str):
+        return self._connect(file_config)
+
+    def _get_catalogs(self, connection, file_config: DuckDBConfigFile) -> list[str]:
+        rows = self._fetchall_dicts(connection, "SELECT database_name FROM duckdb_databases();", None)
+        catalogs = [r["database_name"] for r in rows if r.get("database_name")]
+        catalogs_filtered = [c for c in catalogs if c.lower() not in self._IGNORED_CATALOGS]
+        return catalogs_filtered or [self._resolve_pseudo_catalog_name(file_config)]
+
+    def _sql_list_schemas(self, catalogs: list[str] | None) -> SQLQuery:
+        if not catalogs:
+            return SQLQuery("SELECT schema_name, catalog_name FROM information_schema.schemata", None)
+        sql = "SELECT catalog_name, schema_name FROM information_schema.schemata WHERE catalog_name = ANY(?)"
+        return SQLQuery(sql, (catalogs,))
+
+    def collect_catalog_model(self, connection, catalog: str, schemas: list[str]) -> list[DatabaseSchema] | None:
+        if not schemas:
+            return []
+
+        comps = self._component_queries()
+        results: dict[str, list[dict]] = {cq: [] for cq in comps}
+
+        for cq, sql in comps.items():
+            results[cq] = self._fetchall_dicts(connection, sql, (schemas,))
+
+        return IntrospectionModelBuilder.build_schemas_from_components(
+            schemas=schemas,
+            rels=results.get("relations", []),
+            cols=results.get("columns", []),
+            pk_cols=results.get("pk", []),
+            uq_cols=results.get("uq", []),
+            checks=results.get("checks", []),
+            fk_cols=results.get("fks", []),
+            idx_cols=results.get("idx", []),
+        )
+
+    def _component_queries(self) -> dict[str, str]:
+        return {
+            "relations": self._sql_relations(),
+            "columns": self._sql_columns(),
+            "pk": self._sql_primary_keys(),
+            "uq": self._sql_unique(),
+            "checks": self._sql_checks(),
+            "fks": self._sql_foreign_keys(),
+            "idx": self._sql_indexes(),
+        }
+
+    def _sql_relations(self) -> str:
+        return r"""
+            SELECT
+                table_schema AS schema_name,
+                table_name,
+                CASE table_type
+                    WHEN 'BASE TABLE' THEN 'table'
+                    WHEN 'VIEW' THEN 'view'
+                    WHEN 'MATERIALIZED VIEW' THEN 'materialized_view'
+                    ELSE lower(table_type)
+                END AS kind,
+                NULL::VARCHAR AS description
+            FROM 
+                information_schema.tables
+            WHERE 
+                table_schema = ANY(?)
+            ORDER BY 
+                table_name; 
+        """
+
+    def _sql_columns(self) -> str:
+        return r"""
+            SELECT
+                c.table_schema AS schema_name,
+                c.table_name,
+                c.column_name,
+                c.ordinal_position AS ordinal_position,
+                c.data_type AS data_type,
+                CASE 
+                    WHEN c.is_nullable = 'YES' THEN TRUE 
+                    ELSE FALSE 
+                END AS is_nullable,
+                c.column_default AS default_expression,
+                NULL::VARCHAR AS generated,
+                NULL::VARCHAR AS description
+            FROM 
+                information_schema.columns c
+            WHERE 
+                c.table_schema = ANY(?)
+            ORDER BY 
+                c.table_schema,
+                c.table_name, 
+                c.ordinal_position; 
+        """
+
+    def _sql_primary_keys(self) -> str:
+        return r"""
+            WITH d AS (
+                SELECT 
+                    *
+                FROM 
+                    duckdb_constraints()
+                WHERE 
+                    schema_name = ANY(?)
+                    AND constraint_type = 'PRIMARY KEY'
+            ),
+            cols AS (
+                SELECT
+                    d.schema_name,
+                    d.table_name,
+                    d.constraint_name,
+                    r.pos AS position,
+                    d.constraint_column_names[r.pos] AS column_name
+                FROM 
+                    d,
+                    range(1, length(d.constraint_column_names) + 1) AS r(pos)
+            )
+            SELECT
+                schema_name,
+                table_name,
+                constraint_name,
+                position,
+                column_name
+            FROM 
+                cols
+            ORDER BY
+                schema_name,
+                table_name, 
+                constraint_name, 
+                position;
+        """
+
+    def _sql_unique(self) -> str:
+        return r"""
+            WITH d AS (
+                SELECT 
+                    *
+                FROM 
+                    duckdb_constraints()
+                WHERE 
+                    schema_name = ANY(?)
+                    AND constraint_type = 'UNIQUE'
+            ),
+            cols AS (
+                SELECT
+                    d.schema_name,
+                    d.table_name,
+                    d.constraint_name,
+                    r.pos AS position,
+                    d.constraint_column_names[r.pos] AS column_name
+                FROM 
+                    d,
+                    range(1, length(d.constraint_column_names) + 1) AS r(pos)
+            )
+            SELECT
+                schema_name,
+                table_name,
+                constraint_name,
+                position,
+                column_name
+            FROM 
+                cols
+            ORDER BY
+                schema_name,
+                table_name, 
+                constraint_name, 
+                position;
+        """
+
+    def _sql_checks(self) -> str:
+        return r"""
+            SELECT
+                d.schema_name,
+                d.table_name,
+                d.constraint_name,
+                d.expression        AS expression,
+                TRUE                AS validated
+            FROM 
+                duckdb_constraints() AS d
+            WHERE 
+                d.schema_name = ANY(?)
+                AND d.constraint_type = 'CHECK'
+            ORDER BY 
+                d.schema_name, 
+                d.table_name, 
+                d.constraint_name; 
+           """
+
+    def _sql_foreign_keys(self) -> str:
+        return r"""
+            WITH d AS (
+                SELECT 
+                    *
+                FROM 
+                    duckdb_constraints()
+                WHERE 
+                    schema_name = ANY(?)
+                    AND constraint_type = 'FOREIGN KEY'
+            ),
+            cols AS (
+                SELECT
+                    d.schema_name,
+                    d.table_name,
+                    d.constraint_name,
+                    r.pos AS position,
+                    d.constraint_column_names[r.pos]  AS from_column,
+                    d.referenced_column_names[r.pos]  AS to_column
+                FROM 
+                    d,
+                    range(1, length(d.constraint_column_names) + 1) AS r(pos)
+            ),
+            ref AS (
+            SELECT
+                rc.constraint_schema AS schema_name,
+                rc.constraint_name,
+                tc.table_schema AS ref_schema,
+                tc.table_name   AS ref_table
+            FROM 
+                information_schema.referential_constraints rc
+                JOIN information_schema.table_constraints tc ON 
+                    tc.constraint_schema = rc.unique_constraint_schema 
+                    AND tc.constraint_name = rc.unique_constraint_name
+            ),
+            rules AS (
+                SELECT
+                    constraint_schema AS schema_name,
+                    constraint_name,
+                    lower(update_rule) AS on_update,
+                    lower(delete_rule) AS on_delete
+                FROM 
+                    information_schema.referential_constraints
+            )
+            SELECT
+                c.schema_name,
+                c.table_name,
+                c.constraint_name,
+                c.position,
+                c.from_column,
+                r.ref_schema,
+                r.ref_table,
+                c.to_column,
+                coalesce(u.on_update, 'no action') AS on_update,
+                coalesce(u.on_delete, 'no action') AS on_delete,
+                TRUE AS enforced,
+                TRUE AS validated
+            FROM 
+                cols c JOIN ref r ON r.schema_name = c.schema_name AND r.constraint_name = c.constraint_name
+            LEFT JOIN rules u ON u.schema_name = c.schema_name AND u.constraint_name = c.constraint_name
+            ORDER BY 
+                c.schema_name, 
+                c.table_name, 
+                c.constraint_name, 
+                c.position;
+        """
+
+    def _sql_indexes(self) -> str:
+        return r"""
+            WITH idx AS (
+                SELECT
+                    schema_name,
+                    table_name,
+                    index_name,
+                    is_unique,
+                    string_split(trim(BOTH '[]' FROM expressions), ',') AS expr_list
+                FROM 
+                    duckdb_indexes()
+                WHERE 
+                    schema_name = ANY(?)
+            )
+            SELECT
+                schema_name,
+                table_name,
+                index_name,
+                pos AS position,
+                trim(expr_list[pos]) AS expr,
+                is_unique
+            FROM 
+                idx,
+                range(1, length(expr_list) + 1) AS r(pos)
+            ORDER BY
+                schema_name, 
+                table_name,
+                index_name,
+                position;
+         """
+
+    def _sql_sample_rows(self, catalog: str, schema: str, table: str, limit: int) -> SQLQuery:
+        sql = f'SELECT * FROM "{schema}"."{table}" LIMIT ?'
+        return SQLQuery(sql, (limit,))
+
+    def _quote_literal(self, value: str) -> str:
+        return "'" + str(value).replace("'", "''") + "'"
+
+    def _fetchall_dicts(self, connection, sql: str, params) -> list[dict]:
+        cur = connection.cursor()
+        cur.execute(sql, params or [])
+        columns = [desc[0].lower() for desc in cur.description] if cur.description else []
+        rows = cur.fetchall()
+        return [dict(zip(columns, row)) for row in rows]