relationalai 1.0.0a3__py3-none-any.whl → 1.0.0a5__py3-none-any.whl

v0/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")
 
     #--------------------------------------------------

v0/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):

v0/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
-
+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,9 +12,9 @@ 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',
-    'create_resources_instance',
+    'PRINT_TXN_PROGRESS_FLAG', 'create_resources_instance',
 ]
 
 

v0/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):

v0/relationalai/clients/resources/snowflake/direct_access_resources.py

@@ -12,12 +12,13 @@ from ....environments import runtime_env, SnowbookEnvironment
 from ...config import Config, ConfigStore, ENDPOINT_FILE
 from ...direct_access_client import DirectAccessClient
 from ...types import EngineState
-from ...util import get_pyrel_version, poll_with_specified_overhead, safe_json_loads, ms_to_timestamp
-from ....errors import ResponseStatusException, QueryTimeoutExceededException
+from ...util import get_pyrel_version, safe_json_loads, ms_to_timestamp
+from ....errors import GuardRailsException, ResponseStatusException, QueryTimeoutExceededException, RAIException
 from snowflake.snowpark import Session
 
 # Import UseIndexResources to enable use_index functionality with direct access
 from .use_index_resources import UseIndexResources
+from .snowflake import TxnCreationResult
 
 # Import helper functions from util
 from .util import is_engine_issue as _is_engine_issue, is_database_issue as _is_database_issue, collect_error_messages
@@ -27,6 +28,7 @@ from typing import Iterable
 
 # Constants
 TXN_ABORT_REASON_TIMEOUT = "transaction timeout"
+TXN_ABORT_REASON_GUARD_RAILS = "guard rail violation"
 
 
 class DirectAccessResources(UseIndexResources):
@@ -161,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,
                     )
@@ -217,83 +220,59 @@ class DirectAccessResources(UseIndexResources):
 
         return response
 
-    def _exec_async_v2(
+    def _create_v2_txn(
         self,
         database: str,
-        engine: Union[str, None],
+        engine: str | None,
         raw_code: str,
-        inputs: Dict | None = None,
-        readonly=True,
-        nowait_durable=False,
-        headers: Dict[str, str] | None = None,
-        bypass_index=False,
-        language: str = "rel",
-        query_timeout_mins: int | None = None,
-        gi_setup_skipped: bool = False,
-    ):
-
-        with debugging.span("transaction") as txn_span:
-            with debugging.span("create_v2") as create_span:
-
-                use_graph_index = self.config.get("use_graph_index", USE_GRAPH_INDEX)
-
-                payload = {
-                    "dbname": database,
-                    "engine_name": engine,
-                    "query": raw_code,
-                    "v1_inputs": inputs,
-                    "nowait_durable": nowait_durable,
-                    "readonly": readonly,
-                    "language": language,
-                }
-                if query_timeout_mins is None and (timeout_value := self.config.get("query_timeout_mins", DEFAULT_QUERY_TIMEOUT_MINS)) is not None:
-                    query_timeout_mins = int(timeout_value)
-                if query_timeout_mins is not None:
-                    payload["timeout_mins"] = query_timeout_mins
-                query_params={"use_graph_index": str(use_graph_index and not bypass_index)}
-
-                # Add gi_setup_skipped to headers
-                if headers is None:
-                    headers = {}
-                headers["gi_setup_skipped"] = str(gi_setup_skipped)
-                headers['pyrel_program_id'] = debugging.get_program_span_id() or ""
-
-                response = self._txn_request_with_gi_retry(
-                    payload, headers, query_params, engine
-                )
-
-                artifact_info = {}
-                response_content = response.json()
+        inputs: Dict,
+        headers: Dict[str, str],
+        readonly: bool,
+        nowait_durable: bool,
+        bypass_index: bool,
+        language: str,
+        query_timeout_mins: int | None,
+    ) -> TxnCreationResult:
+        """
+        Create a transaction via direct HTTP access and return the result.
 
-                txn_id = response_content["transaction"]['id']
-                state = response_content["transaction"]['state']
+        This override uses HTTP requests instead of SQL stored procedures.
+        """
+        use_graph_index = self.config.get("use_graph_index", USE_GRAPH_INDEX)
 
-                txn_span["txn_id"] = txn_id
-                create_span["txn_id"] = txn_id
-                debugging.event("transaction_created", txn_span, txn_id=txn_id)
+        payload = {
+            "dbname": database,
+            "engine_name": engine,
+            "query": raw_code,
+            "v1_inputs": inputs,
+            "nowait_durable": nowait_durable,
+            "readonly": readonly,
+            "language": language,
+        }
+        if query_timeout_mins is None and (timeout_value := self.config.get("query_timeout_mins", DEFAULT_QUERY_TIMEOUT_MINS)) is not None:
+            query_timeout_mins = int(timeout_value)
+        if query_timeout_mins is not None:
+            payload["timeout_mins"] = query_timeout_mins
+        query_params = {"use_graph_index": str(use_graph_index and not bypass_index)}
+
+        response = self._txn_request_with_gi_retry(
+            payload, headers, query_params, engine
+        )
 
-            # fast path: transaction already finished
-            if state in ["COMPLETED", "ABORTED"]:
-                if txn_id in self._pending_transactions:
-                    self._pending_transactions.remove(txn_id)
+        response_content = response.json()
 
-                # Process rows to get the rest of the artifacts
-                for result in response_content.get("results", []):
-                    filename = result['filename']
-                    # making keys uppercase to match the old behavior
-                    artifact_info[filename] = {k.upper(): v for k, v in result.items()}
+        txn_id = response_content["transaction"]['id']
+        state = response_content["transaction"]['state']
 
-            # Slow path: transaction not done yet; start polling
-            else:
-                self._pending_transactions.append(txn_id)
-                with debugging.span("wait", txn_id=txn_id):
-                    poll_with_specified_overhead(
-                        lambda: self._check_exec_async_status(txn_id, headers=headers), 0.1
-                    )
-                artifact_info = self._list_exec_async_artifacts(txn_id, headers=headers)
+        # Build artifact_info if transaction completed immediately (fast path)
+        artifact_info: Dict[str, Dict] = {}
+        if state in ["COMPLETED", "ABORTED"]:
+            for result in response_content.get("results", []):
+                filename = result['filename']
+                # making keys uppercase to match the old behavior
+                artifact_info[filename] = {k.upper(): v for k, v in result.items()}
 
-            with debugging.span("fetch"):
-                return self._download_results(artifact_info, txn_id, state)
+        return TxnCreationResult(txn_id=txn_id, state=state, artifact_info=artifact_info)
 
     def _prepare_index(
         self,
@@ -355,15 +334,20 @@ class DirectAccessResources(UseIndexResources):
             if txn_id in self._pending_transactions:
                 self._pending_transactions.remove(txn_id)
 
-        if status == "ABORTED" and transaction.get("abort_reason", "") == TXN_ABORT_REASON_TIMEOUT:
-            config_file_path = getattr(self.config, 'file_path', None)
-            timeout_ms = int(transaction.get("timeout_ms", 0))
-            timeout_mins = timeout_ms // 60000 if timeout_ms > 0 else int(self.config.get("query_timeout_mins", DEFAULT_QUERY_TIMEOUT_MINS) or DEFAULT_QUERY_TIMEOUT_MINS)
-            raise QueryTimeoutExceededException(
-                timeout_mins=timeout_mins,
-                query_id=txn_id,
-                config_file_path=config_file_path,
-            )
+        if status == "ABORTED":
+            reason = transaction.get("abort_reason", "")
+
+            if reason == TXN_ABORT_REASON_TIMEOUT:
+                config_file_path = getattr(self.config, 'file_path', None)
+                timeout_ms = int(transaction.get("timeout_ms", 0))
+                timeout_mins = timeout_ms // 60000 if timeout_ms > 0 else int(self.config.get("query_timeout_mins", DEFAULT_QUERY_TIMEOUT_MINS) or DEFAULT_QUERY_TIMEOUT_MINS)
+                raise QueryTimeoutExceededException(
+                    timeout_mins=timeout_mins,
+                    query_id=txn_id,
+                    config_file_path=config_file_path,
+                )
+            elif reason == TXN_ABORT_REASON_GUARD_RAILS:
+                raise GuardRailsException(response_content.get("progress", {}))
 
         # @TODO: Find some way to tunnel the ABORT_REASON out. Azure doesn't have this, but it's handy
         return status == "COMPLETED" or status == "ABORTED"
@@ -530,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(
@@ -543,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:
@@ -569,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"],
@@ -576,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,
         )
@@ -609,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:
@@ -621,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:
@@ -632,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: