relationalai 0.13.1__py3-none-any.whl → 0.13.2__py3-none-any.whl

relationalai/clients/resources/snowflake/snowflake.py

@@ -39,6 +39,7 @@ from ...types import AvailableModel, EngineState, Import, ImportSource, ImportSo
 from ...config import Config
 from ...client import Client, ExportParams, ProviderBase, ResourcesBase
 from ...util import IdentityParser, escape_for_f_string, get_pyrel_version, get_with_retries, poll_with_specified_overhead, safe_json_loads, sanitize_module_name, scrub_exception, wrap_with_request_id, normalize_datetime
+from .engine_service import EngineServiceSQL, EngineType
 from .util import (
     collect_error_messages,
     process_jinja_template,
@@ -64,6 +65,7 @@ from .error_handlers import (
     ErrorHandler,
     DuoSecurityErrorHandler,
     AppMissingErrorHandler,
+    AppFunctionMissingErrorHandler,
     DatabaseErrorsHandler,
     EngineErrorsHandler,
     ServiceNotStartedErrorHandler,
@@ -90,16 +92,10 @@ from .engine_state_handlers import (
 # Constants
 #--------------------------------------------------
 
-VALID_POOL_STATUS = ["ACTIVE", "IDLE", "SUSPENDED"]
 # transaction list and get return different fields (duration vs timings)
 LIST_TXN_SQL_FIELDS = ["id", "database_name", "engine_name", "state", "abort_reason", "read_only","created_by", "created_on", "finished_at", "duration"]
 GET_TXN_SQL_FIELDS = ["id", "database", "engine", "state", "abort_reason", "read_only","created_by", "created_on", "finished_at", "timings"]
 VALID_ENGINE_STATES = ["READY", "PENDING"]
-
-# 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"]
 # Note: ENGINE_ERRORS, ENGINE_NOT_READY_MSGS, DATABASE_ERRORS moved to util.py
 PYREL_ROOT_DB = 'pyrel_root_db'
 
@@ -184,11 +180,17 @@ class Resources(ResourcesBase):
         self._sproc_models = None
         # Store language for backward compatibility (used by child classes for use_index polling)
         self.language = language
+        # Engine subsystem (composition: keeps engine CRUD isolated from the core Resources class)
+        self._engines = EngineServiceSQL(self)
         # Register error and state handlers
         self._register_handlers()
         # Register atexit callback to cancel pending transactions
         atexit.register(self.cancel_pending_transactions)
 
+    @property
+    def engines(self) -> EngineServiceSQL:
+        return self._engines
+
     #--------------------------------------------------
     # Initialization & Properties
     #--------------------------------------------------
@@ -217,11 +219,12 @@ class Resources(ResourcesBase):
                 return handlers
         """
         return [
-            DuoSecurityErrorHandler(),
             AppMissingErrorHandler(),
+            AppFunctionMissingErrorHandler(),
+            ServiceNotStartedErrorHandler(),
+            DuoSecurityErrorHandler(),
             DatabaseErrorsHandler(),
             EngineErrorsHandler(),
-            ServiceNotStartedErrorHandler(),
             TransactionAbortedErrorHandler(),
         ]
 
@@ -613,7 +616,7 @@ class Resources(ResourcesBase):
 
         # Validate engine size
         if engine_size:
-            is_size_valid, sizes = self.validate_engine_size(engine_size)
+            is_size_valid, sizes = self._engines.validate_engine_size(engine_size)
             if not is_size_valid:
                 error_msg = f"Invalid engine size '{engine_size}'. Valid sizes are: {', '.join(sizes)}"
                 if use_default_size:
@@ -703,67 +706,27 @@ class Resources(ResourcesBase):
         if not handled:
             raise EngineProvisioningFailed(engine_name, error) from error
 
-    def validate_engine_size(self, size: str) -> Tuple[bool, List[str]]:
-        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):
-        sizes = []
-        if cloud_provider is None:
-            cloud_provider = self.get_cloud_provider()
-        if cloud_provider == 'azure':
-            sizes = ENGINE_SIZES_AZURE
-        else:
-            sizes = ENGINE_SIZES_AWS
-        if self.config.show_all_engine_sizes():
-            return INTERNAL_ENGINE_SIZES + sizes
-        else:
-            return sizes
+        return self._engines.get_engine_sizes(cloud_provider=cloud_provider)
 
-    def list_engines(self, state: str | None = None):
-        where_clause = f"WHERE STATUS = '{state.upper()}'" if state else ""
-        statement = f"SELECT NAME, ID, SIZE, STATUS, CREATED_BY, CREATED_ON, UPDATED_ON FROM {APP_NAME}.api.engines {where_clause} ORDER BY NAME ASC;"
-        results = self._exec(statement)
-        if not results:
-            return []
-        return [
-            {
-                "name": row["NAME"],
-                "id": row["ID"],
-                "size": row["SIZE"],
-                "state": row["STATUS"], # callers are expecting 'state'
-                "created_by": row["CREATED_BY"],
-                "created_on": row["CREATED_ON"],
-                "updated_on": row["UPDATED_ON"],
-            }
-            for row in results
-        ]
-
-    def get_engine(self, name: str):
-        results = self._exec(
-            f"SELECT NAME, ID, SIZE, STATUS, CREATED_BY, CREATED_ON, UPDATED_ON, VERSION, AUTO_SUSPEND_MINS, SUSPENDS_AT FROM {APP_NAME}.api.engines WHERE NAME='{name}';"
+    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,
+    ):
+        return self._engines.list_engines(
+            state=state,
+            name=name,
+            type=type,
+            size=size,
+            created_by=created_by,
         )
-        if not results:
-            return None
-        engine = results[0]
-        if not engine:
-            return None
-        engine_state: EngineState = {
-            "name": engine["NAME"],
-            "id": engine["ID"],
-            "size": engine["SIZE"],
-            "state": engine["STATUS"], # callers are expecting '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"]
-        }
-        return engine_state
+
+    def get_engine(self, name: str, type: str):
+        return self._engines.get_engine(name, type)
 
     def get_default_engine_name(self) -> str:
         if self.config.get("engine_name", None) is not None:
@@ -784,60 +747,82 @@ Otherwise, remove it from your '{profile}' configuration profile.
     def is_valid_engine_state(self, name:str):
         return name in VALID_ENGINE_STATES
 
+    # Can be overridden by subclasses (e.g. DirectAccessResources)
     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,
         ):
-        api = "create_engine_async" if is_async else "create_engine"
-        if size is None:
-            size = self.config.get_default_engine_size()
-        # if auto_suspend_mins is None, get the default value from the config
-        if auto_suspend_mins is None:
-            auto_suspend_mins = self.config.get_default_auto_suspend_mins()
-        try:
-            headers = debugging.gen_current_propagation_headers()
-            with debugging.span(api, name=name, size=size, auto_suspend_mins=auto_suspend_mins):
-                # check in case the config default is missing
-                if auto_suspend_mins is None:
-                    self._exec(f"call {APP_NAME}.api.{api}('{name}', '{size}', null, {headers});")
-                else:
-                    self._exec(f"call {APP_NAME}.api.{api}('{name}', '{size}', PARSE_JSON('{{\"auto_suspend_mins\": {auto_suspend_mins}}}'), {headers});")
-        except Exception as e:
-            raise EngineProvisioningFailed(name, e) from e
+        return self._engines._create_engine(
+            name=name,
+            type=type,
+            size=size,
+            auto_suspend_mins=auto_suspend_mins,
+            is_async=is_async,
+            headers=headers,
+            settings=settings,
+        )
 
-    def create_engine(self, name:str, size:str|None=None, auto_suspend_mins:int|None=None, headers: Dict | None = None):
-        self._create_engine(name, size, auto_suspend_mins, headers=headers)
+    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,
+    ):
+        if type is None:
+            type = EngineType.LOGIC
+        # Route through _create_engine so subclasses (e.g. DirectAccessResources)
+        # can override engine creation behavior.
+        return self._create_engine(
+            name=name,
+            type=type,
+            size=size,
+            auto_suspend_mins=auto_suspend_mins,
+            is_async=False,
+            headers=headers,
+            settings=settings,
+        )
 
-    def create_engine_async(self, name:str, size:str|None=None, auto_suspend_mins:int|None=None):
-        self._create_engine(name, size, auto_suspend_mins, True)
+    def create_engine_async(
+        self,
+        name: str,
+        type: str = EngineType.LOGIC,
+        size: str | None = None,
+        auto_suspend_mins: int | None = None,
+    ):
+        # Route through _create_engine so subclasses (e.g. DirectAccessResources)
+        # can override async engine creation behavior.
+        return self._create_engine(
+            name=name,
+            type=type,
+            size=size,
+            auto_suspend_mins=auto_suspend_mins,
+            is_async=True,
+        )
 
-    def delete_engine(self, name:str, force:bool = False, headers: Dict | None = None):
-        request_headers = debugging.add_current_propagation_headers(headers)
-        self._exec(f"call {APP_NAME}.api.delete_engine('{name}', {force},{request_headers});")
+    def delete_engine(self, name: str, type: str):
+        return self._engines.delete_engine(name, type)
 
-    def suspend_engine(self, name:str):
-        self._exec(f"call {APP_NAME}.api.suspend_engine('{name}');")
+    def suspend_engine(self, name: str, type: str | None = None):
+        return self._engines.suspend_engine(name, type)
 
-    def resume_engine(self, name:str, headers: Dict | None = None) -> Dict:
-        request_headers = debugging.add_current_propagation_headers(headers)
-        self._exec(f"call {APP_NAME}.api.resume_engine('{name}',{request_headers});")
-        # returning empty dict to match the expected return type
-        return {}
+    def resume_engine(self, name: str, type: str | None = None, headers: Dict | None = None) -> Dict:
+        return self._engines.resume_engine(name, type=type, headers=headers)
 
-    def resume_engine_async(self, name:str, headers: Dict | None = None) -> Dict:
-        if headers is None:
-            headers = {}
-        self._exec(f"call {APP_NAME}.api.resume_engine_async('{name}',{headers});")
-        # returning empty dict to match the expected return type
-        return {}
+    def resume_engine_async(self, name: str, type: str | None = None, headers: Dict | None = None) -> Dict:
+        return self._engines.resume_engine_async(name, type=type, headers=headers)
 
     def alter_engine_pool(self, size:str|None=None, mins:int|None=None, maxs:int|None=None):
         """Alter engine pool node limits for Snowflake."""
-        self._exec(f"call {APP_NAME}.api.alter_engine_pool_node_limits('{size}', {mins}, {maxs});")
+        return self._engines.alter_engine_pool(size=size, mins=mins, maxs=maxs)
 
     #--------------------------------------------------
     # Graphs
@@ -1871,11 +1856,17 @@ Otherwise, remove it from your '{profile}' configuration profile.
         assert isinstance(user, str), f"current_user() must return a string, not {type(user)}"
         return _sanitize_user_name(user)
 
-    def is_engine_ready(self, engine_name: str):
-        engine = self.get_engine(engine_name)
+    def is_engine_ready(self, engine_name: str, type: str = EngineType.LOGIC):
+        engine = self.get_engine(engine_name, type)
         return engine and engine["state"] == "READY"
 
-    def auto_create_engine(self, name: str | None = None, size: str | None = None, headers: Dict | None = None):
+    def auto_create_engine(
+        self,
+        name: str | None = None,
+        type: str = EngineType.LOGIC,
+        size: str | None = None,
+        headers: Dict | None = None,
+    ):
         """Synchronously create/ensure an engine is ready, blocking until ready."""
         with debugging.span("auto_create_engine", active=self._active_engine) as span:
             active = self._get_active_engine()
@@ -1883,18 +1874,19 @@ Otherwise, remove it from your '{profile}' configuration profile.
                 return active
 
             # Resolve and validate parameters
-            engine_name, engine_size = self._prepare_engine_params(name, size)
+            name, size = self._prepare_engine_params(name, size)
 
             try:
                 # Get current engine state
-                engine = self.get_engine(engine_name)
+                engine = self.get_engine(name, type)
                 if engine:
                     span.update(cast(dict, engine))
 
                 # Create context for state handling
                 context = EngineContext(
-                    engine_name=engine_name,
-                    engine_size=engine_size,
+                    name=name,
+                    size=size,
+                    type=type,
                     headers=headers,
                     requested_size=size,
                     span=span,
@@ -1904,12 +1896,14 @@ Otherwise, remove it from your '{profile}' configuration profile.
                 self._process_engine_state(engine, context, self._sync_engine_state_handlers)
 
             except Exception as e:
-                self._handle_engine_creation_errors(e, engine_name)
+                self._handle_engine_creation_errors(e, name)
 
-            return engine_name
+            return name
 
-    def auto_create_engine_async(self, name: str | None = None):
+    def auto_create_engine_async(self, name: str | None = None, type: str | None = None):
         """Asynchronously create/ensure an engine, returns immediately."""
+        if type is None:
+            type = EngineType.LOGIC
         active = self._get_active_engine()
         if active and (active == name or name is None):
             return active
@@ -1920,16 +1914,17 @@ Otherwise, remove it from your '{profile}' configuration profile.
         ) as spinner:
             with debugging.span("auto_create_engine_async", active=self._active_engine):
                 # Resolve and validate parameters (use_default_size=True for async)
-                engine_name, engine_size = self._prepare_engine_params(name, None, use_default_size=True)
+                name, size = self._prepare_engine_params(name, None, use_default_size=True)
 
                 try:
                     # Get current engine state
-                    engine = self.get_engine(engine_name)
+                    engine = self.get_engine(name, type)
 
                     # Create context for state handling
                     context = EngineContext(
-                        engine_name=engine_name,
-                        engine_size=engine_size,
+                        name=name,
+                        size=size,
+                        type=type,
                         headers=None,
                         requested_size=None,
                         spinner=spinner,
@@ -1940,11 +1935,11 @@ Otherwise, remove it from your '{profile}' configuration profile.
 
                 except Exception as e:
                     spinner.update_messages({
-                        "finished_message": f"Failed to create engine {engine_name}",
+                        "finished_message": f"Failed to create engine {name}",
                     })
-                    self._handle_engine_creation_errors(e, engine_name, preserve_rai_exception=True)
+                    self._handle_engine_creation_errors(e, name, preserve_rai_exception=True)
 
-                return engine_name
+                return name
 
     #--------------------------------------------------
     # Exec

relationalai/clients/types.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 from abc import ABC
 from datetime import datetime
 from typing import Any, Optional, TypedDict
+from typing_extensions import NotRequired
 from pathlib import Path
 from urllib.parse import urlparse
 
@@ -27,6 +28,7 @@ class AvailableModel(TypedDict):
 
 class EngineState(TypedDict):
     name: str
+    type: str
     id: str
     size: str
     state: str
@@ -38,6 +40,9 @@ class EngineState(TypedDict):
     auto_suspend: int|None
     suspends_at: datetime|None
 
+    # Optional JSON settings (engine configuration)
+    settings: NotRequired[dict | None]
+
 class SourceInfo(TypedDict, total=False):
     type: str|None
     state: str

relationalai/errors.py

@@ -1780,7 +1780,7 @@ class SnowflakeRaiAppNotStarted(RAIException):
     def __init__(self, app_name: str):
         self.app_name = app_name
         self.message = f"The RelationalAI app '{app_name}' isn't started."
-        self.name = "The RelationalAI application not started"
+        self.name = "The RelationalAI application isn't running"
         self.content = self.format_message()
 
         super().__init__(self.message, self.name, self.content)

relationalai/semantics/metamodel/typer/typer.py

@@ -1361,6 +1361,9 @@ class Replacer(visitor.Rewriter):
 # Typer pass
 #--------------------------------------------------
 
+# global flag to suppress type errors from being printed
+SUPPRESS_TYPE_ERRORS = False
+
 class InferTypes(compiler.Pass):
     def __init__(self):
         super().__init__()
@@ -1389,7 +1392,8 @@ class InferTypes(compiler.Pass):
         with debugging.span("type.replace"):
             final = Replacer(w.net).walk(model)
 
-        for err in w.net.errors:
-            rich.print(str(err), file=sys.stderr)
+        if not SUPPRESS_TYPE_ERRORS:
+            for err in w.net.errors:
+                rich.print(str(err), file=sys.stderr)
 
         return final