relationalai 0.13.0__py3-none-any.whl → 0.13.0.dev0__py3-none-any.whl

relationalai/config/connections/__init__.py

@@ -0,0 +1,46 @@
+"""
+Connection types for the multi-connection config system.
+
+This module exports:
+- BaseConnection: Base class for all connections
+- Snowflake authenticators: UsernamePasswordAuth, UsernamePasswordMFAAuth, etc.
+- DuckDBConnection: DuckDB connection
+- SnowflakeConnection: Discriminated union of Snowflake authenticators
+- ConnectionConfig: Top-level discriminated union (Snowflake | DuckDB)
+"""
+
+from __future__ import annotations
+
+from typing import Annotated, Union
+from pydantic import Field
+
+from .base import BaseConnection
+from .snowflake import (
+    UsernamePasswordAuth,
+    UsernamePasswordMFAAuth,
+    ExternalBrowserAuth,
+    JWTAuth,
+    OAuthAuth,
+    ProgrammaticAccessTokenAuth,
+    SnowflakeConnection,
+    SnowflakeAuthenticator,
+)
+from .duckdb import DuckDBConnection
+
+ConnectionConfig = Annotated[
+    Union[SnowflakeAuthenticator, DuckDBConnection],
+    Field(discriminator="type")
+]
+
+__all__ = [
+    "BaseConnection",
+    "UsernamePasswordAuth",
+    "UsernamePasswordMFAAuth",
+    "ExternalBrowserAuth",
+    "JWTAuth",
+    "OAuthAuth",
+    "ProgrammaticAccessTokenAuth",
+    "SnowflakeConnection",
+    "DuckDBConnection",
+    "ConnectionConfig",
+]

relationalai/config/connections/base.py

@@ -0,0 +1,23 @@
+"""Base connection protocol for all connection types."""
+
+from __future__ import annotations
+
+from abc import ABC
+from typing import Any
+from pydantic import BaseModel, ConfigDict, PrivateAttr
+
+
+class BaseConnection(BaseModel, ABC):
+    # Note: `type` field is NOT defined here, subclasses must define it
+    # with specific Literal types for discriminated unions to work correctly
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    # Private field for caching session (excluded from serialization)
+    _cached_session: Any = PrivateAttr(default=None)
+
+    def get_session(self) -> Any:
+        raise NotImplementedError("Subclasses must implement get_session()")
+
+    def clear_session_cache(self) -> None:
+        self._cached_session = None

relationalai/config/connections/duckdb.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+from typing import Literal, Any, TYPE_CHECKING
+from pydantic import Field
+
+if TYPE_CHECKING:
+    import duckdb
+
+from .base import BaseConnection
+
+
+class DuckDBConnection(BaseConnection):
+    type: Literal["duckdb"] = "duckdb"
+
+    path: str = Field(..., description="Path to DuckDB database file (or ':memory:' for in-memory)")
+    read_only: bool = Field(False, description="Open database in read-only mode")
+    config: dict[str, Any] | None = Field(None, description="Additional DuckDB configuration options")
+
+    def get_session(self) -> duckdb.DuckDBPyConnection:
+        import duckdb
+
+        if self._cached_session is None:
+            self._cached_session = duckdb.connect(
+                database=self.path,
+                read_only=self.read_only,
+                config=self.config or {}
+            )
+
+        return self._cached_session

relationalai/config/connections/snowflake.py

@@ -0,0 +1,243 @@
+"""Snowflake connection configurations with discriminated authenticators."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Literal, Any, Annotated, Dict, Union, cast, TYPE_CHECKING
+from pydantic import Field, SecretStr, field_validator
+
+if TYPE_CHECKING:
+    import snowflake.snowpark
+
+from .base import BaseConnection
+
+
+class SnowflakeConnectionBase(BaseConnection, ABC):
+    """Base class for all Snowflake connection authenticators with common fields."""
+
+    type: Literal["snowflake"] = "snowflake"
+
+    # Required fields common to all Snowflake authenticators
+    account: str = Field(..., description="Snowflake account identifier")
+    warehouse: str = Field(..., description="Snowflake warehouse name")
+
+    # Optional fields common to all Snowflake authenticators
+    rai_app_name: str = Field(default="RELATIONALAI", description="RAI application name in Snowflake")
+    role: str | None = Field(default=None, description="Snowflake role")
+    database: str | None = Field(default=None, description="Default database")
+    schema_: str | None = Field(default=None, alias="schema", description="Default schema")
+
+    @abstractmethod
+    def _get_connection_params(self) -> Dict[str, Any]:
+        """Return authenticator-specific connection parameters."""
+        pass
+
+    def get_session(self) -> snowflake.snowpark.Session:
+        from snowflake.snowpark import Session
+
+        if self._cached_session is None:
+            connection_params = self._get_connection_params()
+            self._add_common_params(connection_params)
+            self._cached_session = Session.builder.configs(connection_params).create()
+
+        return self._cached_session
+
+    def _add_common_params(self, connection_params: Dict[str, Any]) -> None:
+        if self.role:
+            connection_params["role"] = self.role
+        if self.database:
+            connection_params["database"] = self.database
+        if self.schema_:
+            connection_params["schema"] = self.schema_
+
+
+class UsernamePasswordAuth(SnowflakeConnectionBase):
+    authenticator: Literal["username_password"] = "username_password"
+
+    # Required Snowflake fields
+    user: str = Field(..., description="Snowflake username")
+    password: str | SecretStr = Field(..., description="Snowflake password")
+
+    @field_validator('password', mode='before')
+    @classmethod
+    def convert_password_to_secret(cls, v: Any) -> SecretStr:
+        if isinstance(v, str):
+            return SecretStr(v)
+        return v
+
+    def _get_connection_params(self) -> Dict[str, Any]:
+        return {
+            "user": self.user,
+            "password": cast(SecretStr, self.password).get_secret_value(),
+            "account": self.account,
+            "warehouse": self.warehouse,
+            "authenticator": self.authenticator,
+        }
+
+
+class UsernamePasswordMFAAuth(SnowflakeConnectionBase):
+    authenticator: Literal["username_password_mfa"] = "username_password_mfa"
+
+    # Required Snowflake fields
+    user: str = Field(..., description="Snowflake username")
+    password: str | SecretStr = Field(..., description="Snowflake password")
+    passcode: str = Field(..., description="MFA passcode (6-digit code)")
+
+    @field_validator('password', mode='before')
+    @classmethod
+    def convert_password_to_secret(cls, v: Any) -> SecretStr:
+        """Convert plain str to SecretStr for internal use."""
+        if isinstance(v, str):
+            return SecretStr(v)
+        return v
+
+    def _get_connection_params(self) -> Dict[str, Any]:
+        return {
+            "user": self.user,
+            "password": cast(SecretStr, self.password).get_secret_value(),
+            "account": self.account,
+            "warehouse": self.warehouse,
+            "authenticator": self.authenticator,
+            "passcode": self.passcode,
+        }
+
+
+class ExternalBrowserAuth(SnowflakeConnectionBase):
+    authenticator: Literal["externalbrowser"] = "externalbrowser"
+
+    # Required Snowflake fields
+    user: str = Field(..., description="Snowflake username")
+
+    def _get_connection_params(self) -> Dict[str, Any]:
+        return {
+            "user": self.user,
+            "account": self.account,
+            "warehouse": self.warehouse,
+            "authenticator": self.authenticator,
+        }
+
+
+class JWTAuth(SnowflakeConnectionBase):
+    authenticator: Literal["jwt"] = "jwt"
+
+    # Required Snowflake fields
+    user: str = Field(..., description="Snowflake username")
+    private_key_path: str = Field(..., description="Path to private key file for JWT")
+
+    # Optional Snowflake fields
+    private_key_passphrase: str | SecretStr | None = Field(None, description="Passphrase for private key")
+
+    @field_validator('private_key_passphrase', mode='before')
+    @classmethod
+    def convert_passphrase_to_secret(cls, v: Any) -> SecretStr | None:
+        """Convert plain str to SecretStr for internal use."""
+        if v is None:
+            return None
+        if isinstance(v, str):
+            return SecretStr(v)
+        return v
+
+    def _get_connection_params(self) -> Dict[str, Any]:
+        raise NotImplementedError("JWTAuth uses custom get_session() implementation")
+
+    def get_session(self) -> snowflake.snowpark.Session:
+        if self._cached_session is None:
+            from snowflake.snowpark import Session
+            from cryptography.hazmat.backends import default_backend
+            from cryptography.hazmat.primitives import serialization
+
+            # Load private key
+            with open(self.private_key_path, "rb") as key_file:
+                if self.private_key_passphrase:
+                    p_key = serialization.load_pem_private_key(
+                        key_file.read(),
+                        password=cast(SecretStr, self.private_key_passphrase).get_secret_value().encode(),
+                        backend=default_backend()
+                    )
+                else:
+                    p_key = serialization.load_pem_private_key(
+                        key_file.read(),
+                        password=None,
+                        backend=default_backend()
+                    )
+
+            pkb = p_key.private_bytes(
+                encoding=serialization.Encoding.DER,
+                format=serialization.PrivateFormat.PKCS8,
+                encryption_algorithm=serialization.NoEncryption()
+            )
+
+            connection_params: Dict[str, Any] = {
+                "user": self.user,
+                "account": self.account,
+                "warehouse": self.warehouse,
+                "private_key": pkb,
+            }
+
+            self._add_common_params(connection_params)
+            self._cached_session = Session.builder.configs(connection_params).create()
+
+        return self._cached_session
+
+
+class OAuthAuth(SnowflakeConnectionBase):
+    authenticator: Literal["oauth"] = "oauth"
+
+    # Required Snowflake fields
+    token: str | SecretStr = Field(..., description="OAuth access token")
+
+    @field_validator('token', mode='before')
+    @classmethod
+    def convert_token_to_secret(cls, v: Any) -> SecretStr:
+        """Convert plain str to SecretStr for internal use."""
+        if isinstance(v, str):
+            return SecretStr(v)
+        return v
+
+    def _get_connection_params(self) -> Dict[str, Any]:
+        return {
+            "account": self.account,
+            "warehouse": self.warehouse,
+            "authenticator": "oauth",
+            "token": cast(SecretStr, self.token).get_secret_value(),
+        }
+
+
+class ProgrammaticAccessTokenAuth(SnowflakeConnectionBase):
+    authenticator: Literal["programmatic_access_token"] = "programmatic_access_token"
+
+    # Required Snowflake fields
+    token: str | SecretStr = Field(..., description="Programmatic Access Token")
+
+    @field_validator('token', mode='before')
+    @classmethod
+    def convert_token_to_secret(cls, v: Any) -> SecretStr:
+        """Convert plain str to SecretStr for internal use."""
+        if isinstance(v, str):
+            return SecretStr(v)
+        return v
+
+    def _get_connection_params(self) -> Dict[str, Any]:
+        return {
+            "account": self.account,
+            "warehouse": self.warehouse,
+            "authenticator": "PROGRAMMATIC_ACCESS_TOKEN",
+            "token": cast(SecretStr, self.token).get_secret_value(),
+        }
+
+
+# Union type for all Snowflake authenticators (for discriminated unions in Pydantic)
+SnowflakeAuthenticator = Annotated[
+    Union[
+        UsernamePasswordAuth,
+        UsernamePasswordMFAAuth,
+        ExternalBrowserAuth,
+        JWTAuth,
+        OAuthAuth,
+        ProgrammaticAccessTokenAuth
+    ],
+    Field(discriminator="authenticator")
+]
+
+# Export both the base class (for isinstance checks and type hints) and the union (for Pydantic validation)
+SnowflakeConnection = SnowflakeConnectionBase

relationalai/config/external/__init__.py

@@ -0,0 +1,17 @@
+"""
+External config format models and converters.
+
+This module provides Pydantic models for external configuration formats
+(DBT profiles.yml, Snowflake config.toml) and converters to RAI Config format.
+"""
+
+from .dbt_models import DBTSnowflakeOutput, DBTDuckDBOutput, DBTProfile
+from .snowflake_models import SnowflakeConnection, SnowflakeConfigFile
+
+__all__ = [
+    "DBTSnowflakeOutput",
+    "DBTDuckDBOutput",
+    "DBTProfile",
+    "SnowflakeConnection",
+    "SnowflakeConfigFile",
+]

relationalai/config/external/dbt_converter.py

@@ -0,0 +1,101 @@
+"""
+Converter for DBT profiles.yml to RAI Config format.
+"""
+
+from __future__ import annotations
+
+import os
+import re
+from typing import Any
+
+from .dbt_models import DBTProfile, DBTProfilesFile
+
+
+def render_jinja_value(value: str) -> str:
+    """
+    Render Jinja2 env_var() templates in DBT profiles.
+
+    Supports: {{ env_var('VAR_NAME') }} or {{ env_var("VAR_NAME") }}
+    """
+    pattern = r"\{\{\s*env_var\(['\"]([^'\"]+)['\"]\)\s*\}\}"
+
+    def replace_env_var(match):
+        var_name = match.group(1)
+        env_value = os.environ.get(var_name)
+        if env_value is None:
+            raise ValueError(
+                f"Environment variable '{var_name}' referenced in profiles.yml is not set"
+            )
+        return env_value
+
+    return re.sub(pattern, replace_env_var, value)
+
+
+def process_value(value: Any) -> Any:
+    if isinstance(value, str):
+        return render_jinja_value(value)
+    elif isinstance(value, dict):
+        return {k: process_value(v) for k, v in value.items()}
+    elif isinstance(value, list):
+        return [process_value(v) for v in value]
+    else:
+        return value
+
+
+def convert_dbt_to_rai(
+    dbt_profiles: dict[str, Any],
+    profile_name: str | None = None,
+    target_name: str | None = None
+) -> dict[str, Any]:
+    profiles_file = DBTProfilesFile(**dbt_profiles)
+    dbt_profiles = profiles_file.model_dump()
+
+    dbt_profiles = process_value(dbt_profiles)
+
+    if not profile_name:
+        profile_name = os.environ.get("DBT_PROFILE")
+
+    if not profile_name:
+        profile_names = [k for k in dbt_profiles.keys() if k != "config" and not k.startswith("config_")]
+        if not profile_names:
+            raise ValueError("No profiles found in profiles.yml")
+        profile_name = profile_names[0]
+
+    if profile_name not in dbt_profiles:
+        available = [k for k in dbt_profiles.keys() if k != "config"]
+        raise ValueError(
+            f"Profile '{profile_name}' not found. Available profiles: {available}"
+        )
+
+    profile_data = dbt_profiles[profile_name]
+    profile = DBTProfile(**profile_data)
+
+    # outputs existence validated by DBTProfile model
+    # Get target (from parameter, env var, target field, or 'dev')
+    if not target_name:
+        target_name = os.environ.get("DBT_TARGET")
+    if not target_name:
+        target_name = profile.target
+    if not target_name:
+        # Try 'dev' first, then first output
+        if "dev" in profile.outputs:
+            target_name = "dev"
+        else:
+            target_name = next(iter(profile.outputs.keys()))
+
+    if target_name not in profile.outputs:
+        available = list(profile.outputs.keys())
+        raise ValueError(
+            f"Target '{target_name}' not found in profile '{profile_name}'. "
+            f"Available targets: {available}"
+        )
+
+    output = profile.outputs[target_name]
+
+    connection = output.convert()
+    connection_name = output.provider
+
+    return {
+        "connections": {connection_name: connection},
+        "default_connection": connection_name,
+    }

relationalai/config/external/dbt_models.py

@@ -0,0 +1,93 @@
+"""
+Pydantic models for DBT profiles.yml structure.
+
+These models validate the structure of DBT profiles.yml files before
+converting them to RAI Config format.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Literal, ClassVar
+from pydantic import BaseModel, ConfigDict, Field, model_validator
+
+
+def normalize_snowflake_authenticator(authenticator: str | None) -> str:
+    if not authenticator:
+        return "username_password"
+
+    auth = authenticator.lower()
+
+    mapping = {
+        "snowflake": "username_password",
+        "username_password_mfa": "username_password_mfa",
+        "externalbrowser": "externalbrowser",
+        "jwt": "jwt",
+        "oauth": "oauth",
+        "snowflake_jwt": "jwt",
+    }
+
+    return mapping.get(auth, "username_password")
+
+
+class DBTSnowflakeOutput(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    provider: ClassVar[str] = "snowflake"
+
+    type: Literal["snowflake"]
+    account: str
+    user: str
+    password: str | None = None
+    warehouse: str
+    role: str | None = None
+    database: str | None = None
+    schema_: str | None = Field(default=None, alias="schema")
+    authenticator: str | None = None
+    private_key_path: str | None = None
+    private_key_passphrase: str | None = None
+    token: str | None = None
+
+    def convert(self) -> dict[str, Any]:
+        connection_dict = self.model_dump(
+            exclude_none=True,
+            by_alias=True,
+            exclude={"provider"}
+        )
+
+        connection_dict["authenticator"] = normalize_snowflake_authenticator(self.authenticator)
+
+        return connection_dict
+
+
+class DBTDuckDBOutput(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    provider: ClassVar[str] = "duckdb"
+
+    type: Literal["duckdb"]
+    path: str
+    read_only: bool | None = None
+    config: dict[str, Any] | None = None
+
+    def convert(self) -> dict[str, Any]:
+        return self.model_dump(
+            exclude_none=True,
+            exclude={"provider"}
+        )
+
+
+class DBTProfile(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    target: str | None = None
+    outputs: dict[str, DBTSnowflakeOutput | DBTDuckDBOutput]
+
+    @model_validator(mode='after')
+    def validate_outputs(self):
+        if not self.outputs:
+            raise ValueError("DBT profile must have at least one output in 'outputs' section")
+        return self
+
+
+class DBTProfilesFile(BaseModel):
+    model_config = ConfigDict(extra="allow")

relationalai/config/external/snowflake_converter.py

@@ -0,0 +1,41 @@
+"""
+Converter for Snowflake config.toml to RAI Config format.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from .snowflake_models import SnowflakeConfigFile
+
+
+def convert_snowflake_to_rai(
+    sf_config: dict[str, Any],
+    connection_name: str | None = None
+) -> dict[str, Any]:
+    config_file = SnowflakeConfigFile(**sf_config)
+
+    # Determine connection to use (default_connection_name validity checked by model)
+    if not connection_name:
+        if config_file.default_connection_name:
+            connection_name = config_file.default_connection_name
+        elif "default" in config_file.connections:
+            connection_name = "default"
+        else:
+            connection_name = next(iter(config_file.connections.keys()))
+
+    if connection_name not in config_file.connections:
+        available = list(config_file.connections.keys())
+        raise ValueError(
+            f"Connection '{connection_name}' not found. Available connections: {available}"
+        )
+
+    connection = config_file.connections[connection_name]
+
+    # Convert connection using the model's convert() method
+    connection_dict = connection.convert()
+
+    return {
+        "connections": {"snowflake": connection_dict},
+        "default_connection": "snowflake",
+    }

relationalai/config/external/snowflake_models.py

@@ -0,0 +1,85 @@
+"""
+Pydantic models for Snowflake config.toml structure.
+
+These models validate the structure of Snowflake config.toml files before
+converting them to RAI Config format.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+from pydantic import BaseModel, ConfigDict, Field, model_validator
+
+
+def normalize_authenticator(authenticator: str | None) -> str:
+    if not authenticator:
+        return "username_password"
+
+    auth = authenticator.lower()
+
+    mapping = {
+        "snowflake": "username_password",
+        "username_password_mfa": "username_password_mfa",
+        "externalbrowser": "externalbrowser",
+        "snowflake_jwt": "jwt",
+        "oauth": "oauth",
+        "programmatic_access_token": "programmatic_access_token",
+    }
+
+    return mapping.get(auth, "username_password")
+
+
+class SnowflakeConnection(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    account: str
+    user: str
+    password: str | None = None
+    warehouse: str
+    role: str | None = None
+    database: str | None = None
+    schema_: str | None = Field(default=None, alias="schema")
+    authenticator: str | None = None
+    private_key_path: str | None = None
+    private_key_passphrase: str | None = None
+    token: str | None = None
+    passcode: str | None = None
+    
+    
+    @model_validator(mode='before')
+    def check_required_fields(cls, values):
+        required_fields = ['account', 'user', 'warehouse']
+        missing_fields = [field for field in required_fields if field not in values or values[field] is None]
+        if missing_fields:
+            raise ValueError(f"Missing required fields in Snowflake connection: {', '.join(missing_fields)}")
+        return values
+
+    def convert(self) -> dict[str, Any]:
+        connection_dict = self.model_dump(
+            exclude_none=True,
+            by_alias=True
+        )
+
+        connection_dict["type"] = "snowflake"
+        connection_dict["authenticator"] = normalize_authenticator(self.authenticator)
+
+        return connection_dict
+
+
+class SnowflakeConfigFile(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    default_connection_name: str | None = None
+    connections: dict[str, SnowflakeConnection]
+    
+    @model_validator(mode='before')
+    def check_connections(cls, values):
+        if 'connections' not in values:
+            raise ValueError("Snowflake config.toml is missing 'connections' section")
+        return values
+
+    @model_validator(mode='before')
+    def check_default_connection(cls, values):
+        if 'default_connection_name' not in values:
+            raise ValueError("Snowflake config.toml is missing 'default_connection_name' field")
+        return values