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

relationalai/clients/resources/snowflake/engine_state_handlers.py

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

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