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

relationalai/config/connections/snowflake.py

@@ -1,243 +0,0 @@
-"""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

@@ -1,17 +0,0 @@
-"""
-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

@@ -1,101 +0,0 @@
-"""
-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

@@ -1,93 +0,0 @@
-"""
-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

@@ -1,41 +0,0 @@
-"""
-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

@@ -1,85 +0,0 @@
-"""
-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

relationalai/config/external/utils.py

@@ -1,19 +0,0 @@
-import os
-
-
-def find_dbt_profiles_file() -> str | None:
-    paths = [
-        "./profiles.yml",
-        os.path.expanduser("~/.dbt/profiles.yml"),
-    ]
-    for path in paths:
-        if os.path.exists(path):
-            return path
-    return None
-
-
-def find_snowflake_config_file() -> str | None:
-    path = os.path.expanduser("~/.snowflake/config.toml")
-    if os.path.exists(path):
-        return path
-    return None

relationalai/semantics/backends/lqp/annotations.py

@@ -1,11 +0,0 @@
-from relationalai.semantics.frontend.base import Field, Library
-
-library = Library("lqp")
-
-Symbol = library.Type("Symbol")
-
-export = library.Relation("export", [])
-external = library.Relation("external", [])
-from_cdc = library.Relation("from_cdc", [])
-adhoc = library.Relation("adhoc", [])
-track = library.Relation("track", [Field.input("library", Symbol), Field.input("relation", Symbol)])