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

relationalai/clients/resources/snowflake/engine_service.py

@@ -0,0 +1,381 @@
+from __future__ import annotations
+
+import json
+from typing import Any, Dict, List, Protocol, Tuple
+
+from .... import debugging
+from ....errors import EngineProvisioningFailed
+from ...config import Config
+from ...types import EngineState
+from ...util import poll_with_specified_overhead
+
+# ---------------------------------------------------------------------------
+# CONSTANTS
+# ---------------------------------------------------------------------------
+
+# Keep the same placeholder string as `snowflake.py` so `Resources._exec_sql()` will
+# replace it with the actual app name.
+APP_NAME = "___RAI_APP___"
+ENGINE_SCHEMA = "experimental"
+API_SCHEMA = f"{APP_NAME}.{ENGINE_SCHEMA}"
+
+# Cloud-specific engine sizes
+INTERNAL_ENGINE_SIZES = ["XS", "S", "M", "L"]
+ENGINE_SIZES_AWS = ["HIGHMEM_X64_S", "HIGHMEM_X64_M", "HIGHMEM_X64_L"]
+ENGINE_SIZES_AZURE = ["HIGHMEM_X64_S", "HIGHMEM_X64_M", "HIGHMEM_X64_SL"]
+
+# ---------------------------------------------------------------------------
+# Engine types
+# ---------------------------------------------------------------------------
+
+class EngineType:
+    """Engine type constants with descriptions."""
+
+    LOGIC = "LOGIC"
+    SOLVER = "SOLVER"
+    ML = "ML"
+
+    _LABELS = {
+        LOGIC: "Logic",
+        SOLVER: "Prescriptive",
+        ML: "Predictive",
+    }
+
+    _DESCRIPTIONS = {
+        LOGIC: "Logic engine for deductive reasoning and relational queries",
+        SOLVER: "Optimization engine using mathematical solvers for prescriptive reasoning",
+        ML: "Machine learning engine for pattern recognition and predictive reasoning",
+    }
+
+    @classmethod
+    def get_label(cls, type: str) -> str:
+        """Get the user-friendly label for an engine type."""
+        return cls._LABELS.get(type, type)
+
+    @classmethod
+    def get_label_with_value(cls, type: str) -> str:
+        """Get the user-friendly label with the value in parentheses."""
+        label = cls.get_label(type)
+        return f"{label} ({type})"
+
+    @classmethod
+    def get_description(cls, type: str) -> str:
+        """Get the description for an engine type."""
+        return cls._DESCRIPTIONS.get(type, "Unknown engine type")
+
+    @classmethod
+    def is_valid(cls, type: str) -> bool:
+        """Check if an engine type is valid."""
+        return type in cls._DESCRIPTIONS
+
+    @classmethod
+    def get_all_types(cls) -> List[str]:
+        """Get all valid engine types."""
+        return list(cls._DESCRIPTIONS.keys())
+
+
+class _ExecResources(Protocol):
+    """Minimal surface EngineServiceSQL needs from Resources (composition, no mixins)."""
+
+    config: Config
+
+    def _exec(
+        self,
+        code: str,
+        params: Any | None = None,
+        raw: bool = False,
+        help: bool = True,
+        skip_engine_db_error_retry: bool = False,
+    ) -> Any:
+        """Execute a statement via the owning resources object."""
+        ...
+
+    def get_cloud_provider(self) -> str:
+        """Return the configured cloud provider identifier (e.g. 'aws', 'azure')."""
+        ...
+
+
+# ---------------------------------------------------------------------------
+# Engine Service
+# ---------------------------------------------------------------------------
+class EngineServiceSQL:
+    """Snowflake engine management backed by SQL stored procedures."""
+
+    def __init__(self, resources: _ExecResources):
+        """Create an engine service bound to a resources-like executor."""
+        self._res = resources
+
+    @staticmethod
+    def _parse_settings(val: Any) -> Dict[str, Any] | None:
+        if val is None:
+            return None
+        if isinstance(val, dict):
+            return val
+        if isinstance(val, str):
+            s = val.strip()
+            if not s:
+                return None
+            try:
+                parsed = json.loads(s)
+                return parsed if isinstance(parsed, dict) else {"value": parsed}
+            except Exception:
+                return {"value": val}
+        # Snowflake VARIANT may arrive as a list/tuple/etc; preserve it under a wrapper
+        return {"value": val}
+
+    def list_engines(
+        self,
+        *,
+        state: str | None = None,
+        name: str | None = None,
+        type: str | None = None,
+        size: str | None = None,
+        created_by: str | None = None,
+    ) -> List[Dict[str, Any]]:
+        """
+        List engines with optional filtering.
+
+        Uses parameterized queries (? placeholders) for SQL injection protection.
+        """
+
+        where_conditions: list[str] = []
+        params: list[Any] = []
+
+        if state:
+            where_conditions.append("STATUS = ?")
+            params.append(state.upper())
+        if name:
+            where_conditions.append("UPPER(NAME) LIKE ?")
+            params.append(f"%{name.upper()}%")
+        if type:
+            where_conditions.append("TYPE = ?")
+            params.append(type.upper())
+        if size:
+            where_conditions.append("SIZE = ?")
+            params.append(size)
+        if created_by:
+            where_conditions.append("UPPER(CREATED_BY) LIKE ?")
+            params.append(f"%{created_by.upper()}%")
+
+        where_clause = f"WHERE {' AND '.join(where_conditions)}" if where_conditions else ""
+        where_suffix = f" {where_clause}" if where_clause else ""
+        statement = f"""
+            SELECT
+                NAME, TYPE, ID, SIZE, STATUS, CREATED_BY, CREATED_ON, UPDATED_ON,
+                AUTO_SUSPEND_MINS, SUSPENDS_AT, SETTINGS
+            FROM {API_SCHEMA}.engines{where_suffix}
+            ORDER BY NAME ASC;
+        """
+
+        results = self._res._exec(statement, params)
+        if not results:
+            return []
+
+        return [
+            {
+                "name": row["NAME"],
+                "type": row["TYPE"],
+                "id": row["ID"],
+                "size": row["SIZE"],
+                "state": row["STATUS"],  # callers expect 'state'
+                "created_by": row["CREATED_BY"],
+                "created_on": row["CREATED_ON"],
+                "updated_on": row["UPDATED_ON"],
+                "auto_suspend_mins": row["AUTO_SUSPEND_MINS"],
+                "suspends_at": row["SUSPENDS_AT"],
+                "settings": self._parse_settings(
+                    # Snowpark Row supports dict-style indexing but not `.get()`.
+                    row["SETTINGS"] if "SETTINGS" in row else None
+                ),
+            }
+            for row in results
+        ]
+
+    def get_engine(self, name: str, type: str) -> EngineState | None:
+        """Fetch a single engine by (name, type), returning None if not found."""
+        results = self._res._exec(
+            f"""
+            SELECT
+                NAME, TYPE, ID, SIZE, STATUS, CREATED_BY, CREATED_ON, UPDATED_ON,
+                VERSION, AUTO_SUSPEND_MINS, SUSPENDS_AT, SETTINGS
+            FROM {API_SCHEMA}.engines
+            WHERE NAME = ? AND TYPE = ?;
+            """,
+            [name, type],
+        )
+        if not results:
+            return None
+        engine = results[0]
+        if not engine:
+            return None
+        engine_state: EngineState = {
+            "name": engine["NAME"],
+            "type": engine["TYPE"],
+            "id": engine["ID"],
+            "size": engine["SIZE"],
+            "state": engine["STATUS"],  # callers expect 'state'
+            "created_by": engine["CREATED_BY"],
+            "created_on": engine["CREATED_ON"],
+            "updated_on": engine["UPDATED_ON"],
+            "version": engine["VERSION"],
+            "auto_suspend": engine["AUTO_SUSPEND_MINS"],
+            "suspends_at": engine["SUSPENDS_AT"],
+            "settings": self._parse_settings(engine["SETTINGS"] if "SETTINGS" in engine else None),
+        }
+        return engine_state
+
+    def _create_engine(
+        self,
+        *,
+        name: str,
+        type: str = EngineType.LOGIC,
+        size: str | None = None,
+        auto_suspend_mins: int | None = None,
+        is_async: bool = False,
+        headers: Dict | None = None,
+        settings: Dict[str, Any] | None = None,
+    ) -> None:
+        """Create an engine using the appropriate stored procedure.
+
+        Note: `headers` is accepted for API compatibility; it is not currently used
+        in the Snowflake implementation.
+        """
+        API = "create_engine_async" if is_async else "create_engine"
+        if size is None:
+            size = self._res.config.get_default_engine_size()
+        if auto_suspend_mins is None:
+            auto_suspend_mins = self._res.config.get_default_auto_suspend_mins()
+
+        try:
+            with debugging.span(API, name=name, size=size, auto_suspend_mins=auto_suspend_mins, engine_type=type):
+                payload: Dict[str, Any] | None = None
+                if settings:
+                    payload = {"settings": dict(settings)}
+                if auto_suspend_mins is not None:
+                    payload = dict(payload or {})
+                    payload["auto_suspend_mins"] = auto_suspend_mins
+
+                if payload is None:
+                    self._res._exec(
+                        f"call {API_SCHEMA}.{API}(?, ?, ?, null);",
+                        [type, name, size],
+                    )
+                else:
+                    self._res._exec(
+                        f"call {API_SCHEMA}.{API}(?, ?, ?, PARSE_JSON(?));",
+                        [type, name, size, json.dumps(payload)],
+                    )
+        except Exception as e:
+            raise EngineProvisioningFailed(name, e) from e
+
+    def create_engine(
+        self,
+        name: str,
+        type: str | None = None,
+        size: str | None = None,
+        auto_suspend_mins: int | None = None,
+        headers: Dict | None = None,
+        settings: Dict[str, Any] | None = None,
+    ) -> None:
+        """Create an engine (synchronous variant)."""
+        if type is None:
+            type = EngineType.LOGIC
+        self._create_engine(
+            name=name,
+            type=type,
+            size=size,
+            auto_suspend_mins=auto_suspend_mins,
+            headers=headers,
+            settings=settings,
+        )
+
+    def create_engine_async(
+        self,
+        name: str,
+        type: str = EngineType.LOGIC,
+        size: str | None = None,
+        auto_suspend_mins: int | None = None,
+    ) -> None:
+        """Create an engine asynchronously."""
+        self._create_engine(
+            name=name,
+            type=type,
+            size=size,
+            auto_suspend_mins=auto_suspend_mins,
+            is_async=True,
+        )
+
+    def delete_engine(self, name: str, type: str) -> None:
+        """Delete an engine by (name, type)."""
+        self._res._exec(
+            f"call {API_SCHEMA}.delete_engine(?, ?);",
+            [type, name],
+        )
+
+    def suspend_engine(self, name: str, type: str | None = None) -> None:
+        """Suspend an engine by name (and optional type)."""
+        if type is None:
+            type = EngineType.LOGIC
+        self._res._exec(
+            f"call {API_SCHEMA}.suspend_engine(?, ?);",
+            [type, name],
+        )
+
+    def resume_engine(self, name: str, type: str | None = None, headers: Dict | None = None) -> Dict[str, Any]:
+        """Resume an engine and block until it is READY.
+
+        This preserves historical behavior where `resume_engine` was synchronous.
+        Use `resume_engine_async` for a fire-and-forget resume call.
+        """
+        if type is None:
+            type = EngineType.LOGIC
+        self.resume_engine_async(name, type=type, headers=headers)
+        poll_with_specified_overhead(
+            lambda: (self.get_engine(name, type) or {}).get("state") == "READY",
+            overhead_rate=0.1,
+            max_delay=0.5,
+            timeout=900,
+        )
+        return {}
+
+    def resume_engine_async(self, name: str, type: str | None = None, headers: Dict | None = None) -> Dict[str, Any]:
+        """Resume an engine asynchronously and return an API-compatible payload.
+
+        Note: `headers` is accepted for API compatibility; it is not currently used
+        in the Snowflake implementation.
+        """
+        if type is None:
+            type = EngineType.LOGIC
+        self._res._exec(
+            f"call {API_SCHEMA}.resume_engine_async(?, ?);",
+            [type, name],
+        )
+        return {}
+
+    def validate_engine_size(self, size: str) -> Tuple[bool, List[str]]:
+        """Validate an engine size, returning (is_valid, allowed_sizes_if_invalid)."""
+        if size is not None:
+            sizes = self.get_engine_sizes()
+            if size not in sizes:
+                return False, sizes
+        return True, []
+
+    def get_engine_sizes(self, cloud_provider: str | None = None) -> List[str]:
+        """Return the list of valid engine sizes for the given cloud provider."""
+        if cloud_provider is None:
+            cloud_provider = self._res.get_cloud_provider()
+
+        if cloud_provider == "azure":
+            sizes = ENGINE_SIZES_AZURE
+        else:
+            sizes = ENGINE_SIZES_AWS
+
+        if self._res.config.show_all_engine_sizes():
+            return INTERNAL_ENGINE_SIZES + sizes
+        return sizes
+
+    def alter_engine_pool(self, size: str | None = None, mins: int | None = None, maxs: int | None = None) -> None:
+        """Alter engine pool node limits for Snowflake."""
+        # Keep the exact SQL shape used by Resources for backwards compatibility.
+        self._res._exec(f"call {APP_NAME}.api.alter_engine_pool_node_limits('{size}', {mins}, {maxs});")
+

relationalai/clients/resources/snowflake/engine_state_handlers.py

@@ -0,0 +1,315 @@
+"""
+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."""
+    name: str
+    size: str | None
+    type: str
+    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.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.name, context.type),
+                overhead_rate=0.1,
+                max_delay=0.5,
+                timeout=900
+            )
+
+        # Return updated engine (should be READY now)
+        updated_engine = resources.get_engine(context.name, context.type)
+        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.name}'",
+            f"Engine '{context.name}' resumed",
+            f"Failed to resume engine '{context.name}'"
+        ):
+            try:
+                resources.resume_engine_async(context.name, type=context.type, headers=context.headers)
+                poll_with_specified_overhead(
+                    lambda: resources.is_engine_ready(context.name, context.type),
+                    overhead_rate=0.1,
+                    max_delay=0.5,
+                    timeout=900
+                )
+            except Exception:
+                raise EngineResumeFailed(context.name)
+
+        # Return updated engine (should be READY now)
+        updated_engine = resources.get_engine(context.name, context.type)
+        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.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.name, context.type)
+            # 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.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.name}",
+            f"Auto-created engine {context.name}",
+            "Engine creation failed",
+        ):
+            resources.create_engine(
+                context.name,
+                size=context.size,
+                type=context.type,
+                headers=context.headers,
+            )
+
+        return resources.get_engine(context.name, context.type)
+
+
+# ============================================================================
+# 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.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.name}",
+            })
+
+        try:
+            resources.resume_engine_async(context.name, type=context.type)
+        except Exception:
+            raise EngineResumeFailed(context.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.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.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.name, context.type)
+            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.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.name, size=context.size, type=context.type)
+
+        if context.spinner:
+            context.spinner.update_messages({
+                "finished_message": f"Starting engine {context.name}...",
+            })
+        return None  # Engine is being created asynchronously
+