relationalai 0.13.1__py3-none-any.whl → 0.13.3__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)
@@ -579,7 +614,6 @@ class Client():
                     self._timed_query(
                         "update_registry",
                         dependencies.generate_update_registry(),
-                        readonly=False,
                         abort_on_error=False,
                     )
 
@@ -588,7 +622,6 @@ class Client():
                     self._timed_query(
                         "update_packages",
                         dependencies.generate_update_packages(),
-                        readonly=False,
                         abort_on_error=False,
                     )
             else:
@@ -611,10 +644,11 @@ class Client():
         finally:
             self._database = database_name
 
-    def _timed_query(self, span_name:str, code: str, readonly=True, abort_on_error=True):
+    def _timed_query(self, span_name:str, code: str, abort_on_error=True):
         with debugging.span(span_name, model=self._database) as end_span:
             start = time.perf_counter()
-            res, raw = self._query(code, None, end_span, readonly=readonly, abort_on_error=abort_on_error)
+            # NOTE hardcoding to readonly=False, read-only Rel transactions are deprecated.
+            res, raw = self._query(code, None, end_span, readonly=False, abort_on_error=abort_on_error)
             debugging.time(span_name, time.perf_counter() - start, code=code)
             return res, raw
 

relationalai/clients/exec_txn_poller.py

@@ -27,16 +27,43 @@ class ExecTxnPoller:
 
     def __init__(
         self,
+        print_txn_progress: bool,
         resource: "Resources",
-        txn_id: str,
+        txn_id: Optional[str] = None,
         headers: Optional[Dict] = None,
         txn_start_time: Optional[float] = None,
     ):
+        self.print_txn_progress = print_txn_progress
         self.res = resource
         self.txn_id = txn_id
         self.headers = headers or {}
         self.txn_start_time = txn_start_time or time.time()
 
+    def __enter__(self) -> ExecTxnPoller:
+        if not self.print_txn_progress:
+            return self
+        self.progress = create_progress(
+            description=lambda: self.description_with_timing(),
+            success_message="",  # We'll handle this ourselves
+            leading_newline=False,
+            trailing_newline=False,
+            show_duration_summary=False,
+        )
+        self.progress.__enter__()
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback) -> None:
+        if not self.print_txn_progress or self.txn_id is None:
+            return
+        # Update to success message with duration
+        total_duration = time.time() - self.txn_start_time
+        txn_id = self.txn_id
+        self.progress.update_main_status(
+            query_complete_message(txn_id, total_duration)
+        )
+        self.progress.__exit__(exc_type, exc_value, traceback)
+        return
+
     def poll(self) -> bool:
         """
         Poll for transaction completion with interactive progress display.
@@ -44,48 +71,52 @@ class ExecTxnPoller:
         Returns:
             True if transaction completed successfully, False otherwise
         """
+        if not self.txn_id:
+            raise ValueError("Transaction ID must be provided for polling.")
+        else:
+            txn_id = self.txn_id
+
+        if self.print_txn_progress:
+            # Update the main status to include the new txn_id
+            self.progress.update_main_status_fn(
+                lambda: self.description_with_timing(txn_id),
+            )
 
         # Don't show duration summary - we handle our own completion message
-        with create_progress(
-            description="Evaluating Query...",
-            success_message="",  # We'll handle this ourselves
-            leading_newline=False,
-            trailing_newline=False,
-            show_duration_summary=False,
-        ) as progress:
-            def check_status() -> bool:
-                """Check if transaction is complete."""
-                elapsed = time.time() - self.txn_start_time
-                # Update the main status with elapsed time
-                progress.update_main_status(
-                    query_progress_message(self.txn_id, elapsed)
-                )
-                return self.res._check_exec_async_status(self.txn_id, headers=self.headers)
-
-            with debugging.span("wait", txn_id=self.txn_id):
-                poll_with_specified_overhead(check_status, overhead_rate=POLL_OVERHEAD_RATE)
-
-            # Calculate final duration
-            total_duration = time.time() - self.txn_start_time
-
-            # Update to success message with duration
-            progress.update_main_status(
-                query_complete_message(self.txn_id, total_duration)
-            )
+        def check_status() -> bool:
+            """Check if transaction is complete."""
+            finished = self.res._check_exec_async_status(txn_id, headers=self.headers)
+            return finished
+
+        with debugging.span("wait", txn_id=self.txn_id):
+            poll_with_specified_overhead(check_status, overhead_rate=POLL_OVERHEAD_RATE)
+
 
         return True
 
+    def description_with_timing(self, txn_id: str | None = None) -> str:
+        elapsed = time.time() - self.txn_start_time
+        if txn_id is None:
+            return query_progress_header(elapsed)
+        else:
+            return query_progress_message(txn_id, elapsed)
+
+def query_progress_header(duration: float) -> str:
+    # Don't print sub-second decimals, because it updates too fast and is distracting.
+    duration_str = format_duration(duration, seconds_decimals=False)
+    return f"Evaluating Query... {duration_str:>15}\n"
+
 def query_progress_message(id: str, duration: float) -> str:
     return (
+        query_progress_header(duration) +
         # Print with whitespace to align with the end of the transaction ID
-        f"Evaluating Query... {format_duration(duration):>18}\n" +
-        f"{GRAY_COLOR}Query: {id}{ENDC}"
+        f"{GRAY_COLOR}ID: {id}{ENDC}"
     )
 
 def query_complete_message(id: str, duration: float, status_header: bool = False) -> str:
     return (
         (f"{GREEN_COLOR}✅ " if status_header else "") +
         # Print with whitespace to align with the end of the transaction ID
-        f"Query Complete: {format_duration(duration):>24}\n" +
-        f"{GRAY_COLOR}Query: {id}{ENDC}"
+        f"Query Complete: {format_duration(duration):>21}\n" +
+        f"{GRAY_COLOR}ID: {id}{ENDC}"
     )

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

@@ -13,7 +13,7 @@ from ...config import Config, ConfigStore, ENDPOINT_FILE
 from ...direct_access_client import DirectAccessClient
 from ...types import EngineState
 from ...util import get_pyrel_version, safe_json_loads, ms_to_timestamp
-from ....errors import GuardRailsException, ResponseStatusException, QueryTimeoutExceededException
+from ....errors import GuardRailsException, ResponseStatusException, QueryTimeoutExceededException, RAIException
 from snowflake.snowpark import Session
 
 # Import UseIndexResources to enable use_index functionality with direct access
@@ -163,11 +163,12 @@ class DirectAccessResources(UseIndexResources):
                             headers=headers,
                         )
                         response = _send_request()
-            except requests.exceptions.ConnectionError as e:
+            except (requests.exceptions.ConnectionError, RAIException) as e:
                 messages = collect_error_messages(e)
-                if any("nameresolutionerror" in msg for msg in messages):
-                    # when we can not resolve the service endpoint, we assume it is outdated
-                    # hence, we try to retrieve it again and query again.
+                if any("nameresolutionerror" in msg for msg in messages) or \
+                   any("could not find the service associated with endpoint" in msg for msg in messages):
+                    # when we can not resolve the service endpoint or the service is not found,
+                    # we assume the endpoint is outdated, so we retrieve it again and retry.
                     self.direct_access_client.service_endpoint = self._retrieve_service_endpoint(
                         enforce_update=True,
                     )
@@ -513,7 +514,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 +534,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 +572,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 +580,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 +629,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 +641,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 +654,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: