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

relationalai/clients/resources/snowflake/engine_state_handlers.py

@@ -0,0 +1,309 @@
+"""
+Engine state handlers for auto_create_engine methods using Strategy Pattern.
+
+Each state handler encapsulates the logic for handling a specific engine state
+(PENDING, SUSPENDED, READY, GONE, or None/missing). Handlers are separated
+for sync and async modes since they have different behaviors.
+"""
+from __future__ import annotations
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Dict, Any, cast, Union
+
+from ....errors import EngineNotFoundException, EngineProvisioningFailed, EngineResumeFailed, EngineSizeMismatchWarning
+from ....tools.cli_controls import Spinner
+from ...util import poll_with_specified_overhead
+
+if TYPE_CHECKING:
+    from .snowflake import Resources
+    from ...types import EngineState
+    EngineDict = Union[EngineState, Dict[str, Any]]
+else:
+    EngineDict = Dict[str, Any]
+
+
+@dataclass
+class EngineContext:
+    """Context for engine state handling."""
+    engine_name: str
+    engine_size: str | None
+    headers: Dict | None
+    requested_size: str | None  # Size explicitly requested by user
+    spinner: Spinner | None = None  # For async mode UI updates
+    span: Any = None  # For sync mode debugging span
+
+
+class EngineStateHandler(ABC):
+    """Base class for engine state handlers using Strategy Pattern."""
+
+    @abstractmethod
+    def handles_state(self, state: str | None) -> bool:
+        """Check if this handler can process the given engine state."""
+        pass
+
+    @abstractmethod
+    def handle(self, engine: EngineDict | None, context: EngineContext, resources: 'Resources') -> EngineDict | None:
+        """
+        Handle the engine state and return updated engine dict or None.
+
+        Returns:
+            - Updated engine dict if engine should remain
+            - None if engine should be deleted/recreated
+        """
+        pass
+
+
+# ============================================================================
+# Sync Mode Handlers
+# ============================================================================
+
+class SyncPendingStateHandler(EngineStateHandler):
+    """Handle PENDING state in sync mode - poll until ready."""
+
+    def handles_state(self, state: str | None) -> bool:
+        return state == "PENDING"
+
+    def handle(self, engine: EngineDict | None, context: EngineContext, resources: 'Resources') -> EngineDict | None:
+        if not engine:
+            return None
+
+        # Warn if requested size doesn't match pending engine size
+        if context.requested_size is not None and engine.get("size") != context.requested_size:
+            existing_size = engine.get("size") or ""
+            EngineSizeMismatchWarning(context.engine_name, existing_size, context.requested_size)
+
+        # Poll until engine is ready
+        with Spinner(
+            "Waiting for engine to be initialized",
+            "Engine ready",
+        ):
+            poll_with_specified_overhead(
+                lambda: resources.is_engine_ready(context.engine_name),
+                overhead_rate=0.1,
+                max_delay=0.5,
+                timeout=900
+            )
+
+        # Return updated engine (should be READY now)
+        updated_engine = resources.get_engine(context.engine_name)
+        return cast(EngineDict, updated_engine) if updated_engine else None
+
+
+class SyncSuspendedStateHandler(EngineStateHandler):
+    """Handle SUSPENDED state in sync mode - resume and poll until ready."""
+
+    def handles_state(self, state: str | None) -> bool:
+        return state == "SUSPENDED"
+
+    def handle(self, engine: EngineDict | None, context: EngineContext, resources: 'Resources') -> EngineDict | None:
+        if not engine:
+            return None
+
+        with Spinner(
+            f"Resuming engine '{context.engine_name}'",
+            f"Engine '{context.engine_name}' resumed",
+            f"Failed to resume engine '{context.engine_name}'"
+        ):
+            try:
+                resources.resume_engine_async(context.engine_name, headers=context.headers)
+                poll_with_specified_overhead(
+                    lambda: resources.is_engine_ready(context.engine_name),
+                    overhead_rate=0.1,
+                    max_delay=0.5,
+                    timeout=900
+                )
+            except Exception:
+                raise EngineResumeFailed(context.engine_name)
+
+        # Return updated engine (should be READY now)
+        updated_engine = resources.get_engine(context.engine_name)
+        return cast(EngineDict, updated_engine) if updated_engine else None
+
+
+class SyncReadyStateHandler(EngineStateHandler):
+    """Handle READY state in sync mode - set active and return."""
+
+    def handles_state(self, state: str | None) -> bool:
+        return state == "READY"
+
+    def handle(self, engine: EngineDict | None, context: EngineContext, resources: 'Resources') -> EngineDict | None:
+        if not engine:
+            return None
+
+        # Warn if requested size doesn't match ready engine size
+        if context.requested_size is not None and engine.get("size") != context.requested_size:
+            existing_size = engine.get("size") or ""
+            EngineSizeMismatchWarning(context.engine_name, existing_size, context.requested_size)
+
+        # Cast to EngineState for _set_active_engine
+        if TYPE_CHECKING:
+            from ...types import EngineState
+            resources._set_active_engine(cast(EngineState, engine))
+        else:
+            resources._set_active_engine(engine)  # type: ignore[arg-type]
+        return engine
+
+
+class SyncGoneStateHandler(EngineStateHandler):
+    """Handle GONE state in sync mode - delete and return None to trigger recreation."""
+
+    def handles_state(self, state: str | None) -> bool:
+        return state == "GONE"
+
+    def handle(self, engine: EngineDict | None, context: EngineContext, resources: 'Resources') -> EngineDict | None:
+        if not engine:
+            return None
+
+        try:
+            # "Gone" is abnormal condition when metadata and SF service don't match
+            # Therefore, we have to delete the engine and create a new one
+            # It could be case that engine is already deleted, so we have to catch the exception
+            resources.delete_engine(context.engine_name, headers=context.headers)
+            # After deleting the engine, return None so that we can create a new engine
+            return None
+        except Exception as e:
+            # If engine is already deleted, we will get an exception
+            # We can ignore this exception and create a new engine
+            if isinstance(e, EngineNotFoundException):
+                return None
+            else:
+                raise EngineProvisioningFailed(context.engine_name, e) from e
+
+
+class SyncMissingEngineHandler(EngineStateHandler):
+    """Handle missing engine (None) in sync mode - create synchronously."""
+
+    def handles_state(self, state: str | None) -> bool:
+        return state is None
+
+    def handle(self, engine: EngineDict | None, context: EngineContext, resources: 'Resources') -> EngineDict | None:
+        # This handler is called when engine doesn't exist
+        # Create engine synchronously with spinner
+        with Spinner(
+            f"Auto-creating engine {context.engine_name}",
+            f"Auto-created engine {context.engine_name}",
+            "Engine creation failed",
+        ):
+            resources.create_engine(context.engine_name, size=context.engine_size, headers=context.headers)
+
+        return resources.get_engine(context.engine_name)
+
+
+# ============================================================================
+# Async Mode Handlers
+# ============================================================================
+
+class AsyncPendingStateHandler(EngineStateHandler):
+    """Handle PENDING state in async mode - just update spinner, don't poll."""
+
+    def handles_state(self, state: str | None) -> bool:
+        return state == "PENDING"
+
+    def handle(self, engine: EngineDict | None, context: EngineContext, resources: 'Resources') -> EngineDict | None:
+        if not engine:
+            return None
+
+        # In async mode, just update spinner - use_index will wait for engine to be ready
+        if context.spinner:
+            context.spinner.update_messages({
+                "finished_message": f"Starting engine {context.engine_name}",
+            })
+
+        return engine
+
+
+class AsyncSuspendedStateHandler(EngineStateHandler):
+    """Handle SUSPENDED state in async mode - resume asynchronously."""
+
+    def handles_state(self, state: str | None) -> bool:
+        return state == "SUSPENDED"
+
+    def handle(self, engine: EngineDict | None, context: EngineContext, resources: 'Resources') -> EngineDict | None:
+        if not engine:
+            return None
+
+        if context.spinner:
+            context.spinner.update_messages({
+                "finished_message": f"Resuming engine {context.engine_name}",
+            })
+
+        try:
+            resources.resume_engine_async(context.engine_name)
+        except Exception:
+            raise EngineResumeFailed(context.engine_name)
+
+        return engine
+
+
+class AsyncReadyStateHandler(EngineStateHandler):
+    """Handle READY state in async mode - set active."""
+
+    def handles_state(self, state: str | None) -> bool:
+        return state == "READY"
+
+    def handle(self, engine: EngineDict | None, context: EngineContext, resources: 'Resources') -> EngineDict | None:
+        if not engine:
+            return None
+
+        if context.spinner:
+            context.spinner.update_messages({
+                "finished_message": f"Engine {context.engine_name} initialized",
+            })
+
+        # Cast to EngineState for _set_active_engine
+        if TYPE_CHECKING:
+            from ...types import EngineState
+            resources._set_active_engine(cast(EngineState, engine))
+        else:
+            resources._set_active_engine(engine)  # type: ignore[arg-type]
+        return engine
+
+
+class AsyncGoneStateHandler(EngineStateHandler):
+    """Handle GONE state in async mode - delete and return None to trigger recreation."""
+
+    def handles_state(self, state: str | None) -> bool:
+        return state == "GONE"
+
+    def handle(self, engine: EngineDict | None, context: EngineContext, resources: 'Resources') -> EngineDict | None:
+        if not engine:
+            return None
+
+        if context.spinner:
+            context.spinner.update_messages({
+                "message": f"Restarting engine {context.engine_name}",
+            })
+
+        try:
+            # "Gone" is abnormal condition when metadata and SF service don't match
+            # Therefore, we have to delete the engine and create a new one
+            # It could be case that engine is already deleted, so we have to catch the exception
+            # Set it to None so that we can create a new engine asynchronously
+            resources.delete_engine(context.engine_name)
+            return None
+        except Exception as e:
+            # If engine is already deleted, we will get an exception
+            # We can ignore this exception and create a new engine asynchronously
+            if isinstance(e, EngineNotFoundException):
+                return None
+            else:
+                raise EngineProvisioningFailed(context.engine_name, e) from e
+
+
+class AsyncMissingEngineHandler(EngineStateHandler):
+    """Handle missing engine (None) in async mode - create asynchronously."""
+
+    def handles_state(self, state: str | None) -> bool:
+        return state is None
+
+    def handle(self, engine: EngineDict | None, context: EngineContext, resources: 'Resources') -> EngineDict | None:
+        # This handler is called when engine doesn't exist
+        # Create engine asynchronously
+        resources.create_engine_async(context.engine_name, size=context.engine_size)
+
+        if context.spinner:
+            context.spinner.update_messages({
+                "finished_message": f"Starting engine {context.engine_name}...",
+            })
+        return None  # Engine is being created asynchronously
+

relationalai/clients/resources/snowflake/error_handlers.py

@@ -0,0 +1,199 @@
+"""
+Error handlers for Snowflake Resources using Strategy Pattern.
+
+Each error handler encapsulates the detection logic and exception creation
+for a specific type of error. Handlers are processed in order until one matches.
+"""
+from __future__ import annotations
+import re
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any
+
+from ....errors import (
+    DuoSecurityFailed,
+    EngineNotFoundException,
+    EnginePending,
+    EngineNameValidationException,
+    EngineProvisioningFailed,
+    EngineResumeFailed,
+    RAIAbortedTransactionError,
+    RAIException,
+    SnowflakeAppMissingException,
+    SnowflakeDatabaseException,
+    SnowflakeRaiAppNotStarted,
+)
+
+if TYPE_CHECKING:
+    from .snowflake import ExecContext
+    from .snowflake import Resources
+
+from .util import is_database_issue, is_engine_issue, collect_error_messages
+
+class ErrorHandler(ABC):
+    """Base class for error handlers using Strategy Pattern."""
+
+    @abstractmethod
+    def matches(self, error: Exception, message: str, ctx: 'ExecContext', resources: 'Resources') -> bool:
+        """Check if this handler can process the error."""
+        pass
+
+    @abstractmethod
+    def handle(self, error: Exception, ctx: 'ExecContext', resources: 'Resources') -> Any | None:
+        """Handle the error and either raise an exception or return an alternate result."""
+        pass
+
+
+class DuoSecurityErrorHandler(ErrorHandler):
+    """Handle Duo security authentication errors."""
+
+    def matches(self, error: Exception, message: str, ctx: 'ExecContext', resources: 'Resources') -> bool:
+        messages = collect_error_messages(error)
+        return any("duo security" in msg for msg in messages)
+
+    def handle(self, error: Exception, ctx: 'ExecContext', resources: 'Resources') -> Any | None:
+        raise DuoSecurityFailed(error)
+
+
+class AppMissingErrorHandler(ErrorHandler):
+    """Handle missing RAI app database errors."""
+
+    def matches(self, error: Exception, message: str, ctx: 'ExecContext', resources: 'Resources') -> bool:
+        rai_app = resources.config.get("rai_app_name", "")
+        if not isinstance(rai_app, str):
+            return False
+        pattern = f"database '{rai_app}' does not exist or not authorized."
+        messages = collect_error_messages(error)
+        return any(re.search(pattern.lower(), msg) for msg in messages)
+
+    def handle(self, error: Exception, ctx: 'ExecContext', resources: 'Resources') -> Any | None:
+        rai_app = resources.config.get("rai_app_name", "")
+        assert isinstance(rai_app, str), f"rai_app_name must be a string, not {type(rai_app)}"
+        raise SnowflakeAppMissingException(rai_app, resources.config.get("role"))
+
+
+class DatabaseErrorsHandler(ErrorHandler):
+    """Handle database-related errors from Snowflake/RAI."""
+
+    def matches(self, error: Exception, message: str, ctx: 'ExecContext', resources: 'Resources') -> bool:
+        """Check if error is database-related."""
+        # Use collect_error_messages to get all messages from exception chain
+        messages = collect_error_messages(error)
+        return any(is_database_issue(msg) for msg in messages)
+
+    def handle(self, error: Exception, ctx: 'ExecContext', resources: 'Resources') -> Any | None:
+        """Handle database errors and raise appropriate exception."""
+        raise SnowflakeDatabaseException(error)
+
+
+class EngineErrorsHandler(ErrorHandler):
+    """Handle all engine-related errors from Snowflake/RAI."""
+
+    def matches(self, error: Exception, message: str, ctx: 'ExecContext', resources: 'Resources') -> bool:
+        """Check if error is engine-related."""
+        # Use collect_error_messages to get all messages from exception chain
+        messages = collect_error_messages(error)
+        return any(is_engine_issue(msg) for msg in messages)
+
+    def handle(self, error: Exception, ctx: 'ExecContext', resources: 'Resources') -> Any | None:
+        """Handle engine errors and raise appropriate exception."""
+        # Use collect_error_messages to get all messages from exception chain
+        messages = collect_error_messages(error)
+        engine = resources.get_default_engine_name()
+        assert isinstance(engine, str), f"engine must be a string, not {type(engine)}"
+
+        # Check all collected messages for engine error patterns
+        for message in messages:
+            if "engine is in pending" in message or "engine is provisioning" in message:
+                raise EnginePending(engine)
+            elif "engine not found" in message or "no engines found" in message:
+                raise EngineNotFoundException(engine, str(error))
+            elif "engine was deleted" in message:
+                raise EngineNotFoundException(engine, "Engine was deleted")
+            elif "engine is suspended" in message:
+                raise EngineResumeFailed(engine)
+            elif "create/resume" in message:
+                raise EngineProvisioningFailed(engine, error)
+
+        # Generic engine error - use the original error message
+        raise RAIException(str(error))
+
+
+class ServiceNotStartedErrorHandler(ErrorHandler):
+    """Handle RAI service not started errors."""
+
+    def matches(self, error: Exception, message: str, ctx: 'ExecContext', resources: 'Resources') -> bool:
+        messages = collect_error_messages(error)
+        return any(re.search(r"service has not been started", msg) for msg in messages)
+
+    def handle(self, error: Exception, ctx: 'ExecContext', resources: 'Resources') -> Any | None:
+        rai_app = resources.config.get("rai_app_name", "")
+        assert isinstance(rai_app, str), f"rai_app_name must be a string, not {type(rai_app)}"
+        raise SnowflakeRaiAppNotStarted(rai_app)
+
+
+class TransactionAbortedErrorHandler(ErrorHandler):
+    """Handle transaction aborted errors with problem details."""
+
+    def matches(self, error: Exception, message: str, ctx: 'ExecContext', resources: 'Resources') -> bool:
+        messages = collect_error_messages(error)
+        return any(re.search(r"state:\s*aborted", msg) for msg in messages)
+
+    def handle(self, error: Exception, ctx: 'ExecContext', resources: 'Resources') -> Any | None:
+        # Use collect_error_messages to get all messages from exception chain
+        messages = collect_error_messages(error)
+        # Check all collected messages for transaction ID
+        for message in messages:
+            txn_id_match = re.search(r"id:\s*([0-9a-f\-]+)", message)
+            if txn_id_match:
+                problems = resources.get_transaction_problems(txn_id_match.group(1))
+                if problems:
+                    # Extract problem details (handle both dict and object formats)
+                    for problem in problems:
+                        if isinstance(problem, dict):
+                            type_field = problem.get('TYPE')
+                            message_field = problem.get('MESSAGE')
+                            report_field = problem.get('REPORT')
+                        else:
+                            type_field = problem.TYPE
+                            message_field = problem.MESSAGE
+                            report_field = problem.REPORT
+                        raise RAIAbortedTransactionError(type_field, message_field, report_field)
+        raise RAIException(str(error))
+
+
+class UseIndexRetryErrorHandler(ErrorHandler):
+    """Handle engine/database errors by polling use_index and retrying the execution.
+
+    Intended for UseIndexResources and subclasses. Register this handler *before*
+    the standard Database/Engine error handlers.
+    """
+
+    def matches(self, error: Exception, message: str, ctx: 'ExecContext', resources: 'Resources') -> bool:
+        if ctx.skip_engine_db_error_retry:
+            return False
+        messages = collect_error_messages(error)
+        return any(msg and (is_database_issue(msg) or is_engine_issue(msg)) for msg in messages)
+
+    def handle(self, error: Exception, ctx: 'ExecContext', resources: 'Resources') -> Any | None:
+        poll_use_index = getattr(resources, "_poll_use_index", None)
+        if not callable(poll_use_index):
+            return None
+
+        engine = resources.get_default_engine_name()
+        engine_size = resources.config.get_default_engine_size()
+        assert isinstance(engine, str), f"engine must be a string, not {type(engine)}"
+
+        model = getattr(resources, "database", "")
+        try:
+            poll_use_index(
+                app_name=resources.get_app_name(),
+                sources=resources.sources,
+                model=model,
+                engine_name=engine,
+                engine_size=engine_size,
+            )
+        except EngineNameValidationException as e:
+            raise EngineNameValidationException(engine) from e
+
+        return ctx.re_execute(resources)
+

relationalai/clients/{export_procedure.py.jinja → resources/snowflake/export_procedure.py.jinja}

@@ -243,7 +243,7 @@ def handle(session{{py_inputs}}, save_as_table="", passed_engine=""):
     except Exception as e:
         logging.debug(f"Error occurred: {str(e)}")
         msg = str(e).lower()
-        if "No columns returned".lower() in msg or "Columns of results could not be determined".lower() in msg:
+        if "no columns returned" in msg or "columns of results could not be determined" in msg:
             logging.debug("No results returned - creating empty dataframe")
             return session.createDataFrame([], StructType([{{ py_outs }}]))
         raise e

relationalai/clients/resources/snowflake/resources_factory.py

@@ -0,0 +1,99 @@
+"""
+Factory function for creating Resources instances based on configuration.
+
+This module provides a factory function that selects the appropriate Resources
+class based on configuration settings (platform, use_direct_access, use_graph_index).
+"""
+from __future__ import annotations
+from typing import Union
+from snowflake.snowpark import Session
+
+from ...config import Config
+from ...local import LocalResources
+from . import Resources, DirectAccessResources, UseIndexResources
+from ....tools.constants import USE_DIRECT_ACCESS, USE_GRAPH_INDEX, Generation
+
+
+def create_resources_instance(
+    config: Config | None = None,
+    profile: str | None = None,
+    dry_run: bool = False,
+    connection: Session | None = None,
+    reset_session: bool = False,
+    language: str = "rel",
+    generation: Generation = Generation.QB,
+) -> Union[LocalResources, DirectAccessResources, UseIndexResources, Resources]:
+    """
+    Factory function that creates the appropriate Resources instance based on config.
+
+    This function selects the Resources class based on:
+    1. Platform (local -> LocalResources)
+    2. use_direct_access flag (DirectAccessResources)
+    3. use_graph_index flag (UseIndexResources)
+    4. Default (base Resources)
+
+    Args:
+        config: Configuration object (optional, will create from profile if not provided)
+        profile: Profile name (optional, used if config is not provided)
+        dry_run: Whether to run in dry-run mode (default: False)
+        connection: Optional Snowflake session connection
+        reset_session: Whether to reset the session (default: False)
+        language: Language for the resources instance (default: "rel")
+        generation: Generation for the resources instance (default: Generation.QB)
+
+    Returns:
+        Appropriate Resources instance based on config settings:
+        - LocalResources if platform is "local"
+        - DirectAccessResources if use_direct_access is enabled
+        - UseIndexResources if use_graph_index is enabled
+        - Resources (base) otherwise
+    """
+    # Create config from profile if not provided
+    if config is None:
+        config = Config(profile)
+
+    platform = config.get("platform", "")
+    if platform == "local":
+        return LocalResources(
+            profile=profile,
+            config=config,
+            dry_run=dry_run,
+            generation=generation,
+            connection=connection,
+            reset_session=reset_session,
+            language=language,
+        )
+
+    if config.get("use_direct_access", USE_DIRECT_ACCESS):
+        return DirectAccessResources(
+            profile=profile,
+            config=config,
+            dry_run=dry_run,
+            generation=generation,
+            connection=connection,
+            reset_session=reset_session,
+            language=language,
+        )
+
+    if config.get("use_graph_index", USE_GRAPH_INDEX):
+        return UseIndexResources(
+            profile=profile,
+            config=config,
+            dry_run=dry_run,
+            generation=generation,
+            connection=connection,
+            reset_session=reset_session,
+            language=language,
+        )
+
+    return Resources(
+        profile=profile,
+        config=config,
+        dry_run=dry_run,
+        generation=generation,
+        connection=connection,
+        reset_session=reset_session,
+        language=language,
+    )
+
+