relationalai 0.13.4__py3-none-any.whl → 1.0.0a1__py3-none-any.whl

relationalai/__init__.py

@@ -1,262 +1,7 @@
 """
-The RelationalAI Python SDK.
+The RelationalAI Python SDK, also known as PyRel.
 """
 
-from __future__ import annotations
 import importlib.metadata
 
-from typing import cast
-
-from .clients.config import Config, save_config
-from . import clients
-from . import dsl
-from . import debugging
-from . import metamodel
-from . import rel
-from .loaders import csv
-from . import analysis
-from . import tools
-from .util.otel_configuration import configure_otel
-from snowflake.snowpark import Session
-from .errors import RAIException, handle_missing_integration
-from .environments import runtime_env, SessionEnvironment, SnowbookEnvironment
-from relationalai.tools.constants import Generation
-
-import __main__
-
-# Define cfg after all imports since it depends on clients being imported
-cfg = clients.config
-
 __version__ = importlib.metadata.version(__package__ or __name__)
-
-def Model(
-    name: str,
-    *,
-    profile: str | None = None,
-    config: Config | None = None,
-    dry_run: bool | None = False,
-    debug: bool | None = None,
-    debug_host: str | None = None,
-    debug_port: int | None = None,
-    connection: Session | None = None,
-    keep_model: bool | None = None,
-    isolated: bool | None = None,
-    nowait_durable: bool | None = None,
-    use_package_manager: bool | None = None,
-    ensure_change_tracking: bool | None = None,
-    enable_otel_handler: bool | None = None,
-    format: str = "default",
-):
-    config = config or Config(profile=profile)
-    if use_package_manager is not None:
-        config.set("use_package_manager", use_package_manager)
-    if ensure_change_tracking is not None:
-        config.set("ensure_change_tracking", ensure_change_tracking)
-
-    if isinstance(runtime_env, SessionEnvironment):
-        connection = runtime_env.configure_session(config, connection)
-
-    if debug is None:
-        config_debug = config.get("debug", True)
-        if isinstance(config_debug, dict):
-            debug = True
-        elif isinstance(config_debug, bool):
-            debug = config_debug
-        else:
-            raise Exception("Invalid value specified for `debug`, expected `true` or `false`.")
-
-    if debug_host is None:
-        # Our get function isn't robust to allowing `debug = true/false` or `[debug]\n  port=...`
-        # Went with the lowest impact solve for now which is handling it locally.
-        try:
-            debug_host = config.get("debug.host", None)
-        except AttributeError:
-            pass
-
-    if debug_port is None:
-        try:
-            config_debug_port = config.get("debug.port", 8080)
-            if not isinstance(config_debug_port, int):
-                raise Exception("Invalid value specified for `debug.port`, expected `int`.")
-            debug_port = config_debug_port
-        except AttributeError:
-            pass
-
-    if debug and not runtime_env.remote:
-        from relationalai.tools.debugger_client import start_debugger_session
-        start_debugger_session(config, host=debug_host, port=debug_port)
-
-    main_path = getattr(__main__, "__file__", None)
-    debugging.create_program_span_if_not_exists(main_path, config)
-
-    if not config.file_path:
-        if cfg.legacy_config_exists():
-            message = (
-                "Use `rai init` to migrate your configuration file "
-                "to the new format (raiconfig.toml)"
-            )
-        else:
-            message = "No configuration file found. Please run `rai init` to create one."
-        raise Exception(message)
-    if config.get("platform", None) is None:
-        config.set("platform", "snowflake")
-    platform = config.get("platform")
-    if platform != "snowflake" and connection is not None:
-        raise ValueError("The `connection` parameter is only supported with the Snowflake platform")
-    if dry_run is None:
-        dry_run = config.get_bool("compiler.dry_run", False)
-    if keep_model is None:
-        keep_model = config.get_bool("model.keep", False)
-    if isolated is None:
-        isolated = config.get_bool("model.isolated", True)
-    if nowait_durable is None:
-        nowait_durable = config.get_bool("model.nowait_durable", True)
-    if enable_otel_handler is None:
-        enable_otel_handler = config.get_bool("enable_otel_handler", False)
-
-    try:
-        if platform == "azure":
-            from relationalai.clients.resources.azure.azure import Graph
-            from .util.otel_handler import disable_otel_handling, is_otel_initialized
-            model = Graph(
-                name,
-                profile=profile,
-                config=config,
-                dry_run=dry_run,
-                isolated=isolated,
-                keep_model=keep_model,
-                format=format,
-            )
-            if is_otel_initialized:
-                disable_otel_handling()
-        elif platform == "snowflake":
-            try:
-                from relationalai.clients.resources.snowflake import Graph
-            except ImportError as e:
-                # Provide a helpful error message for Snowflake notebook environments
-                if isinstance(runtime_env, SnowbookEnvironment):
-                    raise ImportError(
-                        "Failed to import relationalai.clients.resources.snowflake. "
-                        "This may indicate that the relationalai.zip package structure is incomplete. "
-                        "Please ensure the zip file includes the full clients/resources/snowflake directory structure."
-                    ) from e
-                raise
-            model = Graph(
-                name,
-                profile=profile,
-                config=config,
-                dry_run=dry_run,
-                isolated=isolated,
-                connection=connection,
-                keep_model=keep_model,
-                nowait_durable=nowait_durable,
-                format=format,
-            )
-
-            configure_otel(enable_otel_handler, config, model._client.resources)
-
-        else:
-            raise Exception(f"Unknown platform: {platform}")
-    except RAIException as e:
-        raise e.clone(config) from None
-    except Exception as e:
-        handle_missing_integration(e)
-        raise e
-    return model
-
-
-def Resources(
-    profile: str | None = None,
-    config: Config | None = None,
-    connection: Session | None = None,
-    # TODO: This is required because creating a unified Snowflake session is not possible. Ticket here: https://app.snowflake.com/support/case/01038599
-    reset_session: bool = False,
-    generation: Generation | None = Generation.V0,
-):
-    config = config or Config(profile)
-    platform = config.get("platform", "snowflake")
-    if platform == "azure":
-        from relationalai.clients.resources.azure.azure import Resources
-        return Resources(config=config)
-    elif platform == "snowflake":
-        try:
-            from relationalai.clients.resources.snowflake.resources_factory import create_resources_instance
-        except ImportError as e:
-            # Provide a helpful error message for Snowflake notebook environments
-            from relationalai.environments import runtime_env, SnowbookEnvironment
-            if isinstance(runtime_env, SnowbookEnvironment):
-                raise ImportError(
-                    "Failed to import relationalai.clients.resources.snowflake. "
-                    "This may indicate that the relationalai.zip package structure is incomplete. "
-                    "Please ensure the zip file includes the full clients/resources/snowflake directory structure."
-                ) from e
-            raise
-        return create_resources_instance(
-            config=config,
-            profile=profile,
-            connection=connection,
-            reset_session=reset_session,
-            generation=generation or Generation.V0,
-            dry_run=False,
-            language="rel",
-        )
-    elif platform == "local":
-        from relationalai.clients.resources.snowflake.resources_factory import create_resources_instance
-        return create_resources_instance(
-            config=config,
-            profile=profile,
-            connection=connection,
-            reset_session=reset_session,
-            generation=generation or Generation.V0,
-            dry_run=False,
-            language="rel",
-        )
-    else:
-        raise Exception(f"Unknown platform: {platform}")
-
-def Provider(
-    profile: str | None = None,
-    config: Config | None = None,
-    connection: Session | None = None,
-    generation: Generation | None = Generation.V0,
-):
-    resources = Resources(profile, config, connection, generation=generation)
-    platform = resources.config.get("platform", "snowflake")
-    if platform == "azure":
-        from relationalai.clients.resources.azure.azure import Resources as AzureResources, Provider as AzureProvider
-        resources = cast(AzureResources, resources)
-        return AzureProvider(
-            resources=resources
-        )
-    elif platform == "snowflake":
-        try:
-            from relationalai.clients.resources.snowflake import Resources as SnowflakeResources, Provider as SnowflakeProvider
-        except ImportError as e:
-            # Provide a helpful error message for Snowflake notebook environments
-            from relationalai.environments import runtime_env, SnowbookEnvironment
-            if isinstance(runtime_env, SnowbookEnvironment):
-                raise ImportError(
-                    "Failed to import relationalai.clients.resources.snowflake. "
-                    "This may indicate that the relationalai.zip package structure is incomplete. "
-                    "Please ensure the zip file includes the full clients/resources/snowflake directory structure."
-                ) from e
-            raise
-        resources = cast(SnowflakeResources, resources)
-        return SnowflakeProvider(
-            resources=resources,
-            generation=generation
-        )
-    elif platform == "local":
-        import relationalai.clients.local
-        resources = cast(relationalai.clients.local.LocalResources, resources)
-        return relationalai.clients.local.LocalProvider(
-            resources=resources
-        )
-    else:
-        raise Exception(f"Unknown platform: {platform}")
-
-
-def Graph(name:str, dry_run:bool=False):
-    return Model(name, profile=None, dry_run=dry_run)
-
-__all__ = ['Model', 'Config', 'Resources', 'Provider', 'dsl', 'rel', 'debugging', 'metamodel', 'csv', 'analysis', 'tools', 'save_config']

relationalai/config/__init__.py

@@ -0,0 +1,56 @@
+"""
+RelationalAI Multi-Connection Configuration System
+
+This package contains the new Pydantic-based configuration system with:
+- Multi-connection support (Snowflake, DuckDB)
+- Multi-profile support (dev, prod, staging, etc.)
+- Discriminated unions for type-safe connection handling
+- Auto-loading from multiple sources (raiconfig.yaml, Snowflake connections, DBT profiles)
+- Environment variable templating support (YAML)
+- Profile override system (profile values override config defaults)
+- Type-safe validation via Pydantic
+"""
+
+from .config import Config, _Config, RAIConfig, ConfigFromDBT, ConfigFromSnowflake
+from .config_fields import EngineConfig, DataConfig, CompilerConfig, ModelConfig, ReasonerConfig, DebugConfig
+from .connections import (
+    ConnectionConfig,
+    SnowflakeConnection,
+    DuckDBConnection,
+    UsernamePasswordAuth,
+    UsernamePasswordMFAAuth,
+    ExternalBrowserAuth,
+    JWTAuth,
+    OAuthAuth,
+    ProgrammaticAccessTokenAuth,
+)
+
+__all__ = [
+    # Main classes
+    "Config",
+    "_Config",
+    "RAIConfig",
+    "ConfigFromDBT",
+    "ConfigFromSnowflake",
+
+    # Nested config models
+    "EngineConfig",
+    "DataConfig",
+    "CompilerConfig",
+    "ModelConfig",
+    "ReasonerConfig",
+    "DebugConfig",
+
+    # Connection types
+    "ConnectionConfig",
+    "SnowflakeConnection",
+    "DuckDBConnection",
+
+    # Snowflake authenticators
+    "UsernamePasswordAuth",
+    "UsernamePasswordMFAAuth",
+    "ExternalBrowserAuth",
+    "JWTAuth",
+    "OAuthAuth",
+    "ProgrammaticAccessTokenAuth",
+]

relationalai/config/config.py

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

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