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

databao_context_engine/pluginlib/build_plugin.py

@@ -0,0 +1,107 @@
+from dataclasses import dataclass
+from io import BufferedReader
+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
+    """
+
+    embeddable_text: str
+    """
+    The text to embed as a vector for search usage
+    """
+    content: Any
+    """
+    The content to return as a response when the embeddings has been selected in a search
+    """
+
+
+class BaseBuildPlugin(Protocol):
+    id: str
+    name: str
+
+    def supported_types(self) -> set[str]: ...
+
+    """
+    Returns the list of all supported types for this plugin.
+    If the plugin supports multiple types, they should check the type given in the `full_type` argument when `execute` is called.
+    """
+
+    def divide_context_into_chunks(self, context: Any) -> list[EmbeddableChunk]: ...
+
+    """
+    A method dividing the data source context into meaningful chunks that will be used when searching the context from an AI prompt.
+    """
+
+
+@runtime_checkable
+class BuildDatasourcePlugin[T](BaseBuildPlugin, Protocol):
+    config_file_type: type[T]
+
+    def build_context(self, full_type: str, datasource_name: str, file_config: T) -> Any: ...
+
+    """
+    The method that will be called when a config file has been found for a data source supported by this plugin.
+    """
+
+    def check_connection(self, full_type: str, datasource_name: str, file_config: T) -> None:
+        """
+        Checks 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
+        """
+        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.
+    """
+
+    config_file_type: type[dict[str, Any]] = dict[str, Any]
+
+
+@runtime_checkable
+class BuildFilePlugin(BaseBuildPlugin, Protocol):
+    def build_file_context(self, full_type: str, file_name: str, file_buffer: BufferedReader) -> Any: ...
+
+    """
+    The method that will be called when a file has been found as a data source supported by this plugin.
+    """
+
+
+class NotSupportedError(RuntimeError):
+    """Exception raised by methods not supported by a plugin"""
+
+
+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]
+
+    @property
+    def config_folder(self) -> str:
+        return self.main_type
+
+    @property
+    def subtype(self) -> str:
+        return self.full_type.split("/")[1]
+
+    @staticmethod
+    def from_main_and_subtypes(main_type: str, subtype: str) -> "DatasourceType":
+        return DatasourceType(full_type=f"{main_type}/{subtype}")

databao_context_engine/pluginlib/config.py

@@ -0,0 +1,37 @@
+from dataclasses import dataclass
+from typing import Any, Protocol, runtime_checkable
+
+from pydantic import BaseModel, Field
+
+
+class DuckDBSecret(BaseModel):
+    name: str | None = Field(default=None)
+    type: str = Field(
+        description="DuckDB secret type. Examples: s3, postgres, iceberg, etc. See https://duckdb.org/docs/stable/configuration/secrets_manager#types-of-secrets"
+    )
+    properties: dict[str, Any] = Field(
+        default={},
+        description="Key/Value pairs which will be used to create a duckdb secret. "
+        "See https://duckdb.org/docs/stable/configuration/secrets_manager",
+    )
+
+
+@dataclass(kw_only=True)
+class ConfigPropertyDefinition:
+    property_key: str
+    required: bool
+    property_type: type | None = str
+    default_value: str | None = None
+    nested_properties: list["ConfigPropertyDefinition"] | None = None
+
+
+@dataclass(kw_only=True)
+class ConfigPropertyAnnotation:
+    required: bool = False
+    default_value: str | None = None
+    ignored_for_config_wizard: bool = False
+
+
+@runtime_checkable
+class CustomiseConfigProperties(Protocol):
+    def get_config_file_properties(self) -> list[ConfigPropertyDefinition]: ...

databao_context_engine/pluginlib/plugin_utils.py

@@ -0,0 +1,68 @@
+import json
+import logging
+import os
+from pathlib import Path
+from typing import Any, Mapping
+
+from pydantic import TypeAdapter
+
+from databao_context_engine.pluginlib.build_plugin import BuildDatasourcePlugin, BuildFilePlugin, DatasourceType
+
+logger = logging.getLogger(__name__)
+
+
+def execute_datasource_plugin(
+    plugin: BuildDatasourcePlugin, datasource_type: DatasourceType, config: Mapping[str, Any], datasource_name: str
+) -> Any:
+    if not isinstance(plugin, BuildDatasourcePlugin):
+        raise ValueError("This method can only execute a BuildDatasourcePlugin")
+
+    validated_config = _validate_datasource_config_file(config, plugin)
+
+    return plugin.build_context(
+        full_type=datasource_type.full_type,
+        datasource_name=datasource_name,
+        file_config=validated_config,
+    )
+
+
+def check_connection_for_datasource(
+    plugin: BuildDatasourcePlugin, datasource_type: DatasourceType, config: Mapping[str, Any], datasource_name: str
+) -> None:
+    if not isinstance(plugin, BuildDatasourcePlugin):
+        raise ValueError("Connection checks can only be performed on BuildDatasourcePlugin")
+
+    validated_config = _validate_datasource_config_file(config, plugin)
+
+    plugin.check_connection(
+        full_type=datasource_type.full_type,
+        datasource_name=datasource_name,
+        file_config=validated_config,
+    )
+
+
+def _validate_datasource_config_file(config: Mapping[str, Any], plugin: BuildDatasourcePlugin) -> Any:
+    return TypeAdapter(plugin.config_file_type).validate_python(config)
+
+
+def execute_file_plugin(plugin: BuildFilePlugin, datasource_type: DatasourceType, file_path: Path) -> Any:
+    with file_path.open("rb") as fh:
+        return plugin.build_file_context(
+            full_type=datasource_type.full_type,
+            file_name=file_path.name,
+            file_buffer=fh,
+        )
+
+
+def generate_json_schema(plugin: BuildDatasourcePlugin, pretty_print: bool = True) -> str | None:
+    if plugin.config_file_type == dict[str, Any]:
+        logger.debug(f"Skipping json schema generation for plugin {plugin.id}: no custom config_file_type provided")
+        return None
+
+    json_schema = TypeAdapter(plugin.config_file_type).json_schema(mode="serialization")
+
+    return json.dumps(json_schema, indent=4 if pretty_print else None)
+
+
+def format_json_schema_for_output(plugin: BuildDatasourcePlugin, json_schema: str) -> str:
+    return os.linesep.join([f"JSON Schema for plugin {plugin.id}:", json_schema])

databao_context_engine/plugins/athena_db_plugin.py

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

databao_context_engine/plugins/base_db_plugin.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+from typing import Any, TypeVar
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from databao_context_engine.pluginlib.build_plugin import (
+    BuildDatasourcePlugin,
+    EmbeddableChunk,
+)
+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
+
+
+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")
+
+
+T = TypeVar("T", bound=BaseDatabaseConfigFile)
+
+
+class BaseDatabasePlugin(BuildDatasourcePlugin[T]):
+    name: str
+    supported: set[str]
+
+    def __init__(self, introspector: BaseIntrospector):
+        self._introspector = introspector
+
+    def supported_types(self) -> set[str]:
+        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
+
+    def check_connection(self, full_type: str, datasource_name: str, file_config: T) -> None:
+        self._introspector.check_connection(file_config)
+
+    def divide_context_into_chunks(self, context: Any) -> list[EmbeddableChunk]:
+        return build_database_chunks(context)

databao_context_engine/plugins/clickhouse_db_plugin.py

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

databao_context_engine/plugins/databases/athena_introspector.py

@@ -0,0 +1,101 @@
+from __future__ import annotations
+
+from typing import Any, Mapping
+
+from pyathena import connect
+from pyathena.cursor import DictCursor
+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 AthenaConfigFile(BaseDatabaseConfigFile):
+    type: str = Field(default="databases/athena")
+    connection: dict[str, Any] = Field(
+        description="Connection parameters for the Athena database. It can contain any of the keys supported by the Athena connection library"
+    )
+
+
+class AthenaIntrospector(BaseIntrospector[AthenaConfigFile]):
+    _IGNORED_SCHEMAS = {
+        "information_schema",
+    }
+    supports_catalogs = True
+
+    def _connect(self, file_config: AthenaConfigFile):
+        connection = file_config.connection
+        if not isinstance(connection, Mapping):
+            raise ValueError("Invalid YAML config: 'connection' must be a mapping of connection parameters")
+
+        return connect(**connection, 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.get("catalog", self._resolve_pseudo_catalog_name(file_config))
+        return [catalog]
+
+    def _connect_to_catalog(self, file_config: AthenaConfigFile, catalog: str):
+        self._connect(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)
+        catalog = catalogs[0]
+        sql = "SELECT schema_name, catalog_name FROM information_schema.schemata WHERE catalog_name = %(catalog)s"
+        return SQLQuery(sql, {"catalog": catalog})
+
+    # TODO: Incomplete plugin. Awaiting permission access to AWS to properly develop
+    def collect_catalog_model(self, connection, catalog: str, schemas: list[str]) -> list[DatabaseSchema] | None:
+        if not schemas:
+            return []
+
+        comps = {"columns": self._sql_columns(catalog, schemas)}
+        results: dict[str, list[dict]] = {}
+
+        for name, q in comps.items():
+            results[name] = self._fetchall_dicts(connection, q.sql, q.params)
+
+        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 _sql_columns(self, catalog: str, schemas: list[str]) -> SQLQuery:
+        sql = 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})
+        ORDER BY
+            table_schema,
+            table_name,
+            ordinal_position
+        """
+        return SQLQuery(sql, {"schema": schemas})
+
+    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})

databao_context_engine/plugins/databases/base_introspector.py

@@ -0,0 +1,144 @@
+from __future__ import annotations
+
+import logging
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any, Mapping, Protocol, Sequence, Union
+
+from databao_context_engine.plugins.databases.databases_types import (
+    DatabaseCatalog,
+    DatabaseIntrospectionResult,
+    DatabaseSchema,
+)
+from databao_context_engine.plugins.databases.introspection_scope import IntrospectionScope
+from databao_context_engine.plugins.databases.introspection_scope_matcher import IntrospectionScopeMatcher
+
+logger = logging.getLogger(__name__)
+
+
+class SupportsIntrospectionScope(Protocol):
+    introspection_scope: IntrospectionScope | None
+
+
+class BaseIntrospector[T: SupportsIntrospectionScope](ABC):
+    supports_catalogs: bool = True
+    _IGNORED_SCHEMAS: set[str] = {"information_schema"}
+    _SAMPLE_LIMIT: int = 5
+
+    def check_connection(self, file_config: T) -> None:
+        with self._connect(file_config) as connection:
+            self._fetchall_dicts(connection, "SELECT 1 as test", None)
+
+    def introspect_database(self, file_config: T) -> DatabaseIntrospectionResult:
+        scope_matcher = IntrospectionScopeMatcher(
+            file_config.introspection_scope,
+            ignored_schemas=self._ignored_schemas(),
+        )
+
+        with self._connect(file_config) as root_connection:
+            catalogs = self._get_catalogs_adapted(root_connection, file_config)
+
+        discovered_schemas_per_catalog: dict[str, list[str]] = {}
+        for catalog in catalogs:
+            with self._connect_to_catalog(file_config, 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)
+
+        introspected_catalogs: list[DatabaseCatalog] = []
+        for catalog in scope.catalogs:
+            schemas_to_introspect = scope.schemas_per_catalog.get(catalog, [])
+            if not schemas_to_introspect:
+                continue
+
+            with self._connect_to_catalog(file_config, catalog) as catalog_connection:
+                introspected_schemas = self.collect_catalog_model(catalog_connection, catalog, schemas_to_introspect)
+
+                if not introspected_schemas:
+                    continue
+
+                for schema in introspected_schemas:
+                    for table in schema.tables:
+                        table.samples = self._collect_samples_for_table(
+                            catalog_connection, catalog, schema.name, table.name
+                        )
+
+                introspected_catalogs.append(DatabaseCatalog(name=catalog, schemas=introspected_schemas))
+        return DatabaseIntrospectionResult(catalogs=introspected_catalogs)
+
+    def _get_catalogs_adapted(self, connection, file_config: T) -> list[str]:
+        if self.supports_catalogs:
+            return self._get_catalogs(connection, file_config)
+        return [self._resolve_pseudo_catalog_name(file_config)]
+
+    def _sql_list_schemas(self, catalogs: list[str] | None) -> SQLQuery:
+        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)
+
+    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)
+        rows = self._fetchall_dicts(connection, sql_query.sql, sql_query.params)
+
+        schemas: list[str] = []
+        for row in rows:
+            schema_name = row.get("schema_name")
+            if schema_name:
+                schemas.append(schema_name)
+
+        return schemas
+
+    @abstractmethod
+    def collect_catalog_model(self, connection, catalog: str, schemas: list[str]) -> list[DatabaseSchema] | None:
+        raise NotImplementedError
+
+    def _collect_samples_for_table(self, connection, catalog: str, schema: str, table: str) -> list[dict[str, Any]]:
+        samples: list[dict[str, Any]] = []
+        if self._SAMPLE_LIMIT > 0:
+            try:
+                sql_query = self._sql_sample_rows(catalog, schema, table, self._SAMPLE_LIMIT)
+                samples = self._fetchall_dicts(connection, sql_query.sql, sql_query.params)
+            except NotImplementedError:
+                samples = []
+            except Exception as e:
+                logger.warning("Failed to fetch samples for %s.%s (catalog=%s): %s", schema, table, catalog, e)
+                samples = []
+        return samples
+
+    @abstractmethod
+    def _connect(self, file_config: T):
+        raise NotImplementedError
+
+    @abstractmethod
+    def _fetchall_dicts(self, connection, sql: str, params) -> list[dict]:
+        raise NotImplementedError
+
+    @abstractmethod
+    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
+
+    def _resolve_pseudo_catalog_name(self, file_config: T) -> str:
+        return "default"
+
+    def _ignored_schemas(self) -> set[str]:
+        return self._IGNORED_SCHEMAS
+
+
+@dataclass
+class SQLQuery:
+    sql: str
+    params: ParamsType = None
+
+
+ParamsType = Union[Mapping[str, Any], Sequence[Any], None]