databao-context-engine 0.1.4.dev1__py3-none-any.whl → 0.1.6__py3-none-any.whl

databao_context_engine/plugins/databases/{postgresql_introspector.py → postgresql/postgresql_introspector.py}

@@ -5,7 +5,7 @@ import asyncpg
 from pydantic import BaseModel, Field
 
 from databao_context_engine.pluginlib.config import ConfigPropertyAnnotation
-from databao_context_engine.plugins.base_db_plugin import BaseDatabaseConfigFile
+from databao_context_engine.plugins.databases.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
@@ -81,12 +81,14 @@ class PostgresqlIntrospector(BaseIntrospector[PostgresConfigFile]):
         if self.supports_catalogs:
             sql = "SELECT catalog_name, schema_name FROM information_schema.schemata WHERE catalog_name = ANY($1)"
             return SQLQuery(sql, (catalogs,))
-        else:
-            sql = "SELECT schema_name FROM information_schema.schemata"
-            return SQLQuery(sql, None)
 
-    def _connect(self, file_config: PostgresConfigFile):
+        sql = "SELECT schema_name FROM information_schema.schemata"
+        return SQLQuery(sql, None)
+
+    def _connect(self, file_config: PostgresConfigFile, *, catalog: str | None = None):
         kwargs = self._create_connection_kwargs(file_config.connection)
+        if catalog:
+            kwargs["database"] = catalog
         return _SyncAsyncpgConnection(kwargs)
 
     def _fetchall_dicts(self, connection: _SyncAsyncpgConnection, sql: str, params) -> list[dict]:
@@ -97,13 +99,7 @@ class PostgresqlIntrospector(BaseIntrospector[PostgresConfigFile]):
         if database is not None:
             return [database]
 
-        rows = connection.fetch_scalar_values("SELECT datname FROM pg_catalog.pg_database WHERE datistemplate = false")
-        return rows
-
-    def _connect_to_catalog(self, file_config: PostgresConfigFile, catalog: str):
-        cfg = file_config.model_copy(deep=True)
-        cfg.connection.database = catalog
-        return self._connect(cfg)
+        return connection.fetch_scalar_values("SELECT datname FROM pg_catalog.pg_database WHERE datistemplate = false")
 
     def collect_catalog_model(
         self, connection: _SyncAsyncpgConnection, catalog: str, schemas: list[str]
@@ -408,10 +404,7 @@ class PostgresqlIntrospector(BaseIntrospector[PostgresConfigFile]):
         }
         connection_parts.update(connection_config.additional_properties)
 
-        connection_string = " ".join(
-            f"{k}={_escape_pg_value(str(v))}" for k, v in connection_parts.items() if v is not None
-        )
-        return connection_string
+        return " ".join(f"{k}={_escape_pg_value(str(v))}" for k, v in connection_parts.items() if v is not None)
 
     def _create_connection_kwargs(self, connection_config: PostgresConnectionProperties) -> dict[str, Any]:
         kwargs: dict[str, Any] = {

databao_context_engine/plugins/databases/snowflake/snowflake_db_plugin.py

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

databao_context_engine/plugins/databases/{snowflake_introspector.py → snowflake/snowflake_introspector.py}

@@ -7,7 +7,7 @@ from pydantic import BaseModel, Field
 from snowflake.connector import DictCursor
 
 from databao_context_engine.pluginlib.config import ConfigPropertyAnnotation
-from databao_context_engine.plugins.base_db_plugin import BaseDatabaseConfigFile
+from databao_context_engine.plugins.databases.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
@@ -20,7 +20,7 @@ class SnowflakePasswordAuth(BaseModel):
 class SnowflakeKeyPairAuth(BaseModel):
     private_key_file: str | None = None
     private_key_file_pwd: str | None = None
-    private_key: Annotated[str, ConfigPropertyAnnotation(secret=True)]
+    private_key: Annotated[str | None, ConfigPropertyAnnotation(secret=True)] = None
 
 
 class SnowflakeSSOAuth(BaseModel):
@@ -61,18 +61,17 @@ class SnowflakeIntrospector(BaseIntrospector[SnowflakeConfigFile]):
     _IGNORED_CATALOGS = {"STREAMLIT_APPS"}
     supports_catalogs = True
 
-    def _connect(self, file_config: SnowflakeConfigFile):
+    def _connect(self, file_config: SnowflakeConfigFile, *, catalog: str | None = None):
         connection = file_config.connection
         snowflake.connector.paramstyle = "qmark"
+        connection_kwargs = connection.to_snowflake_kwargs()
+        if catalog:
+            connection_kwargs["database"] = catalog
+
         return snowflake.connector.connect(
-            **connection.to_snowflake_kwargs(),
+            **connection_kwargs,
         )
 
-    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.database
         if database:

databao_context_engine/plugins/databases/sqlite/sqlite_db_plugin.py

@@ -0,0 +1,12 @@
+from databao_context_engine.plugins.databases.base_db_plugin import BaseDatabasePlugin
+from databao_context_engine.plugins.databases.sqlite.sqlite_introspector import SQLiteConfigFile, SQLiteIntrospector
+
+
+class SQLiteDbPlugin(BaseDatabasePlugin[SQLiteConfigFile]):
+    id = "jetbrains/sqlite"
+    name = "SQLite Plugin"
+    supported = {"sqlite"}
+    config_file_type = SQLiteConfigFile
+
+    def __init__(self):
+        super().__init__(SQLiteIntrospector())

databao_context_engine/plugins/databases/sqlite/sqlite_introspector.py

@@ -0,0 +1,241 @@
+import sqlite3
+
+from pydantic import BaseModel, Field
+
+from databao_context_engine.plugins.databases.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 SQLiteConnectionConfig(BaseModel):
+    database_path: str = Field(description="Path to the SQLite database file")
+
+
+class SQLiteConfigFile(BaseDatabaseConfigFile):
+    type: str = Field(default="sqlite")
+    connection: SQLiteConnectionConfig
+
+
+class SQLiteIntrospector(BaseIntrospector[SQLiteConfigFile]):
+    _IGNORED_SCHEMAS = {"temp", "information_schema"}
+    _PSEUDO_SCHEMA = "main"
+    supports_catalogs = False
+
+    def _connect(self, file_config: SQLiteConfigFile, *, catalog: str | None = None):
+        database_path = str(file_config.connection.database_path)
+        conn = sqlite3.connect(database_path)
+        conn.text_factory = str
+        return conn
+
+    def _get_catalogs(self, connection, file_config: SQLiteConfigFile) -> list[str]:
+        return [self._resolve_pseudo_catalog_name(file_config)]
+
+    def _fetchall_dicts(self, connection, sql: str, params) -> list[dict]:
+        cur = connection.cursor()
+        if params is None:
+            cur.execute(sql)
+        else:
+            cur.execute(sql, params)
+
+        if cur.description is None:
+            return []
+
+        rows = cur.fetchall()
+        out: list[dict] = []
+        for r in rows:
+            if isinstance(r, sqlite3.Row):
+                out.append({k.lower(): r[k] for k in r.keys()})
+            else:
+                cols = [d[0].lower() for d in cur.description]
+                out.append(dict(zip(cols, r)))
+        return out
+
+    def _list_schemas_for_catalog(self, connection, catalog: str) -> list[str]:
+        return [self._PSEUDO_SCHEMA]
+
+    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]] = {name: [] for name in comps}
+
+        for name, sql in comps.items():
+            results[name] = self._fetchall_dicts(connection, sql, None) or []
+
+        return IntrospectionModelBuilder.build_schemas_from_components(
+            schemas=[self._PSEUDO_SCHEMA],
+            rels=results.get("relations", []),
+            cols=results.get("columns", []),
+            pk_cols=results.get("pk", []),
+            uq_cols=results.get("uq", []),
+            checks=[],
+            fk_cols=results.get("fks", []),
+            idx_cols=results.get("idx", []),
+            partitions=[],
+        )
+
+    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(),
+            "fks": self._sql_foreign_keys(),
+            "idx": self._sql_indexes(),
+        }
+
+    def _sql_relations(self) -> str:
+        return f"""
+            SELECT
+                '{self._PSEUDO_SCHEMA}' AS schema_name,
+                m.name AS table_name,
+                CASE m.type
+                    WHEN 'view' THEN 'view'
+                    ELSE 'table'
+                END AS kind,
+                NULL AS description
+            FROM 
+                sqlite_master m
+            WHERE
+                m.type IN ('table', 'view')
+                AND m.name NOT LIKE 'sqlite_%'
+            ORDER BY
+                m.name;
+        """
+
+    def _sql_columns(self) -> str:
+        return f"""
+            SELECT
+                '{self._PSEUDO_SCHEMA}' AS schema_name,
+                m.name AS table_name,
+                c.name AS column_name,
+                (c.cid + 1) AS ordinal_position,
+                COALESCE(c.type,'') AS data_type,
+                CASE 
+                    WHEN c.pk > 0 THEN 0 
+                    WHEN c."notnull" = 0 THEN 1 
+                    ELSE 0 
+                END AS is_nullable,
+                c.dflt_value AS default_expression,
+                CASE
+                    WHEN c.hidden IN (2,3) THEN 'computed'
+                END AS generated,
+                NULL AS description
+            FROM 
+                sqlite_master m
+                JOIN pragma_table_xinfo(m.name) c
+            WHERE 
+                m.type IN ('table','view')
+                AND m.name NOT LIKE 'sqlite_%'
+            ORDER BY 
+                m.name, 
+                c.cid;
+        """
+
+    def _sql_primary_keys(self) -> str:
+        return f"""
+            SELECT
+                '{self._PSEUDO_SCHEMA}' AS schema_name,
+                m.name AS table_name,
+                ('pk_' || m.name) AS constraint_name,
+                c.pk AS position,
+                c.name AS column_name
+            FROM 
+                sqlite_master m
+                JOIN pragma_table_info(m.name) c
+            WHERE
+                m.type = 'table'
+                AND m.name NOT LIKE 'sqlite_%'
+                AND c.pk > 0
+            ORDER BY
+                m.name,
+                c.pk;
+        """
+
+    def _sql_unique(self) -> str:
+        return f"""
+            SELECT
+                '{self._PSEUDO_SCHEMA}' AS schema_name,
+                m.name AS table_name,
+                il.name AS constraint_name,
+                (ii.seqno + 1) AS position,
+                ii.name AS column_name
+            FROM 
+                sqlite_master m
+                JOIN pragma_index_list(m.name) il
+                JOIN pragma_index_info(il.name) ii
+            WHERE
+                m.type = 'table'
+                AND m.name NOT LIKE 'sqlite_%'
+                AND il."unique" = 1
+                AND il.origin = 'u'
+            ORDER BY
+                m.name,
+                il.name,
+                ii.seqno;
+        """
+
+    def _sql_foreign_keys(self) -> str:
+        return f"""
+            SELECT
+                '{self._PSEUDO_SCHEMA}' AS schema_name,
+                m.name AS table_name,
+                ('fk_' || m.name || '_' || fk.id) AS constraint_name,
+                (fk.seq + 1) AS position,
+                fk."from" AS from_column,
+                'main' AS ref_schema,
+                fk."table" AS ref_table,
+                fk."to" AS to_column,
+                lower(fk.on_update) AS on_update,
+                lower(fk.on_delete) AS on_delete,
+                1 AS enforced,
+                1 AS validated
+            FROM sqlite_master m
+            JOIN pragma_foreign_key_list(m.name) fk
+            WHERE
+                m.type = 'table'
+                AND m.name NOT LIKE 'sqlite_%'
+            ORDER BY
+                m.name,
+                fk.id,
+                fk.seq;
+        """
+
+    def _sql_indexes(self) -> str:
+        return f"""
+            SELECT
+                '{self._PSEUDO_SCHEMA}' AS schema_name,
+                m.name AS table_name,
+                il.name AS index_name,
+                (ix.seqno + 1) AS position,
+                ix.name AS expr,
+                il."unique" AS is_unique,
+                NULL AS method,
+                CASE
+                    WHEN il.partial = 1 AND sm.sql IS NOT NULL AND instr(upper(sm.sql), 'WHERE') > 0
+                    THEN trim(substr(sm.sql, instr(upper(sm.sql), 'WHERE') + length('WHERE')))
+                END AS predicate
+            FROM 
+                sqlite_master m
+                JOIN pragma_index_list(m.name) il
+                JOIN pragma_index_xinfo(il.name) ix
+                LEFT JOIN sqlite_master sm ON sm.type = 'index' AND sm.name = il.name
+            WHERE 
+                m.type='table'
+                AND m.name NOT LIKE 'sqlite_%'
+                AND lower(il.origin) = 'c'
+                AND ix.key = 1
+            ORDER BY 
+                m.name, 
+                il.name, 
+                ix.seqno;
+        """
+
+    def _sql_sample_rows(self, catalog: str, schema: str, table: str, limit: int) -> SQLQuery:
+        sql = f"SELECT * FROM {self._quote_ident(table)} LIMIT ?"
+        return SQLQuery(sql, (limit,))
+
+    def _quote_ident(self, ident: str) -> str:
+        return '"' + str(ident).replace('"', '""') + '"'

databao_context_engine/plugins/dbt/dbt_chunker.py

@@ -0,0 +1,47 @@
+from dataclasses import dataclass
+
+from databao_context_engine.pluginlib.build_plugin import EmbeddableChunk
+from databao_context_engine.plugins.dbt.types import DbtColumn, DbtContext, DbtModel
+
+
+@dataclass
+class DbtColumnChunkContent:
+    database_name: str
+    schema_name: str
+    model_name: str
+    column: DbtColumn
+
+
+def build_dbt_chunks(context: DbtContext) -> list[EmbeddableChunk]:
+    chunks = []
+
+    for model in context.models:
+        chunks.append(_create_model_chunk(model))
+
+        for column in model.columns:
+            chunks.append(_create_column_chunk(model, column))
+
+    return chunks
+
+
+def _create_model_chunk(model: DbtModel) -> EmbeddableChunk:
+    return EmbeddableChunk(embeddable_text=_build_model_chunk_text(model), content=model)
+
+
+def _build_model_chunk_text(model: DbtModel) -> str:
+    # TODO: Use description and potentially other infos?
+    return f"Model {model.name} in database {model.database} and schema {model.schema}, with unique id {model.id}"
+
+
+def _create_column_chunk(model: DbtModel, column: DbtColumn) -> EmbeddableChunk:
+    return EmbeddableChunk(
+        embeddable_text=_build_column_chunk_text(model, column),
+        content=DbtColumnChunkContent(
+            database_name=model.database, schema_name=model.schema, model_name=model.name, column=column
+        ),
+    )
+
+
+def _build_column_chunk_text(model: DbtModel, column: DbtColumn) -> str:
+    # TODO: Use description and potentially other infos?
+    return f"Column {column.name} in model {model.id}"

databao_context_engine/plugins/dbt/dbt_context_extractor.py

@@ -0,0 +1,106 @@
+from pathlib import Path
+
+from databao_context_engine.plugins.dbt.types import (
+    DbtColumn,
+    DbtConfigFile,
+    DbtContext,
+    DbtMaterialization,
+    DbtModel,
+)
+from databao_context_engine.plugins.dbt.types_artifacts import (
+    DbtArtifacts,
+    DbtCatalog,
+    DbtCatalogColumn,
+    DbtCatalogNode,
+    DbtManifest,
+    DbtManifestColumn,
+    DbtManifestModel,
+)
+
+
+def check_connection(config_file: DbtConfigFile) -> None:
+    _read_dbt_artifacts(config_file.dbt_target_folder_path.expanduser())
+
+
+def extract_context(config_file: DbtConfigFile) -> DbtContext:
+    artifacts = _read_dbt_artifacts(config_file.dbt_target_folder_path.expanduser())
+
+    return _extract_context_from_artifacts(artifacts)
+
+
+def _read_dbt_artifacts(dbt_target_folder_path: Path) -> DbtArtifacts:
+    if not dbt_target_folder_path.is_dir():
+        raise ValueError(f'Invalid "dbt_target_folder_path": not a directory ({dbt_target_folder_path})')
+
+    # TODO: Check the manifest schema version?
+    manifest_file = dbt_target_folder_path.joinpath("manifest.json")
+    if not manifest_file.is_file():
+        raise ValueError(f'Invalid "dbt_target_folder_path": missing manifest.json file ({manifest_file})')
+
+    manifest = DbtManifest.model_validate_json(manifest_file.read_text())
+
+    catalog_file = dbt_target_folder_path.joinpath("catalog.json")
+    catalog = DbtCatalog.model_validate_json(catalog_file.read_text()) if catalog_file.is_file() else None
+
+    return DbtArtifacts(manifest=manifest, catalog=catalog)
+
+
+def _extract_context_from_artifacts(artifacts: DbtArtifacts) -> DbtContext:
+    manifest_models = [
+        manifest_model
+        for manifest_model in artifacts.manifest.nodes.values()
+        if isinstance(manifest_model, DbtManifestModel)
+    ]
+
+    catalog_nodes = artifacts.catalog.nodes if artifacts.catalog else {}
+
+    # TODO: Extract the stages? Or at least the "highest-level" models (= marts?)
+    # TODO: Extract the constraints
+    # TODO: Organize the models by schemas? Or by stages?
+    return DbtContext(
+        models=[
+            _manifest_model_to_dbt_model(manifest_model, catalog_nodes.get(manifest_model.unique_id, None))
+            for manifest_model in manifest_models
+        ],
+    )
+
+
+def _manifest_model_to_dbt_model(manifest_model: DbtManifestModel, catalog_node: DbtCatalogNode | None) -> DbtModel:
+    catalog_columns = catalog_node.columns if catalog_node else {}
+
+    return DbtModel(
+        id=manifest_model.unique_id,
+        name=manifest_model.name,
+        database=manifest_model.database,
+        schema=manifest_model.schema_,
+        description=manifest_model.description,
+        columns=[
+            _manifest_column_to_dbt_column(manifest_column, catalog_columns.get(manifest_column.name))
+            for manifest_column in manifest_model.columns.values()
+        ],
+        materialization=_manifest_materialization_to_dbt_materializaton(
+            manifest_model.config.materialized if manifest_model.config else None
+        ),
+        primary_key=manifest_model.primary_key,
+        depends_on_nodes=manifest_model.depends_on.get("nodes", []) if manifest_model.depends_on else [],
+    )
+
+
+def _manifest_column_to_dbt_column(
+    manifest_column: DbtManifestColumn, catalog_column: DbtCatalogColumn | None
+) -> DbtColumn:
+    return DbtColumn(
+        name=manifest_column.name,
+        description=manifest_column.description,
+        type=catalog_column.type if catalog_column else manifest_column.data_type,
+    )
+
+
+def _manifest_materialization_to_dbt_materializaton(materialized: str | None) -> DbtMaterialization | None:
+    if materialized is None:
+        return None
+
+    try:
+        return DbtMaterialization(materialized)
+    except ValueError:
+        return None

databao_context_engine/plugins/dbt/dbt_plugin.py

@@ -0,0 +1,25 @@
+from typing import Any
+
+from databao_context_engine import BuildDatasourcePlugin
+from databao_context_engine.pluginlib.build_plugin import EmbeddableChunk
+from databao_context_engine.plugins.dbt.dbt_chunker import build_dbt_chunks
+from databao_context_engine.plugins.dbt.dbt_context_extractor import check_connection, extract_context
+from databao_context_engine.plugins.dbt.types import DbtConfigFile
+
+
+class DbtPlugin(BuildDatasourcePlugin[DbtConfigFile]):
+    id = "jetbrains/dbt"
+    name = "Dbt Plugin"
+    config_file_type = DbtConfigFile
+
+    def supported_types(self) -> set[str]:
+        return {"dbt"}
+
+    def build_context(self, full_type: str, datasource_name: str, file_config: DbtConfigFile) -> Any:
+        return extract_context(file_config)
+
+    def check_connection(self, full_type: str, datasource_name: str, file_config: DbtConfigFile) -> None:
+        check_connection(file_config)
+
+    def divide_context_into_chunks(self, context: Any) -> list[EmbeddableChunk]:
+        return build_dbt_chunks(context)

databao_context_engine/plugins/dbt/types.py

@@ -0,0 +1,44 @@
+from dataclasses import dataclass
+from enum import Enum
+from pathlib import Path
+
+from pydantic import BaseModel, Field
+
+
+class DbtConfigFile(BaseModel):
+    name: str | None = Field(default=None)
+    type: str = Field(default="dbt")
+    dbt_target_folder_path: Path
+
+
+class DbtMaterialization(str, Enum):
+    TABLE = "table"
+    VIEW = "view"
+
+    def __str__(self):
+        return self.value
+
+
+@dataclass(kw_only=True)
+class DbtColumn:
+    name: str
+    type: str | None = None
+    description: str | None = None
+
+
+@dataclass(kw_only=True)
+class DbtModel:
+    id: str
+    name: str
+    database: str
+    schema: str
+    columns: list[DbtColumn]
+    description: str | None = None
+    materialization: DbtMaterialization | None = None
+    primary_key: list[str] | None = None
+    depends_on_nodes: list[str]
+
+
+@dataclass(kw_only=True)
+class DbtContext:
+    models: list[DbtModel]

databao_context_engine/plugins/dbt/types_artifacts.py

@@ -0,0 +1,58 @@
+from dataclasses import dataclass
+from typing import Annotated, Literal
+
+from pydantic import BaseModel, Discriminator, Field
+
+
+class DbtManifestNodeConfig(BaseModel):
+    materialized: str
+
+
+class DbtManifestColumn(BaseModel):
+    name: str
+    description: str | None = None
+    data_type: str | None = None
+
+
+class DbtManifestModel(BaseModel):
+    resource_type: Literal["model"]
+    unique_id: str
+    name: str
+    database: str
+    schema_: str = Field(alias="schema")
+    description: str | None = None
+    config: DbtManifestNodeConfig | None = None
+    columns: dict[str, DbtManifestColumn]
+    depends_on: dict[str, list[str]] | None = None
+    primary_key: list[str] | None = None
+
+
+class DbtManifestOtherNode(BaseModel):
+    resource_type: Literal["seed", "analysis", "test", "operation", "sql_operation", "snapshot"]
+
+
+DbtManifestNode = Annotated[DbtManifestModel | DbtManifestOtherNode, Discriminator("resource_type")]
+
+
+class DbtManifest(BaseModel):
+    nodes: dict[str, DbtManifestNode]
+
+
+class DbtCatalogColumn(BaseModel):
+    name: str
+    type: str
+
+
+class DbtCatalogNode(BaseModel):
+    unique_id: str | None = None
+    columns: dict[str, DbtCatalogColumn]
+
+
+class DbtCatalog(BaseModel):
+    nodes: dict[str, DbtCatalogNode]
+
+
+@dataclass(kw_only=True)
+class DbtArtifacts:
+    manifest: DbtManifest
+    catalog: DbtCatalog | None

databao_context_engine/plugins/{unstructured_files_plugin.py → files/unstructured_files_plugin.py}

@@ -24,7 +24,7 @@ class InternalUnstructuredFilesPlugin(BuildFilePlugin):
         self.tokens_overlap = tokens_overlap or self._DEFAULT_TOKENS_OVERLAP
 
     def supported_types(self) -> set[str]:
-        return {f"files/{extension}" for extension in self._SUPPORTED_FILES_EXTENSIONS}
+        return self._SUPPORTED_FILES_EXTENSIONS
 
     def build_file_context(self, full_type: str, file_name: str, file_buffer: BufferedReader) -> Any:
         file_content = self._read_file(file_buffer)