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

relationalai/config/config.py

@@ -1,289 +0,0 @@
-"""
-Main configuration class for PyRel (YAML-based).
-
-"""
-
-from __future__ import annotations
-
-from abc import ABC
-from typing import Any, TypeVar, overload, Literal, TYPE_CHECKING
-from pydantic import Field, model_validator
-from pydantic_settings import SettingsConfigDict
-
-try:
-    from confocal import BaseConfig
-except ImportError:
-    # Confocal not yet published to PyPI - use pydantic_settings as fallback
-    # Config system not actively used yet, this is just to prevent import errors
-    from pydantic_settings import BaseSettings as BaseConfig  # type: ignore[misc, assignment]
-
-if TYPE_CHECKING:
-    import snowflake.snowpark
-    import duckdb
-
-from .connections import ConnectionConfig, BaseConnection, SnowflakeConnection, DuckDBConnection
-from .config_fields import EngineConfig, DataConfig, CompilerConfig, ModelConfig, ReasonerConfig, DebugConfig
-from .external.dbt_converter import convert_dbt_to_rai
-from .external.snowflake_converter import convert_snowflake_to_rai
-from .external.utils import find_dbt_profiles_file, find_snowflake_config_file
-
-# TypeVar for generic connection retrieval
-T = TypeVar('T', bound=BaseConnection)
-
-
-class _Config(BaseConfig, ABC): # type: ignore for now until we publish confocal
-    """Base configuration class with common fields and methods."""
-
-    active_profile: str | None = Field(
-        default=None,
-        description="Currently active profile name (auto-defaults if only one profile)"
-    )
-
-    profiles: dict[str, dict[str, Any]] = Field(
-        default_factory=dict,
-        alias="profile",
-        description="Profiles available (dev, prod, staging, etc.) - profiles overlay config fields"
-    )
-
-    connections: dict[str, ConnectionConfig] = Field(
-        default_factory=dict,
-        description="Connections available (snowflake, duckdb, etc.)"
-    )
-
-    default_connection: str | None = Field(
-        default=None,
-        description="Default connection name (auto-defaults if only one connection)"
-    )
-
-    use_graph_index: bool = Field(
-        default=True,
-        description="Enabling graph index"
-    )
-    use_direct_access: bool = Field(
-        default=False,
-        description="Use direct access mode"
-    )
-    describe_optimization: str | None = Field(default=None, description="Describe optimization")
-    enable_otel_handler: bool = Field(
-        default=False,
-        description="Enable OpenTelemetry handler"
-    )
-
-    # Nested config objects
-    engine: EngineConfig = Field(
-        default_factory=EngineConfig,
-        description="Engine configuration"
-    )
-    data: DataConfig = Field(
-        default_factory=DataConfig,
-        description="Data loading and streaming configuration"
-    )
-    compiler: CompilerConfig = Field(
-        default_factory=CompilerConfig,
-        description="Compiler configuration"
-    )
-    model: ModelConfig = Field(
-        default_factory=ModelConfig,
-        description="Model configuration"
-    )
-    reasoner: ReasonerConfig = Field(
-        default_factory=ReasonerConfig,
-        description="Reasoner configuration"
-    )
-    debug: DebugConfig = Field(
-        default_factory=DebugConfig,
-        description="Debug configuration"
-    )
-
-    @model_validator(mode='before')
-    @classmethod
-    def set_default_snowflake_authenticator(cls, data: Any):
-        """Set default authenticator for Snowflake connections in raw data."""
-        if not isinstance(data, dict):
-            return data
-
-        def process_connections(connections):
-            if not isinstance(connections, dict):
-                return
-            for conn_data in connections.values():
-                if isinstance(conn_data, dict) and conn_data.get('type') == 'snowflake':
-                    conn_data.setdefault('authenticator', 'username_password')
-
-        # Process top-level connections
-        if 'connections' in data:
-            process_connections(data['connections'])
-
-        # Process connections in profiles (before overlay)
-        if 'profiles' in data and isinstance(data['profiles'], dict):
-            for profile_data in data['profiles'].values():
-                if isinstance(profile_data, dict) and 'connections' in profile_data:
-                    process_connections(profile_data['connections'])
-
-        return data
-
-    @model_validator(mode='after')
-    def validate_connections(self):
-        if not self.connections:
-            raise ValueError("Config must have at least one connection")
-
-        # Auto-set default_connection if only one connection exists and not already set
-        if len(self.connections) == 1 and self.default_connection is None:
-            self.default_connection = next(iter(self.connections.keys()))
-
-        return self
-
-    def get_default_connection(self) -> ConnectionConfig:
-        default_conn_name = self.default_connection
-
-        if default_conn_name is None:
-            if len(self.connections) == 1:
-                default_conn_name = list(self.connections.keys())[0]
-            else:
-                raise ValueError(
-                    f"Multiple connections available but no default_connection specified. "
-                    f"Available: {list(self.connections.keys())}"
-                )
-
-        if default_conn_name not in self.connections:
-            raise ValueError(
-                f"default_connection '{default_conn_name}' not found. "
-                f"Available: {list(self.connections.keys())}"
-            )
-
-        return self.connections[default_conn_name]
-
-    def get_connection(self, connection_type: type[T], name: str | None = None) -> T:
-        if name is None:
-            connection = self.get_default_connection()
-        else:
-            if name not in self.connections:
-                raise ValueError(
-                    f"Connection '{name}' not found. "
-                    f"Available: {list(self.connections.keys())}"
-                )
-            connection = self.connections[name]
-
-        if not isinstance(connection, connection_type):
-            # Provide helpful error message with the actual connection class hierarchy
-            actual_type = type(connection).__name__
-            # Check if it's a Snowflake connection
-            if isinstance(connection, SnowflakeConnection):
-                connection_class = "SnowflakeConnection"
-            elif isinstance(connection, DuckDBConnection):
-                connection_class = "DuckDBConnection"
-            else:
-                connection_class = actual_type
-
-            raise ValueError(
-                f"Connection '{name or self.default_connection}' is not of type {connection_type.__name__}. "
-                f"Got: {connection_class} (actual: {actual_type})"
-            )
-
-        return connection
-
-    @overload
-    def get_session(self, connection_type: type[SnowflakeConnection]) -> snowflake.snowpark.Session: ...
-
-    @overload
-    def get_session(self, connection_type: type[DuckDBConnection]) -> duckdb.DuckDBPyConnection: ...
-
-    @overload
-    def get_session(self, connection_type: None = None) -> snowflake.snowpark.Session | duckdb.DuckDBPyConnection: ...
-
-    def get_session(self, connection_type: type[SnowflakeConnection] | type[DuckDBConnection] | None = None) -> snowflake.snowpark.Session | duckdb.DuckDBPyConnection:
-        if connection_type is None:
-            connection = self.get_default_connection()
-            return connection.get_session()
-        else:
-            connection = self.get_connection(connection_type)
-            return connection.get_session()
-
-
-# =============================================================================
-# Config Source Classes
-# =============================================================================
-
-class RAIConfig(_Config):
-    """Config loaded from raiconfig.yaml."""
-
-    source: Literal["rai"] = Field(default="rai", exclude=True)
-
-    model_config = SettingsConfigDict(
-        yaml_file="raiconfig.yaml",
-        extra="ignore",
-        nested_model_default_partial_update=True,
-    )
-
-
-class ConfigFromSnowflake(_Config):
-    source: Literal["snowflake"] = Field(default="snowflake", exclude=True)
-
-    model_config = SettingsConfigDict(
-        toml_file=find_snowflake_config_file(),
-        extra="ignore",
-        nested_model_default_partial_update=True,
-    )
-
-    @model_validator(mode='before')
-    @classmethod
-    def convert_snowflake_structure(cls, data: Any):
-        if not isinstance(data, dict):
-            return data
-        return convert_snowflake_to_rai(data)
-
-class ConfigFromDBT(_Config):
-    source: Literal["dbt"] = Field(default="dbt", exclude=True)
-
-    model_config = SettingsConfigDict(
-        yaml_file=find_dbt_profiles_file(),
-        extra="ignore",
-        nested_model_default_partial_update=True,
-    )
-
-    @model_validator(mode='before')
-    @classmethod
-    def convert_dbt_structure(cls, data: Any):
-        """Convert DBT profiles.yml format to RAI Config format."""
-        if not isinstance(data, dict):
-            return data
-
-        return convert_dbt_to_rai(data)
-
-
-# =============================================================================
-# Config Factory - Tries each source in order
-# =============================================================================
-
-def Config(**data) -> _Config:
-    """
-    Create Config by trying multiple sources in priority order:
-    1. RAIConfig (raiconfig.yaml) - or direct data if provided
-    2. ConfigFromSnowflake (config.toml) - only if no data provided
-    3. ConfigFromDBT (profiles.yml) - only if no data provided
-    """
-    sources = [
-        ("RAIConfig (raiconfig.yaml)", RAIConfig, True),      # passes **data
-        ("ConfigFromSnowflake (config.toml)", ConfigFromSnowflake, False),
-        ("ConfigFromDBT (profiles.yml)", ConfigFromDBT, False),
-    ]
-
-    errors = []
-
-    for name, source_cls, accepts_data in sources:
-        try:
-            if accepts_data:
-                return source_cls(**data)
-            else:
-                return source_cls()
-        except Exception as e:
-            errors.append((name, str(e)))
-
-    # All failed, create helpful error message
-    error_lines = [f"  ❌ {name}: {error}" for name, error in errors]
-    error_summary = "\n".join(error_lines)
-
-    raise FileNotFoundError(
-        f"Could not load config from any source:\n\n{error_summary}\n\n"
-        f"To fix: Correct one of the config files above, or pass config programmatically:\n"
-        f"   Config(connections={{...}})"
-    )

relationalai/config/config_fields.py

@@ -1,86 +0,0 @@
-"""
-Nested configuration field models for the multi-connection config system.
-
-This module defines nested config models:
-- EngineConfig: Engine execution settings
-- DataConfig: Data loading and streaming settings
-- CompilerConfig: Compiler-specific settings
-- ModelConfig: Model execution settings
-- ReasonerConfig: Reasoner execution settings (including nested ReasonerRuleConfig)
-- DebugConfig: Debug server settings
-
-These models are used in Config (with default instances) and Profile (as optional overrides).
-"""
-
-from __future__ import annotations
-
-from typing import Literal
-
-from pydantic import BaseModel, Field
-
-
-class EngineConfig(BaseModel):
-    name: str | None = Field(default=None, description="Engine name for execution")
-    size: str | None = Field(default=None, description="Size of the engine")
-    auto_suspend_mins: int | None = Field(default=None, description="Auto-suspend engine after N minutes")
-    show_all_sizes: bool = Field(
-        default=False,
-        description="Show all available engine sizes"
-    )
-
-
-class DataConfig(BaseModel):
-    wait_for_stream_sync: bool = Field(default=True, description="Wait for stream synchronization before processing")
-    ensure_change_tracking: bool = Field(default=False, description="Enable change tracking for data modifications")
-    data_freshness_mins: int | None = Field(default=None, description="Data freshness timeout in minutes")
-    query_timeout_mins: int | None = Field(
-        default=None,
-        description="Query timeout in minutes - aborts queries that exceed this duration"
-    )
-    download_url_type: Literal["internal", "external"] | None = Field(
-        default=None,
-        description="Type of download URL for data exports (internal or external)"
-    )
-    check_column_types: bool = Field(
-        default=True,
-        description="Check column types during data loading"
-    )
-
-
-class CompilerConfig(BaseModel):
-    use_monotype_operators: bool = Field(default=False, description="Use monotype operators in compilation")
-    show_corerel_errors: bool = Field(default=True, description="Show CoreRel error messages")
-    dry_run: bool = Field(default=False, description="Run compilation in dry-run mode")
-    inspect_df: bool = Field(default=False, description="Inspect DataFrame during compilation")
-    use_value_types: bool = Field(default=False, description="Use value types in compilation")
-    debug_hidden_keys: bool = Field(default=False, description="Debug hidden keys in compilation")
-    wide_outputs: bool = Field(default=False, description="Use wide output format for query results")
-    strict: bool = Field(default=False, description="Enable strict validation mode")
-
-    # Experimental compiler optimizations
-    use_inlined_intermediates: bool = Field(default=False, description="Use inlined intermediate results")
-    inline_value_maps: bool = Field(default=False, description="Inline value maps in weaver")
-    inline_entity_maps: bool = Field(default=False, description="Inline entity maps in weaver")
-
-
-class ModelConfig(BaseModel):
-    keep: bool = Field(default=False, description="Keep model after execution")
-    isolated: bool = Field(default=True, description="Run model in isolated mode")
-    nowait_durable: bool = Field(default=True, description="Don't wait for durable operations")
-
-
-class ReasonerRuleConfig(BaseModel):
-    use_lqp: bool = Field(default=True, description="Use LQP for reasoner rule execution")
-
-
-class ReasonerConfig(BaseModel):
-    rule: ReasonerRuleConfig = Field(default_factory=ReasonerRuleConfig)
-    use_sql: bool = Field(default=False, description="Use SQL execution instead of LQP")
-
-
-class DebugConfig(BaseModel):
-    enabled: bool = Field(default=True, description="Enable debug mode")
-    host: str | None = Field(default=None, description="Debug server host")
-    port: int = Field(default=8080, description="Debug server port")
-    show_debug_logs: bool = Field(default=False, description="Show debug log messages")
-    show_full_traces: bool = Field(default=False, description="Show full stack traces in errors")

relationalai/config/connections/__init__.py

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

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

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