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

databao_context_engine/plugins/databases/snowflake_introspector.py

@@ -0,0 +1,287 @@
+from __future__ import annotations
+
+from typing import Any, Dict, List, Mapping
+
+import snowflake.connector
+from pydantic import Field
+from snowflake.connector import DictCursor
+
+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 SnowflakeConfigFile(BaseDatabaseConfigFile):
+    type: str = Field(default="databases/snowflake")
+    connection: dict[str, Any] = Field(
+        description="Connection parameters for Snowflake. It can contain any of the keys supported by the Snowflake connection library"
+    )
+
+
+class SnowflakeIntrospector(BaseIntrospector[SnowflakeConfigFile]):
+    _IGNORED_SCHEMAS = {
+        "information_schema",
+    }
+    _IGNORED_CATALOGS = {"STREAMLIT_APPS"}
+    supports_catalogs = True
+
+    def _connect(self, file_config: SnowflakeConfigFile):
+        connection = file_config.connection
+        if not isinstance(connection, Mapping):
+            raise ValueError("Invalid YAML config: 'connection' must be a mapping of connection parameters")
+        snowflake.connector.paramstyle = "qmark"
+        return snowflake.connector.connect(
+            **connection,
+        )
+
+    def _connect_to_catalog(self, file_config: SnowflakeConfigFile, catalog: str):
+        cfg = dict(file_config.connection or {})
+        cfg["database"] = catalog
+        return snowflake.connector.connect(**cfg)
+
+    def _get_catalogs(self, connection, file_config: SnowflakeConfigFile) -> list[str]:
+        database = file_config.connection.get("database")
+        if database:
+            return [database]
+
+        rows = self._fetchall_dicts(connection, "SHOW DATABASES", None)
+        return [r["name"] for r in rows if r["name"] and r["name"].upper() not in self._IGNORED_CATALOGS]
+
+    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)
+        parts = []
+        for catalog in catalogs:
+            parts.append(f"SELECT schema_name, catalog_name FROM {catalog}.information_schema.schemata")
+        return SQLQuery(" UNION ALL ".join(parts), None)
+
+    def collect_catalog_model(self, connection, catalog: str, schemas: list[str]) -> list[DatabaseSchema] | None:
+        if not schemas:
+            return []
+
+        comps = self._component_queries(catalog, schemas)
+
+        statements = [c["sql"].rstrip().rstrip(";") for c in comps]
+        batch_sql = ";\n".join(statements)
+
+        results: dict[str, list[dict]] = {
+            "relations": [],
+            "columns": [],
+            "pk": [],
+            "fks": [],
+            "uq": [],
+        }
+
+        with connection.cursor(DictCursor) as cur:
+            cur.execute(batch_sql, num_statements=len(statements))
+
+            for ix, comp in enumerate(comps, start=1):
+                name = comp["name"]
+
+                rows = self._lower_keys(cur.fetchall()) if cur.description else []
+
+                if name:
+                    results[name].extend(rows)
+
+                if ix < len(comps):
+                    ok = cur.nextset()
+                    if not ok:
+                        raise RuntimeError(
+                            f"Snowflake multi-statement batch ended early after component #{ix} '{name}'"
+                        )
+
+        return IntrospectionModelBuilder.build_schemas_from_components(
+            schemas=schemas,
+            rels=results["relations"],
+            cols=results["columns"],
+            pk_cols=results["pk"],
+            uq_cols=results["uq"],
+            checks=[],
+            fk_cols=results["fks"],
+            idx_cols=[],
+        )
+
+    def _component_queries(self, catalog: str, schemas: list[str]) -> list[dict]:
+        schemas_in = ", ".join(self._quote_literal(s) for s in schemas)
+
+        return [
+            {"name": "relations", "sql": self._sql_relations(catalog, schemas_in)},
+            {"name": "columns", "sql": self._sql_columns(catalog, schemas_in)},
+            {"name": None, "sql": self._sql_pk_show(catalog)},
+            {"name": "pk", "sql": self._sql_pk_select(schemas_in)},
+            {"name": None, "sql": self._sql_fk_show(catalog)},
+            {"name": "fks", "sql": self._sql_fk_select(schemas_in)},
+            {"name": None, "sql": self._sql_uq_show(catalog)},
+            {"name": "uq", "sql": self._sql_uq_select(schemas_in)},
+        ]
+
+    def _sql_relations(self, catalog: str, schemas_in: str) -> str:
+        isq = self._qual_is(catalog)
+
+        return f"""
+            SELECT
+                t.TABLE_SCHEMA AS "schema_name",
+                t.TABLE_NAME AS "table_name",
+                CASE t.TABLE_TYPE
+                    WHEN 'BASE TABLE'        THEN 'table'
+                    WHEN 'VIEW'              THEN 'view'
+                    WHEN 'MATERIALIZED VIEW' THEN 'materialized_view'
+                    WHEN 'EXTERNAL TABLE'    THEN 'external_table'
+                    ELSE LOWER(t.TABLE_TYPE)
+                END AS "kind",
+                t.COMMENT AS "description"
+            FROM 
+                {isq}.TABLES AS t
+            WHERE 
+                t.TABLE_SCHEMA IN ({schemas_in})
+            ORDER BY 
+                t.TABLE_SCHEMA,
+                t.TABLE_NAME
+        """
+
+    def _sql_columns(self, catalog: str, schemas_in: str) -> str:
+        isq = self._qual_is(catalog)
+
+        return f"""
+            SELECT
+            c.TABLE_SCHEMA AS "schema_name",
+                c.TABLE_NAME       AS "table_name",
+                c.COLUMN_NAME      AS "column_name",
+                c.ORDINAL_POSITION AS "ordinal_position",
+                c.DATA_TYPE        AS "data_type",
+                IFF(c.IS_NULLABLE = 'YES', TRUE, FALSE) AS "is_nullable",
+                c.COLUMN_DEFAULT   AS "default_expression",
+                IFF(c.IS_IDENTITY = 'YES', 'identity', NULL) AS "generated",
+                c.COMMENT          AS "description"
+            FROM 
+                {isq}.COLUMNS AS c
+            WHERE 
+                c.TABLE_SCHEMA IN ({schemas_in})
+            ORDER BY 
+                c.TABLE_SCHEMA,
+                c.TABLE_NAME, 
+                c.ORDINAL_POSITION
+        """
+
+    def _sql_pk_show(self, catalog: str) -> str:
+        return f"""
+            SHOW PRIMARY KEYS IN DATABASE {self._quote_ident(catalog)}
+        """
+
+    def _sql_pk_select(self, schemas_in: str) -> str:
+        return f"""
+               SELECT
+                   "schema_name" AS schema_name,
+                   "table_name" AS table_name,
+                   "constraint_name"   AS constraint_name,
+                   "column_name"       AS column_name,
+                   "key_sequence"::INT AS position
+               FROM 
+                   TABLE(RESULT_SCAN(LAST_QUERY_ID()))
+               WHERE
+                    "schema_name" IN ({schemas_in})
+               ORDER BY 
+                   table_name, 
+                   constraint_name, 
+                   position
+               """
+
+    def _sql_fk_show(self, catalog: str) -> str:
+        return f"""
+            SHOW IMPORTED KEYS IN DATABASE {self._quote_ident(catalog)}
+        """
+
+    def _sql_fk_select(self, schemas_in: str) -> str:
+        return f"""
+               SELECT
+                   "fk_schema_name"     AS "schema_name",
+                   "fk_table_name"      AS "table_name",
+                   "fk_name"            AS "constraint_name",
+                   "key_sequence"::INT  AS "position",
+                   "fk_column_name"     AS "from_column",
+                   "pk_schema_name"     AS "ref_schema",
+                   "pk_table_name"      AS "ref_table",
+                   "pk_column_name"     AS "to_column",
+                   LOWER("update_rule") AS "on_update",
+                   LOWER("delete_rule") AS "on_delete",
+                   NULL::BOOLEAN        AS "enforced",
+                   NULL::BOOLEAN        AS "validated"
+               FROM 
+                   TABLE(RESULT_SCAN(LAST_QUERY_ID()))
+               WHERE
+                   "fk_schema_name" IN ({schemas_in})
+               ORDER BY 
+                   "schema_name",
+                   "table_name", 
+                   "constraint_name", 
+                   "position"
+               """
+
+    def _sql_uq_show(self, catalog: str) -> str:
+        return f"""
+            SHOW UNIQUE KEYS IN DATABASE {self._quote_ident(catalog)}
+        """
+
+    def _sql_uq_select(self, schemas_in: str) -> str:
+        return f"""
+               SELECT
+                   "schema_name"       AS "schema_name",
+                   "table_name"        AS "table_name",
+                   "constraint_name"   AS "constraint_name",
+                   "column_name"       AS "column_name",
+                   "key_sequence"::INT AS "position"
+               FROM 
+                   TABLE(RESULT_SCAN(LAST_QUERY_ID()))
+               WHERE
+                   "schema_name" IN ({schemas_in})
+               ORDER BY 
+                   "schema_name",
+                   "table_name", 
+                   "constraint_name", 
+                   "position"
+               """
+
+    def _sql_checks(self, catalog: str, schemas_in: str) -> str:
+        isq = self._qual_is(catalog)
+
+        return f"""
+            SELECT
+                tc.TABLE_SCHEMA    AS "schema_name",
+                tc.TABLE_NAME      AS "table_name",
+                tc.CONSTRAINT_NAME AS "constraint_name",
+                NULL::VARCHAR      AS "expression",
+                TRUE               AS "validated"
+            FROM 
+                {isq}.TABLE_CONSTRAINTS AS tc
+            WHERE 
+                tc.TABLE_SCHEMA IN ({schemas_in})
+                AND tc.CONSTRAINT_TYPE = 'CHECK'
+            ORDER BY 
+                tc.TABLE_SCHEMA, 
+                tc.TABLE_NAME, 
+                tc.CONSTRAINT_NAME
+        """
+
+    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 _fetchall_dicts(self, connection, sql: str, params) -> list[dict]:
+        with connection.cursor(snowflake.connector.DictCursor) as cur:
+            cur.execute(sql, params)
+            rows = cur.fetchall()
+            return [{k.lower(): v for k, v in row.items()} for row in rows]
+
+    def _quote_literal(self, value: str) -> str:
+        return "'" + str(value).replace("'", "''") + "'"
+
+    def _quote_ident(self, ident: str) -> str:
+        return '"' + ident.replace('"', '""') + '"'
+
+    def _qual_is(self, catalog: str) -> str:
+        return f"{self._quote_ident(catalog)}.INFORMATION_SCHEMA"
+
+    @staticmethod
+    def _lower_keys(rows: List[Dict]) -> List[Dict]:
+        return [{k.lower(): v for k, v in row.items()} for row in rows]

databao_context_engine/plugins/duckdb_db_plugin.py

@@ -0,0 +1,12 @@
+from databao_context_engine.plugins.base_db_plugin import BaseDatabasePlugin
+from databao_context_engine.plugins.databases.duckdb_introspector import DuckDBConfigFile, DuckDBIntrospector
+
+
+class DuckDbPlugin(BaseDatabasePlugin[DuckDBConfigFile]):
+    id = "jetbrains/duckdb"
+    name = "DuckDB Plugin"
+    supported = {"databases/duckdb"}
+    config_file_type = DuckDBConfigFile
+
+    def __init__(self):
+        super().__init__(DuckDBIntrospector())

databao_context_engine/plugins/mssql_db_plugin.py

@@ -0,0 +1,12 @@
+from databao_context_engine.plugins.base_db_plugin import BaseDatabasePlugin
+from databao_context_engine.plugins.databases.mssql_introspector import MSSQLConfigFile, MSSQLIntrospector
+
+
+class MSSQLDbPlugin(BaseDatabasePlugin[MSSQLConfigFile]):
+    id = "jetbrains/mssql"
+    name = "MSSQL DB Plugin"
+    supported = {"databases/mssql"}
+    config_file_type = MSSQLConfigFile
+
+    def __init__(self):
+        super().__init__(MSSQLIntrospector())

databao_context_engine/plugins/mysql_db_plugin.py

@@ -0,0 +1,12 @@
+from databao_context_engine.plugins.base_db_plugin import BaseDatabasePlugin
+from databao_context_engine.plugins.databases.mysql_introspector import MySQLConfigFile, MySQLIntrospector
+
+
+class MySQLDbPlugin(BaseDatabasePlugin[MySQLConfigFile]):
+    id = "jetbrains/mysql"
+    name = "MySQL DB Plugin"
+    supported = {"databases/mysql"}
+    config_file_type = MySQLConfigFile
+
+    def __init__(self):
+        super().__init__(MySQLIntrospector())

databao_context_engine/plugins/parquet_plugin.py

@@ -0,0 +1,32 @@
+from typing import Any
+
+from databao_context_engine.pluginlib.build_plugin import BuildDatasourcePlugin, EmbeddableChunk
+from databao_context_engine.plugins.resources.parquet_chunker import build_parquet_chunks
+from databao_context_engine.plugins.resources.parquet_introspector import (
+    ParquetConfigFile,
+    ParquetIntrospector,
+    parquet_type,
+)
+
+
+class ParquetPlugin(BuildDatasourcePlugin[ParquetConfigFile]):
+    id = "jetbrains/parquet"
+    name = "Parquet Plugin"
+    config_file_type = ParquetConfigFile
+
+    def __init__(self):
+        self._introspector = ParquetIntrospector()
+
+    def supported_types(self) -> set[str]:
+        return {parquet_type}
+
+    def divide_context_into_chunks(self, context: Any) -> list[EmbeddableChunk]:
+        return build_parquet_chunks(context)
+
+    def build_context(self, full_type: str, datasource_name: str, file_config: ParquetConfigFile) -> Any:
+        introspection_result = self._introspector.introspect(file_config)
+
+        return introspection_result
+
+    def check_connection(self, full_type: str, datasource_name: str, file_config: ParquetConfigFile) -> None:
+        self._introspector.check_connection(file_config)

databao_context_engine/plugins/plugin_loader.py

@@ -0,0 +1,110 @@
+import logging
+
+from databao_context_engine.pluginlib.build_plugin import (
+    BuildDatasourcePlugin,
+    BuildFilePlugin,
+    BuildPlugin,
+    DatasourceType,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class DuplicatePluginTypeError(RuntimeError):
+    """Raised when two plugins register the same <main>/<sub> plugin key."""
+
+
+PluginList = dict[DatasourceType, BuildPlugin]
+
+
+def get_all_available_plugin_types(exclude_file_plugins: bool = False) -> set[DatasourceType]:
+    return set(load_plugins(exclude_file_plugins=exclude_file_plugins).keys())
+
+
+def get_plugin_for_type(datasource_type: DatasourceType) -> BuildPlugin:
+    all_plugins = load_plugins()
+
+    if datasource_type not in all_plugins:
+        raise ValueError(f"No plugin found for type '{datasource_type.full_type}'")
+
+    return load_plugins()[datasource_type]
+
+
+def load_plugins(exclude_file_plugins: bool = False) -> PluginList:
+    """
+    Loads both builtin and external plugins and merges them into one list
+    """
+    builtin_plugins = _load_builtin_plugins(exclude_file_plugins)
+    external_plugins = _load_external_plugins(exclude_file_plugins)
+    plugins = merge_plugins(builtin_plugins, external_plugins)
+
+    return plugins
+
+
+def _load_builtin_plugins(exclude_file_plugins: bool = False) -> list[BuildPlugin]:
+    all_builtin_plugins: list[BuildPlugin] = []
+
+    all_builtin_plugins += _load_builtin_datasource_plugins()
+
+    if not exclude_file_plugins:
+        all_builtin_plugins += _load_builtin_file_plugins()
+
+    return all_builtin_plugins
+
+
+def _load_builtin_file_plugins() -> list[BuildFilePlugin]:
+    from databao_context_engine.plugins.unstructured_files_plugin import InternalUnstructuredFilesPlugin
+
+    return [
+        InternalUnstructuredFilesPlugin(),
+    ]
+
+
+def _load_builtin_datasource_plugins() -> list[BuildDatasourcePlugin]:
+    """
+    Statically register built-in plugins
+    """
+    from databao_context_engine.plugins.athena_db_plugin import AthenaDbPlugin
+    from databao_context_engine.plugins.clickhouse_db_plugin import ClickhouseDbPlugin
+    from databao_context_engine.plugins.duckdb_db_plugin import DuckDbPlugin
+    from databao_context_engine.plugins.mssql_db_plugin import MSSQLDbPlugin
+    from databao_context_engine.plugins.mysql_db_plugin import MySQLDbPlugin
+    from databao_context_engine.plugins.parquet_plugin import ParquetPlugin
+    from databao_context_engine.plugins.postgresql_db_plugin import PostgresqlDbPlugin
+    from databao_context_engine.plugins.snowflake_db_plugin import SnowflakeDbPlugin
+
+    return [
+        AthenaDbPlugin(),
+        ClickhouseDbPlugin(),
+        DuckDbPlugin(),
+        MSSQLDbPlugin(),
+        MySQLDbPlugin(),
+        PostgresqlDbPlugin(),
+        SnowflakeDbPlugin(),
+        ParquetPlugin(),
+    ]
+
+
+def _load_external_plugins(exclude_file_plugins: bool = False) -> list[BuildPlugin]:
+    """
+    Discover external plugins via entry points
+    """
+    # TODO: implement external plugin loading
+    return []
+
+
+def merge_plugins(*plugin_lists: list[BuildPlugin]) -> PluginList:
+    """
+    Merge multiple plugin maps
+    """
+    registry: PluginList = {}
+    for plugins in plugin_lists:
+        for plugin in plugins:
+            for full_type in plugin.supported_types():
+                datasource_type = DatasourceType(full_type=full_type)
+                if datasource_type in registry:
+                    raise DuplicatePluginTypeError(
+                        f"Plugin type '{datasource_type.full_type}' is provided by both {type(registry[datasource_type]).__name__} and {type(plugin).__name__}"
+                    )
+                registry[datasource_type] = plugin
+    return registry

databao_context_engine/plugins/postgresql_db_plugin.py

@@ -0,0 +1,12 @@
+from databao_context_engine.plugins.base_db_plugin import BaseDatabasePlugin
+from databao_context_engine.plugins.databases.postgresql_introspector import PostgresConfigFile, PostgresqlIntrospector
+
+
+class PostgresqlDbPlugin(BaseDatabasePlugin[PostgresConfigFile]):
+    id = "jetbrains/postgres"
+    name = "PostgreSQL DB Plugin"
+    supported = {"databases/postgres"}
+    config_file_type = PostgresConfigFile
+
+    def __init__(self):
+        super().__init__(PostgresqlIntrospector())

databao_context_engine/plugins/resources/parquet_chunker.py

@@ -0,0 +1,23 @@
+from dataclasses import dataclass
+
+from databao_context_engine.pluginlib.build_plugin import EmbeddableChunk
+from databao_context_engine.plugins.resources.parquet_introspector import ParquetColumn, ParquetIntrospectionResult
+
+
+@dataclass
+class ParquetColumnChunkContent:
+    file_name: str
+    column: ParquetColumn
+
+
+def build_parquet_chunks(result: ParquetIntrospectionResult) -> list[EmbeddableChunk]:
+    chunks = []
+    for file in result.files:
+        for column in file.columns:
+            chunks.append(
+                EmbeddableChunk(
+                    embeddable_text=f"Column [name = {column.name}, type = {column.type}, number of values = {column.num_values}] in parquet file {file.name}",
+                    content=ParquetColumnChunkContent(file_name=file.name, column=column),
+                )
+            )
+    return chunks

databao_context_engine/plugins/resources/parquet_introspector.py

@@ -0,0 +1,154 @@
+import contextlib
+import logging
+import uuid
+from collections import defaultdict
+from dataclasses import dataclass, replace
+from urllib.parse import urlparse
+
+import duckdb
+from _duckdb import DuckDBPyConnection
+from pydantic import BaseModel, Field
+
+from databao_context_engine.pluginlib.config import DuckDBSecret
+
+parquet_type = "resources/parquet"
+
+logger = logging.getLogger(__name__)
+
+
+class ParquetConfigFile(BaseModel):
+    name: str | None = Field(default=None)
+    type: str = Field(default=parquet_type)
+    url: str = Field(
+        default=type,
+        description="Parquet resource location. Should be a valid URL or a path to a local file. "
+        "Examples: s3://your_bucket/file.parquet, s3://your-bucket/*.parquet, https://some.url/some_file.parquet, ~/path_to/file.parquet",
+    )
+    duckdb_secret: DuckDBSecret | None = None
+
+
+@dataclass
+class ParquetColumn:
+    name: str
+    type: str
+    row_groups: int
+    num_values: int
+    stats_min: str
+    stats_max: str
+    stats_null_count: int | None
+    stats_distinct_count: int | None
+
+
+@dataclass
+class ParquetFile:
+    name: str
+    columns: list[ParquetColumn]
+
+
+@dataclass
+class ParquetIntrospectionResult:
+    files: list[ParquetFile]
+
+
+def generate_create_secret_sql(secret_name, duckdb_secret: DuckDBSecret) -> str:
+    parameters = [("type", duckdb_secret.type)] + list(duckdb_secret.properties.items())
+    return f"""CREATE SECRET {secret_name} (
+    {", ".join([f"{k} {v}" for (k, v) in parameters])}
+);
+"""
+
+
+@contextlib.contextmanager
+def _create_secret(conn: DuckDBPyConnection, duckdb_secret: DuckDBSecret):
+    secret_name = duckdb_secret.name or "gen_secret_" + str(uuid.uuid4()).replace("-", "_")
+    create_secret_sql = generate_create_secret_sql(secret_name, duckdb_secret)
+    try:
+        logger.debug(f"About to create duckdb secret '{secret_name}' with type {duckdb_secret.type}")
+        conn.sql(create_secret_sql)
+        yield conn
+    finally:
+        logger.debug(f"Dropping duckdb secret '{secret_name}'")
+        conn.sql(f"DROP SECRET IF EXISTS {secret_name};")
+
+
+def _resolve_url(file_config: ParquetConfigFile) -> str:
+    parquet_url = urlparse(file_config.url)
+    if parquet_url.scheme == "file":
+        return parquet_url.netloc
+    return file_config.url
+
+
+class ParquetIntrospector:
+    @contextlib.contextmanager
+    def _connect(self, file_config: ParquetConfigFile):
+        duckdb_secret = file_config.duckdb_secret
+        with duckdb.connect() as conn:
+            if duckdb_secret is not None:
+                if duckdb_secret.type == "s3":
+                    conn.execute("INSTALL httpfs;")
+                    conn.execute("LOAD httpfs;")
+                    conn.execute("INSTALL s3;")
+                    conn.execute("LOAD s3;")
+                with _create_secret(conn, duckdb_secret):
+                    yield conn
+            else:
+                yield conn
+
+    def check_connection(self, file_config: ParquetConfigFile) -> None:
+        with self._connect(file_config) as conn:
+            with conn.cursor() as cur:
+                resolved_url = _resolve_url(file_config)
+                cur.execute(f"SELECT * FROM parquet_file_metadata('{resolved_url}') LIMIT 1")
+                columns = [desc[0].lower() for desc in cur.description] if cur.description else []
+                rows = cur.fetchall()
+                parquet_file_metadata = [dict(zip(columns, row)) for row in rows]
+                if not parquet_file_metadata:
+                    raise ValueError(f"No parquet files found by url {resolved_url}")
+                if not parquet_file_metadata or not parquet_file_metadata[0]["file_name"]:
+                    raise ValueError("Parquet resource introspection failed")
+
+    def introspect(self, file_config: ParquetConfigFile) -> ParquetIntrospectionResult:
+        with self._connect(file_config) as conn:
+            with conn.cursor() as cur:
+                resolved_url = _resolve_url(file_config)
+                cur.execute(f"SELECT * from parquet_metadata('{resolved_url}')")
+                cols = [desc[0].lower() for desc in cur.description] if cur.description else []
+                rows = cur.fetchall()
+                file_metas = [dict(zip(cols, row)) for row in rows]
+
+                columns_per_file: dict[str, dict[int, ParquetColumn]] = defaultdict(defaultdict)
+                for file_meta in file_metas:
+                    file_name = file_meta["file_name"]
+                    column_id = file_meta["column_id"]
+                    column_name = file_meta["path_in_schema"]
+                    column_type = file_meta.get("type") or ""
+                    num_values: int = file_meta["num_values"]
+                    stats_min = file_meta.get("stats_min") or ""
+                    stats_max = file_meta.get("stats_max") or ""
+                    stats_null_count: int | None = file_meta.get("stats_null_count")
+                    stats_distinct_count: int | None = file_meta.get("stats_distinct_count")
+
+                    columns: dict[int, ParquetColumn] = columns_per_file[file_name]
+                    column: ParquetColumn | None = columns.get(column_id)
+                    if column:
+                        columns[column_id] = replace(
+                            column, num_values=column.num_values + num_values, row_groups=column.row_groups + 1
+                        )
+                    else:
+                        columns[column_id] = ParquetColumn(
+                            name=column_name,
+                            type=column_type,
+                            row_groups=1,
+                            num_values=num_values,
+                            stats_min=stats_min,
+                            stats_max=stats_max,
+                            stats_null_count=stats_null_count,
+                            stats_distinct_count=stats_distinct_count,
+                        )
+
+                return ParquetIntrospectionResult(
+                    files=[
+                        ParquetFile(file_name, columns=list(columns.values()))
+                        for (file_name, columns) in columns_per_file.items()
+                    ]
+                )

databao_context_engine/plugins/snowflake_db_plugin.py

@@ -0,0 +1,12 @@
+from databao_context_engine.plugins.base_db_plugin import BaseDatabasePlugin
+from databao_context_engine.plugins.databases.snowflake_introspector import SnowflakeConfigFile, SnowflakeIntrospector
+
+
+class SnowflakeDbPlugin(BaseDatabasePlugin[SnowflakeConfigFile]):
+    id = "jetbrains/snowflake"
+    name = "Snowflake DB Plugin"
+    supported = {"databases/snowflake"}
+    config_file_type = SnowflakeConfigFile
+
+    def __init__(self):
+        super().__init__(SnowflakeIntrospector())