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

databao_context_engine/plugin_loader.py

@@ -0,0 +1,110 @@
+from databao_context_engine.introspection.property_extract import get_property_list_from_type
+from databao_context_engine.pluginlib.build_plugin import (
+    BuildDatasourcePlugin,
+    BuildFilePlugin,
+    BuildPlugin,
+    DatasourceType,
+)
+from databao_context_engine.pluginlib.config import ConfigPropertyDefinition, CustomiseConfigProperties
+from databao_context_engine.plugins.plugin_loader import (
+    load_plugins,
+)
+
+
+class DatabaoContextPluginLoader:
+    """Loader for plugins installed in the current environment."""
+
+    def __init__(self, plugins_by_type: dict[DatasourceType, BuildPlugin] | None = None):
+        """Initialize the DatabaoContextEngine.
+
+        Args:
+            plugins_by_type: Override the list of plugins loaded from the environment.
+                Typical usage should not provide this argument and leave it as None.
+        """
+        self._all_plugins_by_type = load_plugins() if plugins_by_type is None else plugins_by_type
+
+    def get_all_supported_datasource_types(self, exclude_file_plugins: bool = False) -> set[DatasourceType]:
+        """Return the list of all supported datasource types.
+
+        Args:
+            exclude_file_plugins: If True, do not return datasource types from plugins that deal with raw files.
+
+        Returns:
+            A set of all DatasourceType supported in the current installation environment.
+        """
+        if exclude_file_plugins:
+            return {
+                datasource_type
+                for (datasource_type, plugin) in self._all_plugins_by_type.items()
+                if not isinstance(plugin, BuildFilePlugin)
+            }
+
+        return set(self._all_plugins_by_type.keys())
+
+    def get_plugin_for_datasource_type(self, datasource_type: DatasourceType) -> BuildPlugin:
+        """Return the plugin able to build a context for the given datasource type.
+
+        Args:
+            datasource_type: The type of datasource for which to retrieve the plugin.
+
+        Returns:
+            The plugin able to build a context for the given datasource type.
+
+        Raises:
+            ValueError: If no plugin is found for the given datasource type.
+        """
+        if datasource_type not in self._all_plugins_by_type:
+            raise ValueError(f"No plugin found for type '{datasource_type.full_type}'")
+
+        return self._all_plugins_by_type[datasource_type]
+
+    def get_config_file_type_for_datasource_type(self, datasource_type: DatasourceType) -> type:
+        """Return the type of the config file for the given datasource type.
+
+        Args:
+            datasource_type: The type of datasource for which to retrieve the config file type.
+
+        Returns:
+            The type of the config file for the given datasource type.
+
+        Raises:
+            ValueError: If no plugin is found for the given datasource type.
+            ValueError: If the plugin does not support config files.
+        """
+        plugin = self.get_plugin_for_datasource_type(datasource_type)
+
+        if isinstance(plugin, BuildDatasourcePlugin):
+            return plugin.config_file_type
+
+        raise ValueError(
+            f'Impossible to get a config file type for datasource type "{datasource_type.full_type}". The corresponding plugin is a {type(plugin).__name__} but should be a BuildDatasourcePlugin'
+        )
+
+    def get_config_file_structure_for_datasource_type(
+        self, datasource_type: DatasourceType
+    ) -> list[ConfigPropertyDefinition]:
+        """Return the property structure of the config file for the given datasource type.
+
+        This can be used to generate a form for the user to fill in the config file.
+
+        Args:
+            datasource_type: The type of datasource for which to retrieve the config file structure.
+
+        Returns:
+            The structure of the config file for the given datasource type.
+                This structure is a list of ConfigPropertyDefinition objects.
+                Each object in the list describes a property of the config file and its potential nested properties.
+
+        Raises:
+            ValueError: If no plugin is found for the given datasource type.
+            ValueError: If the plugin does not support config files.
+        """
+        plugin = self.get_plugin_for_datasource_type(datasource_type)
+
+        if isinstance(plugin, CustomiseConfigProperties):
+            return plugin.get_config_file_properties()
+        if isinstance(plugin, BuildDatasourcePlugin):
+            return get_property_list_from_type(plugin.config_file_type)
+        raise ValueError(
+            f'Impossible to create a config for datasource type "{datasource_type.full_type}". The corresponding plugin is a {type(plugin).__name__} but should be a BuildDatasourcePlugin or CustomiseConfigProperties'
+        )

databao_context_engine/pluginlib/build_plugin.py

@@ -5,9 +5,7 @@ from typing import Any, Protocol, runtime_checkable
 
 @dataclass
 class EmbeddableChunk:
-    """
-    A chunk that will be embedded as a vector and used when searching context from a given AI prompt
-    """
+    """A chunk that will be embedded as a vector and used when searching context from a given AI prompt."""
 
     embeddable_text: str
     """
@@ -48,19 +46,19 @@ class BuildDatasourcePlugin[T](BaseBuildPlugin, Protocol):
     """
 
     def check_connection(self, full_type: str, datasource_name: str, file_config: T) -> None:
-        """
-        Checks whether the configuration to the datasource is working.
+        """Check whether the configuration to the datasource is working.
 
         The function is expected to succeed without a result if the connection is working.
         If something is wrong with the connection, the function should raise an Exception
+
+        Raises:
+            NotSupportedError: If the plugin doesn't support this method.
         """
         raise NotSupportedError("This method is not implemented for this plugin")
 
 
 class DefaultBuildDatasourcePlugin(BuildDatasourcePlugin[dict[str, Any]], Protocol):
-    """
-    Use this as a base class for plugins that don't need a specific config file type.
-    """
+    """Use this as a base class for plugins that don't need a specific config file type."""
 
     config_file_type: type[dict[str, Any]] = dict[str, Any]
 
@@ -75,7 +73,7 @@ class BuildFilePlugin(BaseBuildPlugin, Protocol):
 
 
 class NotSupportedError(RuntimeError):
-    """Exception raised by methods not supported by a plugin"""
+    """Exception raised by methods not supported by a plugin."""
 
 
 BuildPlugin = BuildDatasourcePlugin | BuildFilePlugin
@@ -83,25 +81,10 @@ BuildPlugin = BuildDatasourcePlugin | BuildFilePlugin
 
 @dataclass(kw_only=True, frozen=True)
 class DatasourceType:
-    full_type: str
-
-    def __post_init__(self):
-        type_segments = self.full_type.split("/")
-        if len(type_segments) != 2:
-            raise ValueError(f"Invalid DatasourceType: {self.full_type}")
-
-    @property
-    def main_type(self) -> str:
-        return self.full_type.split("/")[0]
+    """The type of a Datasource.
 
-    @property
-    def config_folder(self) -> str:
-        return self.main_type
-
-    @property
-    def subtype(self) -> str:
-        return self.full_type.split("/")[1]
+    Attributes:
+        full_type: The full type of the datasource, in the format `<main_type>/<subtype>`.
+    """
 
-    @staticmethod
-    def from_main_and_subtypes(main_type: str, subtype: str) -> "DatasourceType":
-        return DatasourceType(full_type=f"{main_type}/{subtype}")
+    full_type: str

databao_context_engine/pluginlib/config.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from dataclasses import dataclass
 from typing import Any, Protocol, runtime_checkable
 
@@ -17,12 +19,23 @@ class DuckDBSecret(BaseModel):
 
 
 @dataclass(kw_only=True)
-class ConfigPropertyDefinition:
+class ConfigUnionPropertyDefinition:
+    property_key: str
+    types: tuple[type, ...]
+    type_properties: dict[type, list[ConfigPropertyDefinition]]
+
+
+@dataclass(kw_only=True)
+class ConfigSinglePropertyDefinition:
     property_key: str
     required: bool
     property_type: type | None = str
     default_value: str | None = None
-    nested_properties: list["ConfigPropertyDefinition"] | None = None
+    nested_properties: list[ConfigPropertyDefinition] | None = None
+    secret: bool = False
+
+
+ConfigPropertyDefinition = ConfigSinglePropertyDefinition | ConfigUnionPropertyDefinition
 
 
 @dataclass(kw_only=True)
@@ -30,6 +43,7 @@ class ConfigPropertyAnnotation:
     required: bool = False
     default_value: str | None = None
     ignored_for_config_wizard: bool = False
+    secret: bool = False
 
 
 @runtime_checkable

databao_context_engine/plugins/{athena_db_plugin.py → databases/athena/athena_db_plugin.py}

@@ -1,11 +1,11 @@
-from databao_context_engine.plugins.base_db_plugin import BaseDatabasePlugin
-from databao_context_engine.plugins.databases.athena_introspector import AthenaConfigFile, AthenaIntrospector
+from databao_context_engine.plugins.databases.athena.athena_introspector import AthenaConfigFile, AthenaIntrospector
+from databao_context_engine.plugins.databases.base_db_plugin import BaseDatabasePlugin
 
 
 class AthenaDbPlugin(BaseDatabasePlugin[AthenaConfigFile]):
     id = "jetbrains/athena"
     name = "Athena DB Plugin"
-    supported = {"databases/athena"}
+    supported = {"athena"}
     config_file_type = AthenaConfigFile
 
     def __init__(self):

databao_context_engine/plugins/databases/athena/athena_introspector.py

@@ -0,0 +1,161 @@
+from __future__ import annotations
+
+from typing import Annotated, Any
+
+from pyathena import connect
+from pyathena.cursor import DictCursor
+from pydantic import BaseModel, Field
+
+from databao_context_engine.pluginlib.config import ConfigPropertyAnnotation
+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 AwsProfileAuth(BaseModel):
+    profile_name: str
+
+
+class AwsIamAuth(BaseModel):
+    aws_access_key_id: Annotated[str, ConfigPropertyAnnotation(secret=True)]
+    aws_secret_access_key: Annotated[str, ConfigPropertyAnnotation(secret=True)]
+    session_token: str | None = None
+
+
+class AwsAssumeRoleAuth(BaseModel):
+    role_arn: str | None = None
+    role_session_name: str | None = None
+    source_profile: str | None = None
+
+
+class AwsDefaultAuth(BaseModel):
+    # Uses environment variables, instance profile, ECS task role
+    pass
+
+
+class AthenaConnectionProperties(BaseModel):
+    region_name: str
+    schema_name: str = "default"
+    catalog: str | None = "awsdatacatalog"
+    work_group: str | None = None
+    s3_staging_dir: str | None = None
+    auth: AwsIamAuth | AwsProfileAuth | AwsDefaultAuth | AwsAssumeRoleAuth
+    additional_properties: dict[str, Any] = {}
+
+    def to_athena_kwargs(self) -> dict[str, Any]:
+        kwargs = self.model_dump(
+            exclude={
+                "additional_properties": True,
+            },
+            exclude_none=True,
+        )
+        auth_fields = kwargs.pop("auth", {})
+        kwargs.update(auth_fields)
+        kwargs.update(self.additional_properties)
+        return kwargs
+
+
+class AthenaConfigFile(BaseDatabaseConfigFile):
+    type: str = Field(default="athena")
+    connection: AthenaConnectionProperties
+
+
+class AthenaIntrospector(BaseIntrospector[AthenaConfigFile]):
+    _IGNORED_SCHEMAS = {
+        "information_schema",
+    }
+    supports_catalogs = True
+
+    def _connect(self, file_config: AthenaConfigFile, *, catalog: str | None = None) -> Any:
+        return connect(**file_config.connection.to_athena_kwargs(), cursor_class=DictCursor)
+
+    def _fetchall_dicts(self, connection, sql: str, params) -> list[dict]:
+        with connection.cursor() as cur:
+            cur.execute(sql, params or {})
+            return cur.fetchall()
+
+    def _get_catalogs(self, connection, file_config: AthenaConfigFile) -> list[str]:
+        catalog = file_config.connection.catalog or self._resolve_pseudo_catalog_name(file_config)
+        return [catalog]
+
+    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)
+        catalog = catalogs[0]
+        sql = f"SELECT schema_name, catalog_name FROM {catalog}.information_schema.schemata"
+        return SQLQuery(sql, 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)
+        results: dict[str, list[dict]] = {}
+
+        for name, q in comps.items():
+            results[name] = self._fetchall_dicts(connection, q, 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=[],
+        )
+
+    def _component_queries(self, catalog: str, schemas: list[str]) -> dict[str, str]:
+        schemas_in = ", ".join(self._quote_literal(s) for s in schemas)
+        return {
+            "relations": self._sql_relations(catalog, schemas_in),
+            "columns": self._sql_columns(catalog, schemas_in),
+        }
+
+    def _sql_relations(self, catalog: str, schemas_in: str) -> str:
+        return f"""
+            SELECT
+                table_schema AS schema_name,
+                table_name,
+                CASE table_type
+                    WHEN 'BASE TABLE' THEN 'table'
+                    WHEN 'VIEW' THEN 'view'
+                    ELSE LOWER(table_type)
+                END AS kind,
+                NULL AS description
+            FROM 
+                {catalog}.information_schema.tables
+            WHERE 
+                table_schema IN ({schemas_in})
+        """
+
+    def _sql_columns(self, catalog: str, schemas_in: str) -> str:
+        return f"""
+        SELECT 
+            table_schema AS schema_name,
+            table_name, 
+            column_name, 
+            ordinal_position, 
+            data_type,
+            is_nullable
+        FROM 
+            {catalog}.information_schema.columns
+        WHERE 
+            table_schema IN ({schemas_in})
+        ORDER BY
+            table_schema,
+            table_name,
+            ordinal_position
+        """
+
+    def _resolve_pseudo_catalog_name(self, file_config: AthenaConfigFile) -> str:
+        return "awsdatacatalog"
+
+    def _sql_sample_rows(self, catalog: str, schema: str, table: str, limit: int) -> SQLQuery:
+        sql = f'SELECT * FROM "{schema}"."{table}" LIMIT %(limit)s'
+        return SQLQuery(sql, {"limit": limit})
+
+    def _quote_literal(self, value: str) -> str:
+        return "'" + str(value).replace("'", "''") + "'"

databao_context_engine/plugins/{base_db_plugin.py → databases/base_db_plugin.py}

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, TypeVar
+from typing import Annotated, Any, TypeVar
 
 from pydantic import BaseModel, ConfigDict, Field
 
@@ -8,6 +8,7 @@ from databao_context_engine.pluginlib.build_plugin import (
     BuildDatasourcePlugin,
     EmbeddableChunk,
 )
+from databao_context_engine.pluginlib.config import ConfigPropertyAnnotation
 from databao_context_engine.plugins.databases.base_introspector import BaseIntrospector
 from databao_context_engine.plugins.databases.database_chunker import build_database_chunks
 from databao_context_engine.plugins.databases.introspection_scope import IntrospectionScope
@@ -17,7 +18,9 @@ class BaseDatabaseConfigFile(BaseModel):
     model_config = ConfigDict(populate_by_name=True)
     name: str | None = Field(default=None)
     type: str
-    introspection_scope: IntrospectionScope | None = Field(default=None, alias="introspection-scope")
+    introspection_scope: Annotated[
+        IntrospectionScope | None, ConfigPropertyAnnotation(ignored_for_config_wizard=True)
+    ] = Field(default=None, alias="introspection-scope")
 
 
 T = TypeVar("T", bound=BaseDatabaseConfigFile)
@@ -34,9 +37,7 @@ class BaseDatabasePlugin(BuildDatasourcePlugin[T]):
         return self.supported
 
     def build_context(self, full_type: str, datasource_name: str, file_config: T) -> Any:
-        introspection_result = self._introspector.introspect_database(file_config)
-
-        return introspection_result
+        return self._introspector.introspect_database(file_config)
 
     def check_connection(self, full_type: str, datasource_name: str, file_config: T) -> None:
         self._introspector.check_connection(file_config)

databao_context_engine/plugins/databases/base_introspector.py

@@ -40,7 +40,7 @@ class BaseIntrospector[T: SupportsIntrospectionScope](ABC):
 
         discovered_schemas_per_catalog: dict[str, list[str]] = {}
         for catalog in catalogs:
-            with self._connect_to_catalog(file_config, catalog) as conn:
+            with self._connect(file_config, catalog=catalog) as conn:
                 discovered_schemas_per_catalog[catalog] = self._list_schemas_for_catalog(conn, catalog)
         scope = scope_matcher.filter_scopes(catalogs, discovered_schemas_per_catalog)
 
@@ -50,7 +50,7 @@ class BaseIntrospector[T: SupportsIntrospectionScope](ABC):
             if not schemas_to_introspect:
                 continue
 
-            with self._connect_to_catalog(file_config, catalog) as catalog_connection:
+            with self._connect(file_config, catalog=catalog) as catalog_connection:
                 introspected_schemas = self.collect_catalog_model(catalog_connection, catalog, schemas_to_introspect)
 
                 if not introspected_schemas:
@@ -74,9 +74,9 @@ class BaseIntrospector[T: SupportsIntrospectionScope](ABC):
         if self.supports_catalogs:
             sql = "SELECT catalog_name, schema_name FROM information_schema.schemata WHERE catalog_name = ANY(%s)"
             return SQLQuery(sql, (catalogs,))
-        else:
-            sql = "SELECT schema_name FROM information_schema.schemata"
-            return SQLQuery(sql, None)
+
+        sql = "SELECT schema_name FROM information_schema.schemata"
+        return SQLQuery(sql, None)
 
     def _list_schemas_for_catalog(self, connection: Any, catalog: str) -> list[str]:
         sql_query = self._sql_list_schemas([catalog] if self.supports_catalogs else None)
@@ -108,7 +108,12 @@ class BaseIntrospector[T: SupportsIntrospectionScope](ABC):
         return samples
 
     @abstractmethod
-    def _connect(self, file_config: T):
+    def _connect(self, file_config: T, *, catalog: str | None = None) -> Any:
+        """Connect to the database.
+
+        If the `catalog` argument is provided, the connection is "scoped" to that catalog. For engines that don’t need a new connection,
+        return a connection with the session set/USE’d to that catalog.
+        """
         raise NotImplementedError
 
     @abstractmethod
@@ -119,12 +124,6 @@ class BaseIntrospector[T: SupportsIntrospectionScope](ABC):
     def _get_catalogs(self, connection, file_config: T) -> list[str]:
         raise NotImplementedError
 
-    @abstractmethod
-    def _connect_to_catalog(self, file_config: T, catalog: str):
-        """Return a connection scoped to `catalog`. For engines that
-        don’t need a new connection, return a connection with the
-        session set/USE’d to that catalog."""
-
     def _sql_sample_rows(self, catalog: str, schema: str, table: str, limit: int) -> SQLQuery:
         raise NotImplementedError
 

databao_context_engine/plugins/{clickhouse_db_plugin.py → databases/clickhouse/clickhouse_db_plugin.py}

@@ -1,5 +1,5 @@
-from databao_context_engine.plugins.base_db_plugin import BaseDatabasePlugin
-from databao_context_engine.plugins.databases.clickhouse_introspector import (
+from databao_context_engine.plugins.databases.base_db_plugin import BaseDatabasePlugin
+from databao_context_engine.plugins.databases.clickhouse.clickhouse_introspector import (
     ClickhouseConfigFile,
     ClickhouseIntrospector,
 )
@@ -8,7 +8,7 @@ from databao_context_engine.plugins.databases.clickhouse_introspector import (
 class ClickhouseDbPlugin(BaseDatabasePlugin[ClickhouseConfigFile]):
     id = "jetbrains/clickhouse"
     name = "Clickhouse DB Plugin"
-    supported = {"databases/clickhouse"}
+    supported = {"clickhouse"}
     config_file_type = ClickhouseConfigFile
 
     def __init__(self):

databao_context_engine/plugins/databases/{clickhouse_introspector.py → clickhouse/clickhouse_introspector.py}

@@ -1,21 +1,34 @@
 from __future__ import annotations
 
-from typing import Any, Mapping
+from typing import Annotated, Any
 
 import clickhouse_connect
-from pydantic import Field
+from pydantic import BaseModel, Field
 
-from databao_context_engine.plugins.base_db_plugin import BaseDatabaseConfigFile
+from databao_context_engine.pluginlib.config import ConfigPropertyAnnotation
+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 ClickhouseConnectionProperties(BaseModel):
+    host: Annotated[str, ConfigPropertyAnnotation(default_value="localhost", required=True)]
+    port: int | None = None
+    database: str | None = None
+    username: str | None = None
+    password: Annotated[str, ConfigPropertyAnnotation(secret=True)]
+    additional_properties: dict[str, Any] = {}
+
+    def to_clickhouse_kwargs(self) -> dict[str, Any]:
+        kwargs = self.model_dump(exclude={"additional_properties"}, exclude_none=True)
+        kwargs.update(self.additional_properties)
+        return kwargs
+
+
 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)"
-    )
+    type: str = Field(default="clickhouse")
+    connection: ClickhouseConnectionProperties
 
 
 class ClickhouseIntrospector(BaseIntrospector[ClickhouseConfigFile]):
@@ -23,15 +36,10 @@ class ClickhouseIntrospector(BaseIntrospector[ClickhouseConfigFile]):
 
     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 _connect(self, file_config: ClickhouseConfigFile, *, catalog: str | None = None):
+        return clickhouse_connect.get_client(
+            **file_config.connection.to_clickhouse_kwargs(),
+        )
 
     def _get_catalogs(self, connection, file_config: ClickhouseConfigFile) -> list[str]:
         return ["clickhouse"]

databao_context_engine/plugins/databases/duckdb/duckdb_db_plugin.py

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

databao_context_engine/plugins/databases/{duckdb_introspector.py → duckdb/duckdb_introspector.py}

@@ -3,19 +3,20 @@ 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_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
+from databao_context_engine.plugins.duckdb_tools import fetchall_dicts
 
 
 class DuckDBConfigFile(BaseDatabaseConfigFile):
-    type: str = Field(default="databases/duckdb")
+    type: str = Field(default="duckdb")
     connection: DuckDBConnectionConfig
 
 
 class DuckDBConnectionConfig(BaseModel):
-    database: str = Field(description="Path to the DuckDB database file")
+    database_path: str = Field(description="Path to the DuckDB database file")
 
 
 class DuckDBIntrospector(BaseIntrospector[DuckDBConfigFile]):
@@ -23,13 +24,10 @@ class DuckDBIntrospector(BaseIntrospector[DuckDBConfigFile]):
     _IGNORED_SCHEMAS = {"information_schema", "pg_catalog"}
     supports_catalogs = True
 
-    def _connect(self, file_config: DuckDBConfigFile):
-        database_path = str(file_config.connection.database)
+    def _connect(self, file_config: DuckDBConfigFile, *, catalog: str | None = None):
+        database_path = str(file_config.connection.database_path)
         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")]
@@ -319,7 +317,4 @@ class DuckDBIntrospector(BaseIntrospector[DuckDBConfigFile]):
 
     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]
+        return fetchall_dicts(cur, sql, params)

databao_context_engine/plugins/databases/introspection_model_builder.py

@@ -5,13 +5,13 @@ from databao_context_engine.plugins.databases.databases_types import (
     CheckConstraint,
     DatabaseColumn,
     DatabasePartitionInfo,
+    DatabaseSchema,
     DatabaseTable,
     DatasetKind,
     ForeignKey,
     ForeignKeyColumnMap,
     Index,
     KeyConstraint,
-    DatabaseSchema,
 )