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

relationalai/clients/client.py

@@ -164,11 +164,18 @@ class ResourcesBase(ABC):
         pass
 
     @abstractmethod
-    def list_engines(self, state: str|None = None) -> List[Any]:
+    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[Any]:
         pass
 
     @abstractmethod
-    def get_engine(self, name: str) -> EngineState | None:
+    def get_engine(self, name: str, type: str) -> EngineState | None:
         pass
 
     @abstractmethod
@@ -180,23 +187,31 @@ class ResourcesBase(ABC):
         pass
 
     @abstractmethod
-    def create_engine(self, name: str, size: str|None, auto_suspend_mins: int|None) -> dict | None:
+    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 | None = None,
+    ) -> dict | None:
         pass
 
     @abstractmethod
-    def delete_engine(self, name:str, force:bool=False) -> dict | None:
+    def delete_engine(self, name:str, type: str) -> dict | None:
         pass
 
     @abstractmethod
-    def suspend_engine(self, name: str):
+    def suspend_engine(self, name: str, type: str | None = None):
         pass
 
     @abstractmethod
-    def resume_engine(self, name: str, headers: Dict | None = None) -> dict:
+    def resume_engine(self, name: str, type: str | None = None, headers: Dict | None = None) -> dict:
         pass
 
     @abstractmethod
-    def resume_engine_async(self, name: str) -> dict:
+    def resume_engine_async(self, name: str, type: str | None = None, headers: Dict | None = None) -> dict:
         pass
 
     @abstractmethod
@@ -210,7 +225,7 @@ class ResourcesBase(ABC):
         return engine
 
     @abstractmethod
-    def auto_create_engine_async(self, name: str | None = None) -> str:
+    def auto_create_engine_async(self, name: str | None = None, type: str | None = None) -> str:
         pass
 
     _active_engine: EngineState|None = None
@@ -428,14 +443,34 @@ class ProviderBase(ABC):
 
     resources: ResourcesBase
 
-    def list_engines(self, state: str | None = None):
-        return self.resources.list_engines(state)
+    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.resources.list_engines(state=state, name=name, type=type, size=size, created_by=created_by)
 
-    def create_engine(self, name:str, size:str|None=None, auto_suspend_mins:int|None=None):
-        return self.resources.create_engine(name, size, auto_suspend_mins)
+    def create_engine(
+        self,
+        name: str,
+        type: str | None = None,
+        size: str | None = None,
+        auto_suspend_mins: int | None = None,
+        settings: dict | None = None,
+    ):
+        return self.resources.create_engine(
+            name,
+            type=type,
+            size=size,
+            auto_suspend_mins=auto_suspend_mins,
+            settings=settings,
+        )
 
-    def delete_engine(self, name:str):
-        return self.resources.delete_engine(name)
+    def delete_engine(self, name:str, type: str = "LOGIC"):
+        return self.resources.delete_engine(name, type)
 
     def get_transaction(self, transaction_id:str):
         return self.resources.get_transaction(transaction_id)

relationalai/clients/local.py

@@ -172,10 +172,17 @@ class LocalResources(ResourcesBase):
     def get_engine_sizes(self, cloud_provider: str | None = None):
         raise NotImplementedError("get_engine_sizes not supported in local mode")
 
-    def list_engines(self, state: str | None = None):
+    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,
+    ):
         raise NotImplementedError("list_engines not supported in local mode")
 
-    def get_engine(self, name: str):
+    def get_engine(self, name: str, type: str):
         raise NotImplementedError("get_engine not supported in local mode")
 
     def get_cloud_provider(self) -> str:
@@ -184,25 +191,33 @@ class LocalResources(ResourcesBase):
     def is_valid_engine_state(self, name: str) -> bool:
         raise NotImplementedError("is_valid_engine_state not supported in local mode")
 
-    def create_engine(self, name: str, size:str|None=None, auto_suspend_mins:int|None=None):
+    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 | None = None,
+    ):
         raise NotImplementedError("create_engine not supported in local mode")
 
-    def delete_engine(self, name: str, force: bool = False):
+    def delete_engine(self, name: str, type: str):
         raise NotImplementedError("delete_engine not supported in local mode")
 
-    def suspend_engine(self, name: str):
+    def suspend_engine(self, name: str, type: str | None = None):
         raise NotImplementedError("suspend_engine not supported in local mode")
 
-    def resume_engine(self, name: str, headers: Dict | None = None):
+    def resume_engine(self, name: str, type: str | None = None, headers: Dict | None = None):
         raise NotImplementedError("resume_engine not supported in local mode")
 
-    def resume_engine_async(self, name: str):
+    def resume_engine_async(self, name: str, type: str | None = None, headers: Dict | None = None):
         raise NotImplementedError("resume_engine_async not supported in local mode")
 
     def alter_engine_pool(self, size: str | None = None, mins: int | None = None, maxs: int | None = None):
         raise NotImplementedError("alter_engine_pool not supported in local mode")
 
-    def auto_create_engine_async(self, name: str | None = None) -> str:
+    def auto_create_engine_async(self, name: str | None = None, type: str | None = None) -> str:
         raise NotImplementedError("auto_create_engine_async not supported in local mode")
 
     #--------------------------------------------------

relationalai/clients/resources/azure/azure.py

@@ -78,34 +78,59 @@ class Resources(ResourcesBase):
     def get_cloud_provider(self) -> str:
         return "azure"
 
-    def list_engines(self, state:str|None = None):
-        return api.list_engines(self._api_ctx(), state)
-
-    def get_engine(self, name:str) -> EngineState:
-        return cast(EngineState, api.get_engine(self._api_ctx(), 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,
+    ):
+        # Azure only supports state filtering at the API level; other filters are ignored.
+        engines = api.list_engines(self._api_ctx(), state)
+        # Ensure EngineState shape includes 'type' for callers/tests.
+        for eng in engines:
+            eng.setdefault("type", "LOGIC")
+        return engines
+
+    def get_engine(self, name: str, type: str) -> EngineState:
+        # type is ignored for Azure as it doesn't support multiple engine types
+        engine = cast(EngineState, api.get_engine(self._api_ctx(), name))
+        engine.setdefault("type", "LOGIC")
+        return engine
 
     def is_valid_engine_state(self, name:str):
         return name in VALID_ENGINE_STATES
 
-    def create_engine(self, name:str, size:str|None=None, auto_suspend_mins: int|None=None):
+    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 | None = None,
+    ):
+        # Azure only supports one engine type, so type parameter is ignored
         if size is None:
             size = "M"
         with debugging.span("create_engine", name=name, size=size):
             return api.create_engine_wait(self._api_ctx(), name, size)
 
-    def delete_engine(self, name:str, force:bool=False):
+    def delete_engine(self, name:str, type: str):
         return api.delete_engine(self._api_ctx(), name)
 
-    def suspend_engine(self, name:str):
+    def suspend_engine(self, name:str, type: str | None = None): # type is ignored for Azure
         return api.suspend_engine(self._api_ctx(), name)
 
-    def resume_engine(self, name:str, headers={}):
+    def resume_engine(self, name:str, type: str | None = None, headers={}):
+        # type is ignored for Azure as it doesn't support multiple engine types
         return api.resume_engine_wait(self._api_ctx(), name)
 
-    def resume_engine_async(self, name:str):
+    def resume_engine_async(self, name:str, type: str | None = None, headers: Dict | None = None):
         return api.resume_engine(self._api_ctx(), name)
 
-    def auto_create_engine_async(self, name: str | None = None) -> str:
+    def auto_create_engine_async(self, name: str | None = None, type: str | None = None) -> str:
         raise Exception("Azure doesn't support auto_create_engine_async")
 
     def alter_engine_pool(self, size: str | None = None, mins: int | None = None, maxs: int | None = None):

relationalai/clients/resources/snowflake/__init__.py

@@ -2,8 +2,8 @@
 Snowflake resources module.
 """
 # Import order matters - Resources must be imported first since other classes depend on it
-from .snowflake import Resources, Provider, Graph, SnowflakeClient, APP_NAME, PYREL_ROOT_DB, ExecContext, INTERNAL_ENGINE_SIZES, ENGINE_SIZES_AWS, ENGINE_SIZES_AZURE, PrimaryKey, PRINT_TXN_PROGRESS_FLAG
-
+from .snowflake import Resources, Provider, Graph, SnowflakeClient, APP_NAME, PYREL_ROOT_DB, ExecContext, PrimaryKey, PRINT_TXN_PROGRESS_FLAG
+from .engine_service import EngineType, INTERNAL_ENGINE_SIZES, ENGINE_SIZES_AWS, ENGINE_SIZES_AZURE
 # These imports depend on Resources, so they come after
 from .cli_resources import CLIResources
 from .use_index_resources import UseIndexResources
@@ -12,7 +12,7 @@ from .resources_factory import create_resources_instance
 
 __all__ = [
     'Resources', 'DirectAccessResources', 'Provider', 'Graph', 'SnowflakeClient',
-    'APP_NAME', 'PYREL_ROOT_DB', 'CLIResources', 'UseIndexResources', 'ExecContext',
+    'APP_NAME', 'PYREL_ROOT_DB', 'CLIResources', 'UseIndexResources', 'ExecContext', 'EngineType',
     'INTERNAL_ENGINE_SIZES', 'ENGINE_SIZES_AWS', 'ENGINE_SIZES_AZURE', 'PrimaryKey',
     'PRINT_TXN_PROGRESS_FLAG', 'create_resources_instance',
 ]

relationalai/clients/resources/snowflake/cli_resources.py

@@ -11,6 +11,7 @@ from ....tools.constants import RAI_APP_NAME
 # Import Resources from snowflake - this creates a dependency but no circular import
 # since snowflake.py doesn't import from this file
 from .snowflake import Resources, ExecContext
+from .error_handlers import AppMissingErrorHandler, AppFunctionMissingErrorHandler, ServiceNotStartedErrorHandler
 
 
 class CLIResources(Resources):
@@ -20,7 +21,17 @@ class CLIResources(Resources):
     """
 
     def _handle_standard_exec_errors(self, e: Exception, ctx: ExecContext) -> Any | None:
-        """For CLI resources, re-raise exceptions without transformation."""
+        """
+        For CLI resources, keep exceptions raw except for few specific cases.
+        """
+        message = str(e).lower()
+        for handler in (
+            ServiceNotStartedErrorHandler(),
+            AppMissingErrorHandler(),
+            AppFunctionMissingErrorHandler(),
+        ):
+            if handler.matches(e, message, ctx, self):
+                handler.handle(e, ctx, self)
         raise e
 
     def list_warehouses(self):

relationalai/clients/resources/snowflake/direct_access_resources.py

@@ -513,7 +513,14 @@ class DirectAccessResources(UseIndexResources):
     # Engines
     #--------------------------------------------------
 
-    def list_engines(self, state: str | None = None):
+    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,
+    ):
         response = self.request("list_engines")
         if response.status_code != 200:
             raise ResponseStatusException(
@@ -526,19 +533,31 @@ class DirectAccessResources(UseIndexResources):
             {
                 "name": engine["name"],
                 "id": engine["id"],
+                "type": engine.get("type", "LOGIC"),
                 "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"],
+                # Optional fields (present in newer APIs / service-functions path)
+                "auto_suspend_mins": engine.get("auto_suspend_mins"),
+                "suspends_at": engine.get("suspends_at"),
+                "settings": engine.get("settings"),
             }
             for engine in response_content.get("engines", [])
-            if state is None or engine.get("status") == state
+            if (state is None or engine.get("status") == state)
+            and (name is None or name.upper() in engine.get("name", "").upper())
+            and (type is None or engine.get("type", "LOGIC").upper() == type.upper())
+            and (size is None or engine.get("size") == size)
+            and (created_by is None or created_by.upper() in engine.get("created_by", "").upper())
         ]
         return sorted(engines, key=lambda x: x["name"])
 
-    def get_engine(self, name: str):
-        response = self.request("get_engine", path_params={"engine_name": name, "engine_type": "logic"}, skip_auto_create=True)
+    def get_engine(self, name: str, type: str):
+        if type is None:
+            raise Exception("Engine type is required. Valid types are: LOGIC, SOLVER, ML")
+        engine_type_lower = type.lower()
+        response = self.request("get_engine", path_params={"engine_name": name, "engine_type": engine_type_lower}, skip_auto_create=True)
         if response.status_code == 404: # engine not found return 404
             return None
         elif response.status_code != 200:
@@ -552,6 +571,7 @@ class DirectAccessResources(UseIndexResources):
             "name": engine["name"],
             "id": engine["id"],
             "size": engine["size"],
+            "type": engine.get("type", type),
             "state": engine["status"], # callers are expecting 'state'
             "created_by": engine["created_by"],
             "created_on": engine["created_on"],
@@ -559,31 +579,47 @@ class DirectAccessResources(UseIndexResources):
             "version": engine["version"],
             "auto_suspend": engine["auto_suspend_mins"],
             "suspends_at": engine["suspends_at"],
+            "settings": engine.get("settings"),
         }
         return engine_state
 
     def _create_engine(
             self,
             name: str,
+            type: str = "LOGIC",
             size: str | None = None,
             auto_suspend_mins: int | None = None,
             is_async: bool = False,
-            headers: Dict[str, str] | None = None
+            headers: Dict[str, str] | None = None,
+            settings: Dict[str, Any] | None = None,
         ):
         # only async engine creation supported via direct access
         if not is_async:
-            return super()._create_engine(name, size, auto_suspend_mins, is_async, headers=headers)
+            return super()._create_engine(
+                name,
+                type=type,
+                size=size,
+                auto_suspend_mins=auto_suspend_mins,
+                is_async=is_async,
+                headers=headers,
+                settings=settings,
+            )
+        engine_type_lower = type.lower()
         payload:Dict[str, Any] = {
             "name": name,
         }
+        # Allow passing arbitrary engine settings (API-dependent).
+        if settings:
+            payload["settings"] = settings
         if auto_suspend_mins is not None:
             payload["auto_suspend_mins"] = auto_suspend_mins
-        if size is not None:
-            payload["size"] = size
+        if size is None:
+            size = self.config.get_default_engine_size()
+        payload["size"] = size
         response = self.request(
             "create_engine",
             payload=payload,
-            path_params={"engine_type": "logic"},
+            path_params={"engine_type": engine_type_lower},
             headers=headers,
             skip_auto_create=True,
         )
@@ -592,11 +628,11 @@ class DirectAccessResources(UseIndexResources):
                 f"Failed to create engine {name} with size {size}.", response
             )
 
-    def delete_engine(self, name:str, force:bool = False, headers={}):
+    def delete_engine(self, name: str, type: str):
         response = self.request(
             "delete_engine",
-            path_params={"engine_name": name, "engine_type": "logic"},
-            headers=headers,
+            path_params={"engine_name": name, "engine_type": type.lower()},
+            headers={},
             skip_auto_create=True,
         )
         if response.status_code != 200:
@@ -604,10 +640,12 @@ class DirectAccessResources(UseIndexResources):
                 f"Failed to delete engine {name}.", response
             )
 
-    def suspend_engine(self, name:str):
+    def suspend_engine(self, name: str, type: str | None = None):
+        if type is None:
+            type = "LOGIC"
         response = self.request(
             "suspend_engine",
-            path_params={"engine_name": name, "engine_type": "logic"},
+            path_params={"engine_name": name, "engine_type": type.lower()},
             skip_auto_create=True,
         )
         if response.status_code != 200:
@@ -615,11 +653,13 @@ class DirectAccessResources(UseIndexResources):
                 f"Failed to suspend engine {name}.", response
             )
 
-    def resume_engine_async(self, name:str, headers={}):
+    def resume_engine_async(self, name: str, type: str | None = None, headers: Dict | None = None):
+        if type is None:
+            type = "LOGIC"
         response = self.request(
             "resume_engine",
-            path_params={"engine_name": name, "engine_type": "logic"},
-            headers=headers,
+            path_params={"engine_name": name, "engine_type": type.lower()},
+            headers=headers or {},
             skip_auto_create=True,
         )
         if response.status_code != 200: