databao-context-engine 0.1.2__py3-none-any.whl → 0.1.3__py3-none-any.whl

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,6 +81,12 @@ BuildPlugin = BuildDatasourcePlugin | BuildFilePlugin
 
 @dataclass(kw_only=True, frozen=True)
 class DatasourceType:
+    """The type of a Datasource.
+
+    Attributes:
+        full_type: The full type of the datasource, in the format `<main_type>/<subtype>`.
+    """
+
     full_type: str
 
     def __post_init__(self):
@@ -92,16 +96,28 @@ class DatasourceType:
 
     @property
     def main_type(self) -> str:
+        """The main type of the datasource, aka the folder in which the config or raw file is located."""
         return self.full_type.split("/")[0]
 
     @property
     def config_folder(self) -> str:
+        """The folder in which the config or raw file is located. This is equivalent to `main_type`."""
         return self.main_type
 
     @property
     def subtype(self) -> str:
+        """The subtype of the datasource. This is the actual type declared in the config file or the raw file's extension."""
         return self.full_type.split("/")[1]
 
     @staticmethod
     def from_main_and_subtypes(main_type: str, subtype: str) -> "DatasourceType":
+        """Create a DatasourceType from its main type and subtype.
+
+        Args:
+            main_type: The main type (aka config folder) of the datasource.
+            subtype: The subtype of the datasource.
+
+        Returns:
+            A DatasourceType instance with the specified main type and subtype.
+        """
         return DatasourceType(full_type=f"{main_type}/{subtype}")

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/databases/athena_introspector.py

@@ -1,22 +1,64 @@
 from __future__ import annotations
 
-from typing import Any, Mapping
+from typing import Any, Annotated
 
 from pyathena import connect
 from pyathena.cursor import DictCursor
-from pydantic import Field
+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_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="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"
-    )
+    type: str = Field(default="athena")
+    connection: AthenaConnectionProperties
 
 
 class AthenaIntrospector(BaseIntrospector[AthenaConfigFile]):
@@ -26,11 +68,7 @@ class AthenaIntrospector(BaseIntrospector[AthenaConfigFile]):
     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)
+        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:
@@ -38,29 +76,28 @@ class AthenaIntrospector(BaseIntrospector[AthenaConfigFile]):
             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))
+        catalog = file_config.connection.catalog or self._resolve_pseudo_catalog_name(file_config)
         return [catalog]
 
     def _connect_to_catalog(self, file_config: AthenaConfigFile, catalog: str):
-        self._connect(file_config)
+        return 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})
+        sql = f"SELECT schema_name, catalog_name FROM {catalog}.information_schema.schemata"
+        return SQLQuery(sql, None)
 
-    # 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)}
+        comps = self._component_queries(catalog, schemas)
         results: dict[str, list[dict]] = {}
 
         for name, q in comps.items():
-            results[name] = self._fetchall_dicts(connection, q.sql, q.params)
+            results[name] = self._fetchall_dicts(connection, q, None)
 
         return IntrospectionModelBuilder.build_schemas_from_components(
             schemas=schemas,
@@ -73,8 +110,32 @@ class AthenaIntrospector(BaseIntrospector[AthenaConfigFile]):
             idx_cols=[],
         )
 
-    def _sql_columns(self, catalog: str, schemas: list[str]) -> SQLQuery:
-        sql = f"""
+    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, 
@@ -85,13 +146,12 @@ class AthenaIntrospector(BaseIntrospector[AthenaConfigFile]):
         FROM 
             {catalog}.information_schema.columns
         WHERE 
-            table_schema IN ({schemas})
+            table_schema IN ({schemas_in})
         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"
@@ -99,3 +159,6 @@ class AthenaIntrospector(BaseIntrospector[AthenaConfigFile]):
     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/databases/base_introspector.py

@@ -121,9 +121,11 @@ class BaseIntrospector[T: SupportsIntrospectionScope](ABC):
 
     @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."""
+        """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/databases/clickhouse_introspector.py

@@ -1,21 +1,34 @@
 from __future__ import annotations
 
-from typing import Any, Mapping
+from typing import Any, Annotated
 
 import clickhouse_connect
-from pydantic import Field
+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_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]):
@@ -24,11 +37,9 @@ 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)
+        return clickhouse_connect.get_client(
+            **file_config.connection.to_clickhouse_kwargs(),
+        )
 
     def _connect_to_catalog(self, file_config: ClickhouseConfigFile, catalog: str):
         return self._connect(file_config)

databao_context_engine/plugins/databases/duckdb_introspector.py

@@ -10,7 +10,7 @@ from databao_context_engine.plugins.databases.introspection_model_builder import
 
 
 class DuckDBConfigFile(BaseDatabaseConfigFile):
-    type: str = Field(default="databases/duckdb")
+    type: str = Field(default="duckdb")
     connection: DuckDBConnectionConfig
 
 

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/databases/mssql_introspector.py

@@ -1,21 +1,36 @@
 from __future__ import annotations
 
-from typing import Any, Mapping
+from typing import Any, Mapping, Annotated
 
 from mssql_python import connect  # type: ignore[import-untyped]
-from pydantic import Field
+from pydantic import Field, BaseModel
 
+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_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]):
@@ -42,22 +57,16 @@ class MSSQLIntrospector(BaseIntrospector[MSSQLConfigFile]):
 
     def _connect(self, file_config: MSSQLConfigFile):
         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)
+        connection_string = self._create_connection_string_for_config(connection.to_mssql_kwargs())
         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_string = self._create_connection_string_for_config(cfg_for_db)
-        return connect(connection_string)
+        cfg = file_config.model_copy(deep=True)
+        cfg.connection.database = catalog
+        return self._connect(cfg)
 
     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]
 

databao_context_engine/plugins/databases/mysql_introspector.py

@@ -1,18 +1,33 @@
-from typing import Any, Mapping
+from typing import Any, Annotated
 
 import pymysql
-from pydantic import Field
+from pydantic import BaseModel, Field
 from pymysql.constants import CLIENT
 
+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_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]):
@@ -21,20 +36,16 @@ 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")
-
         return pymysql.connect(
-            **connection,
+            **file_config.connection.to_pymysql_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))
+        cfg = file_config.model_copy(deep=True)
+        cfg.connection.database = catalog
+        return self._connect(cfg)
 
     def _get_catalogs(self, connection, file_config: MySQLConfigFile) -> list[str]:
         with connection.cursor() as cur:

databao_context_engine/plugins/databases/postgresql_introspector.py

@@ -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
 
 

databao_context_engine/plugins/databases/snowflake_introspector.py

@@ -1,22 +1,57 @@
 from __future__ import annotations
 
-from typing import Any, Dict, List, Mapping
+from typing import Any, Dict, List, Annotated
 
 import snowflake.connector
-from pydantic import Field
+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_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, ConfigPropertyAnnotation(secret=True)]
+
+
+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]):
@@ -28,11 +63,9 @@ class SnowflakeIntrospector(BaseIntrospector[SnowflakeConfigFile]):
 
     def _connect(self, file_config: SnowflakeConfigFile):
         connection = file_config.connection
-        if not isinstance(connection, Mapping):
-            raise ValueError("Invalid YAML config: 'connection' must be a mapping of connection parameters")
         snowflake.connector.paramstyle = "qmark"
         return snowflake.connector.connect(
-            **connection,
+            **connection.to_snowflake_kwargs(),
         )
 
     def _connect_to_catalog(self, file_config: SnowflakeConfigFile, catalog: str):
@@ -41,7 +74,7 @@ class SnowflakeIntrospector(BaseIntrospector[SnowflakeConfigFile]):
         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]