relationalai 0.13.5__py3-none-any.whl → 1.0.0a2__py3-none-any.whl

relationalai/clients/resources/snowflake/engine_service.py

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

@@ -1,315 +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."""
-    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
-