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

databao_context_engine/plugins/databases/introspection_scope.py

@@ -6,10 +6,11 @@ from pydantic import BaseModel, ConfigDict, field_validator, model_validator
 
 
 class ScopeIncludeRule(BaseModel):
-    """
-    Allowlist selector.
-    - catalog: optional glob pattern
-    - schemas: optional list of glob patterns (string also accepted and normalized to a list)
+    """Allowlist selector.
+
+    Attributes:
+        catalog: optional glob pattern
+        schemas: optional list of glob patterns (string also accepted and normalized to a list)
 
     A rule must specify at least one of: catalog, schemas.
     """
@@ -36,11 +37,12 @@ class ScopeIncludeRule(BaseModel):
 
 
 class ScopeExcludeRule(BaseModel):
-    """
-    Denylist selector.
-    - catalog: optional glob pattern
-    - schemas: optional list of glob patterns (string also accepted)
-    - except_schemas: optional list of glob patterns (string also accepted)
+    """Denylist selector.
+
+    Attributes:
+        catalog: optional glob pattern
+        schemas: optional list of glob patterns (string also accepted)
+        except_schemas: optional list of glob patterns (string also accepted)
 
     If a target matches the rule but also matches except_schemas, it is NOT excluded by this rule.
     """

databao_context_engine/plugins/databases/introspection_scope_matcher.py

@@ -10,17 +10,14 @@ from databao_context_engine.plugins.databases.introspection_scope import (
 
 @dataclass(frozen=True)
 class ScopeSelection:
-    """
-    The final catalog+schema scope to introspect.
-    """
+    """The final catalog+schema scope to introspect."""
 
     catalogs: list[str]
     schemas_per_catalog: dict[str, list[str]]
 
 
 class IntrospectionScopeMatcher:
-    """
-    Applies include/exclude rules (glob matching, case-insensitive) to a discovered set of catalogs/schemas.
+    """Applies include/exclude rules (glob matching, case-insensitive) to a discovered set of catalogs/schemas.
 
     Semantics:
     - If include is empty => start from "everything"

databao_context_engine/plugins/{mssql_db_plugin.py → databases/mssql/mssql_db_plugin.py}

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

databao_context_engine/plugins/databases/{mssql_introspector.py → mssql/mssql_introspector.py}

@@ -1,21 +1,36 @@
 from __future__ import annotations
 
-from typing import Any, Mapping
+from typing import Annotated, Any, Mapping
 
 from mssql_python import connect  # type: ignore[import-untyped]
-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, DatabaseTable
 from databao_context_engine.plugins.databases.introspection_model_builder import IntrospectionModelBuilder
 
 
+class MSSQLConnectionProperties(BaseModel):
+    host: Annotated[str, ConfigPropertyAnnotation(default_value="localhost", required=True)]
+    port: int | None = None
+    instance_name: str | None = None
+    database: str | None = None
+    user: str | None = None
+    password: Annotated[str, ConfigPropertyAnnotation(secret=True)]
+    encrypt: str | None = None
+    additional_properties: dict[str, Any] = {}
+
+    def to_mssql_kwargs(self) -> dict[str, Any]:
+        kwargs = self.model_dump(exclude={"additional_properties"}, exclude_none=True)
+        kwargs.update(self.additional_properties)
+        return kwargs
+
+
 class MSSQLConfigFile(BaseDatabaseConfigFile):
-    type: str = Field(default="databases/mssql")
-    connection: dict[str, Any] = Field(
-        description="Connection parameters for the Microsoft Server SQL database. It can contain any of the keys supported by the Microsoft Server connection library"
-    )
+    type: str = Field(default="mssql")
+    connection: MSSQLConnectionProperties
 
 
 class MSSQLIntrospector(BaseIntrospector[MSSQLConfigFile]):
@@ -40,24 +55,18 @@ class MSSQLIntrospector(BaseIntrospector[MSSQLConfigFile]):
     )
     supports_catalogs = True
 
-    def _connect(self, file_config: MSSQLConfigFile):
+    def _connect(self, file_config: MSSQLConfigFile, *, catalog: str | None = None):
         connection = file_config.connection
-        if not isinstance(connection, Mapping):
-            raise ValueError("Invalid YAML config: 'connection' must be a mapping of connection parameters")
-
-        connection_string = self._create_connection_string_for_config(connection)
-        return connect(connection_string)
 
-    def _connect_to_catalog(self, file_config: MSSQLConfigFile, catalog: str):
-        base_cfg = file_config.connection or {}
-        cfg_for_db: dict[str, Any] = dict(base_cfg)
-        cfg_for_db["database"] = catalog
+        connection_kwargs = connection.to_mssql_kwargs()
+        if catalog:
+            connection_kwargs["database"] = catalog
 
-        connection_string = self._create_connection_string_for_config(cfg_for_db)
+        connection_string = self._create_connection_string_for_config(connection_kwargs)
         return connect(connection_string)
 
     def _get_catalogs(self, connection, file_config: MSSQLConfigFile) -> list[str]:
-        database = file_config.connection.get("database")
+        database = file_config.connection.database
         if isinstance(database, str) and database:
             return [database]
 
@@ -413,8 +422,7 @@ class MSSQLIntrospector(BaseIntrospector[MSSQLConfigFile]):
             "trust_server_certificate": "yes" if file_config.get("trust_server_certificate") else None,
         }
 
-        connection_string = ";".join(f"{k}={v}" for k, v in connection_parts.items() if v is not None)
-        return connection_string
+        return ";".join(f"{k}={v}" for k, v in connection_parts.items() if v is not None)
 
     def _fetchall_dicts(self, connection, sql: str, params) -> list[dict]:
         with connection.cursor() as cursor:

databao_context_engine/plugins/{mysql_db_plugin.py → databases/mysql/mysql_db_plugin.py}

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

databao_context_engine/plugins/databases/{mysql_introspector.py → mysql/mysql_introspector.py}

@@ -1,18 +1,33 @@
-from typing import Any, Mapping
+from typing import Annotated, Any
 
 import pymysql
-from pydantic import Field
+from pydantic import BaseModel, Field
 from pymysql.constants import CLIENT
 
-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, DatabaseTable
 from databao_context_engine.plugins.databases.introspection_model_builder import IntrospectionModelBuilder
 
 
+class MySQLConnectionProperties(BaseModel):
+    host: Annotated[str, ConfigPropertyAnnotation(default_value="localhost", required=True)]
+    port: int | None = None
+    database: str | None = None
+    user: str | None = None
+    password: Annotated[str, ConfigPropertyAnnotation(secret=True)]
+    additional_properties: dict[str, Any] = {}
+
+    def to_pymysql_kwargs(self) -> dict[str, Any]:
+        kwargs = self.model_dump(exclude={"additional_properties"}, exclude_none=True)
+        kwargs.update(self.additional_properties)
+        return kwargs
+
+
 class MySQLConfigFile(BaseDatabaseConfigFile):
-    connection: dict[str, Any]
-    type: str = Field(default="databases/mysql")
+    connection: MySQLConnectionProperties
+    type: str = Field(default="mysql")
 
 
 class MySQLIntrospector(BaseIntrospector[MySQLConfigFile]):
@@ -20,22 +35,18 @@ class MySQLIntrospector(BaseIntrospector[MySQLConfigFile]):
 
     supports_catalogs = True
 
-    def _connect(self, file_config: MySQLConfigFile):
-        connection = file_config.connection
-        if not isinstance(connection, Mapping):
-            raise ValueError("Invalid YAML config: 'connection' must be a mapping of connection parameters")
+    def _connect(self, file_config: MySQLConfigFile, *, catalog: str | None = None):
+        connection_kwargs = file_config.connection.to_pymysql_kwargs()
+
+        if catalog:
+            connection_kwargs["database"] = catalog
 
         return pymysql.connect(
-            **connection,
+            **connection_kwargs,
             cursorclass=pymysql.cursors.DictCursor,
             client_flag=CLIENT.MULTI_STATEMENTS | CLIENT.MULTI_RESULTS,
         )
 
-    def _connect_to_catalog(self, file_config: MySQLConfigFile, catalog: str):
-        cfg = dict(file_config.connection or {})
-        cfg["database"] = catalog
-        return self._connect(MySQLConfigFile(connection=cfg))
-
     def _get_catalogs(self, connection, file_config: MySQLConfigFile) -> list[str]:
         with connection.cursor() as cur:
             cur.execute(

databao_context_engine/plugins/databases/postgresql/postgresql_db_plugin.py

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

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
@@ -16,12 +16,12 @@ class PostgresConnectionProperties(BaseModel):
     port: int | None = None
     database: str | None = None
     user: str | None = None
-    password: str | None = None
+    password: Annotated[str, ConfigPropertyAnnotation(secret=True)]
     additional_properties: dict[str, Any] = {}
 
 
 class PostgresConfigFile(BaseDatabaseConfigFile):
-    type: str = Field(default="databases/postgres")
+    type: str = Field(default="postgres")
     connection: PostgresConnectionProperties
 
 
@@ -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}

@@ -1,22 +1,57 @@
 from __future__ import annotations
 
-from typing import Any, Dict, List, Mapping
+from typing import Annotated, Any, Dict, List
 
 import snowflake.connector
-from pydantic import Field
+from pydantic import BaseModel, Field
 from snowflake.connector import DictCursor
 
-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 SnowflakePasswordAuth(BaseModel):
+    password: Annotated[str, ConfigPropertyAnnotation(secret=True)]
+
+
+class SnowflakeKeyPairAuth(BaseModel):
+    private_key_file: str | None = None
+    private_key_file_pwd: str | None = None
+    private_key: Annotated[str | None, ConfigPropertyAnnotation(secret=True)] = None
+
+
+class SnowflakeSSOAuth(BaseModel):
+    authenticator: str = Field(description='e.g. "externalbrowser"')
+
+
+class SnowflakeConnectionProperties(BaseModel):
+    account: Annotated[str, ConfigPropertyAnnotation(required=True)]
+    warehouse: str | None = None
+    database: str | None = None
+    user: str | None = None
+    role: str | None = None
+    auth: SnowflakePasswordAuth | SnowflakeKeyPairAuth | SnowflakeSSOAuth
+    additional_properties: dict[str, Any] = {}
+
+    def to_snowflake_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 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"
-    )
+    type: str = Field(default="snowflake")
+    connection: SnowflakeConnectionProperties
 
 
 class SnowflakeIntrospector(BaseIntrospector[SnowflakeConfigFile]):
@@ -26,22 +61,19 @@ 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
-        if not isinstance(connection, Mapping):
-            raise ValueError("Invalid YAML config: 'connection' must be a mapping of connection parameters")
         snowflake.connector.paramstyle = "qmark"
+        connection_kwargs = connection.to_snowflake_kwargs()
+        if catalog:
+            connection_kwargs["database"] = catalog
+
         return snowflake.connector.connect(
-            **connection,
+            **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.get("database")
+        database = file_config.connection.database
         if database:
             return [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())