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

relationalai/clients/exec_txn_poller.py

@@ -0,0 +1,91 @@
+from __future__ import annotations
+
+import time
+from typing import Dict, Optional, TYPE_CHECKING
+
+from relationalai import debugging
+from relationalai.clients.util import poll_with_specified_overhead
+from relationalai.tools.cli_controls import create_progress
+from relationalai.util.format import format_duration
+
+if TYPE_CHECKING:
+    from relationalai.clients.resources.snowflake import Resources
+
+# Polling behavior constants
+POLL_OVERHEAD_RATE = 0.1  # Overhead rate for exponential backoff
+
+# Text color constants
+GREEN_COLOR = '\033[92m'
+GRAY_COLOR = '\033[90m'
+ENDC = '\033[0m'
+
+
+class ExecTxnPoller:
+    """
+    Encapsulates the polling logic for exec_async transaction completion.
+    """
+
+    def __init__(
+        self,
+        resource: "Resources",
+        txn_id: str,
+        headers: Optional[Dict] = None,
+        txn_start_time: Optional[float] = None,
+    ):
+        self.res = resource
+        self.txn_id = txn_id
+        self.headers = headers or {}
+        self.txn_start_time = txn_start_time or time.time()
+
+    def poll(self) -> bool:
+        """
+        Poll for transaction completion with interactive progress display.
+
+        Returns:
+            True if transaction completed successfully, False otherwise
+        """
+
+        # 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)
+            )
+
+        return True
+
+def query_progress_message(id: str, duration: float) -> str:
+    return (
+        # 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}"
+    )
+
+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}"
+    )

relationalai/clients/resources/snowflake/__init__.py

@@ -2,7 +2,7 @@
 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, INTERNAL_ENGINE_SIZES, ENGINE_SIZES_AWS, ENGINE_SIZES_AZURE, PrimaryKey, PRINT_TXN_PROGRESS_FLAG
 
 # These imports depend on Resources, so they come after
 from .cli_resources import CLIResources
@@ -14,7 +14,7 @@ __all__ = [
     'Resources', 'DirectAccessResources', 'Provider', 'Graph', 'SnowflakeClient',
     'APP_NAME', 'PYREL_ROOT_DB', 'CLIResources', 'UseIndexResources', 'ExecContext',
     'INTERNAL_ENGINE_SIZES', 'ENGINE_SIZES_AWS', 'ENGINE_SIZES_AZURE', 'PrimaryKey',
-    'create_resources_instance',
+    'PRINT_TXN_PROGRESS_FLAG', 'create_resources_instance',
 ]
 
 

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
 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):
@@ -217,83 +219,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 +333,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"

relationalai/clients/resources/snowflake/snowflake.py

@@ -15,6 +15,7 @@ import hashlib
 from dataclasses import dataclass
 
 from ....auth.token_handler import TokenHandler
+from relationalai.clients.exec_txn_poller import ExecTxnPoller, query_complete_message
 import snowflake.snowpark
 
 from ....rel_utils import sanitize_identifier, to_fqn_relation_name
@@ -54,7 +55,7 @@ from .util import (
 )
 from ....environments import runtime_env, HexEnvironment, SnowbookEnvironment
 from .... import dsl, rel, metamodel as m
-from ....errors import EngineProvisioningFailed, EngineNameValidationException, Errors, InvalidAliasError, InvalidEngineSizeError, InvalidSourceTypeWarning, RAIException, HexSessionException, SnowflakeChangeTrackingNotEnabledException, SnowflakeDatabaseException, SnowflakeImportMissingException, SnowflakeInvalidSource, SnowflakeMissingConfigValuesException, SnowflakeProxyAPIDeprecationWarning, SnowflakeProxySourceError, ModelNotFoundException, UnknownSourceWarning, RowsDroppedFromTargetTableWarning, QueryTimeoutExceededException
+from ....errors import EngineProvisioningFailed, EngineNameValidationException, Errors, GuardRailsException, InvalidAliasError, InvalidEngineSizeError, InvalidSourceTypeWarning, RAIException, HexSessionException, SnowflakeChangeTrackingNotEnabledException, SnowflakeDatabaseException, SnowflakeImportMissingException, SnowflakeInvalidSource, SnowflakeMissingConfigValuesException, SnowflakeProxyAPIDeprecationWarning, SnowflakeProxySourceError, ModelNotFoundException, UnknownSourceWarning, RowsDroppedFromTargetTableWarning, QueryTimeoutExceededException
 from concurrent.futures import ThreadPoolExecutor
 from datetime import datetime, timedelta
 from snowflake.snowpark.types import StringType, StructField, StructType
@@ -105,6 +106,16 @@ PYREL_ROOT_DB = 'pyrel_root_db'
 TERMINAL_TXN_STATES = ["COMPLETED", "ABORTED"]
 
 TXN_ABORT_REASON_TIMEOUT = "transaction timeout"
+GUARDRAILS_ABORT_REASON = "guard rail violation"
+
+PRINT_TXN_PROGRESS_FLAG = "print_txn_progress"
+
+#--------------------------------------------------
+# Helpers
+#--------------------------------------------------
+
+def should_print_txn_progress(config) -> bool:
+    return bool(config.get(PRINT_TXN_PROGRESS_FLAG, False))
 
 #--------------------------------------------------
 # Resources
@@ -131,6 +142,19 @@ class ExecContext:
             skip_engine_db_error_retry=self.skip_engine_db_error_retry
         )
 
+
+@dataclass
+class TxnCreationResult:
+    """Result of creating a transaction via _create_v2_txn.
+
+    This standardizes the response format between different implementations
+    (SQL stored procedure vs HTTP direct access).
+    """
+    txn_id: str
+    state: str
+    artifact_info: Dict[str, Dict]  # Populated if fast-path (state is COMPLETED/ABORTED)
+
+
 class Resources(ResourcesBase):
     def __init__(
         self,
@@ -1411,15 +1435,18 @@ Otherwise, remove it from your '{profile}' configuration profile.
             if txn_id in self._pending_transactions:
                 self._pending_transactions.remove(txn_id)
 
-        if status == "ABORTED" and response_row.get("ABORT_REASON", "") == TXN_ABORT_REASON_TIMEOUT:
-            config_file_path = getattr(self.config, 'file_path', None)
-            # todo: use the timeout returned alongside the transaction as soon as it's exposed
-            timeout_mins = 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":
+            if response_row.get("ABORT_REASON", "") == TXN_ABORT_REASON_TIMEOUT:
+                config_file_path = getattr(self.config, 'file_path', None)
+                # todo: use the timeout returned alongside the transaction as soon as it's exposed
+                timeout_mins = 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 response_row.get("ABORT_REASON", "") == GUARDRAILS_ABORT_REASON:
+                raise GuardRailsException()
 
         # @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"
@@ -1654,6 +1681,72 @@ Otherwise, remove it from your '{profile}' configuration profile.
             raise Exception("Failed to create transaction")
         return response
 
+    def _create_v2_txn(
+        self,
+        database: str,
+        engine: str | None,
+        raw_code: str,
+        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 and return the result.
+
+        This method handles calling the RAI app stored procedure to create a transaction
+        and parses the response into a standardized TxnCreationResult format.
+
+        This method can be overridden by subclasses (e.g., DirectAccessResources)
+        to use different transport mechanisms (HTTP instead of SQL).
+
+        Args:
+            database: Database/model name
+            engine: Engine name (optional)
+            raw_code: Code to execute (REL, LQP, or SQL)
+            inputs: Input parameters for the query
+            headers: HTTP headers (must be prepared by caller)
+            readonly: Whether the transaction is read-only
+            nowait_durable: Whether to wait for durable writes
+            bypass_index: Whether to bypass graph index setup
+            language: Query language ("rel" or "lqp")
+            query_timeout_mins: Optional query timeout in minutes
+
+        Returns:
+            TxnCreationResult containing txn_id, state, and artifact_info
+        """
+        response = self._exec_rai_app(
+            database=database,
+            engine=engine,
+            raw_code=raw_code,
+            inputs=inputs,
+            readonly=readonly,
+            nowait_durable=nowait_durable,
+            request_headers=headers,
+            bypass_index=bypass_index,
+            language=language,
+            query_timeout_mins=query_timeout_mins,
+        )
+
+        rows = list(iter(response))
+
+        # process the first row since txn_id and state are the same for all rows
+        first_row = rows[0]
+        txn_id = first_row['ID']
+        state = first_row['STATE']
+
+        # Build artifact_info if transaction completed immediately (fast path)
+        artifact_info: Dict[str, Dict] = {}
+        if state in ["COMPLETED", "ABORTED"]:
+            for row in rows:
+                filename = row['FILENAME']
+                artifact_info[filename] = row
+
+        return TxnCreationResult(txn_id=txn_id, state=state, artifact_info=artifact_info)
+
     def _exec_async_v2(
         self,
         database: str,
@@ -1672,15 +1765,20 @@ Otherwise, remove it from your '{profile}' configuration profile.
         High-level async execution method with transaction polling and artifact management.
 
         This is the core method for executing queries asynchronously. It:
-        1. Creates a transaction by calling _exec_rai_app
+        1. Creates a transaction by calling _create_v2_txn
         2. Handles two execution paths:
             - Fast path: Transaction completes immediately (COMPLETED/ABORTED)
             - Slow path: Transaction is pending, requires polling until completion
         3. Manages pending transactions list
         4. Downloads and returns query results/artifacts
 
-        This method is called by _execute_code (base implementation) and can be
-        overridden by child classes (e.g., DirectAccessResources uses HTTP instead).
+        This method is called by _execute_code (base implementation), and calls the
+        following methods that can be overridden by child classes (e.g.,
+        DirectAccessResources uses HTTP instead):
+            - _create_v2_txn
+            - _check_exec_async_status
+            - _list_exec_async_artifacts
+            - _download_results
 
         Args:
             database: Database/model name
@@ -1704,53 +1802,62 @@ Otherwise, remove it from your '{profile}' configuration profile.
         query_attrs_dict = json.loads(request_headers.get("X-Query-Attributes", "{}"))
 
         with debugging.span("transaction", **query_attrs_dict) as txn_span:
+            txn_start_time = time.time()
             with debugging.span("create_v2", **query_attrs_dict) as create_span:
+                # Prepare headers for transaction creation
                 request_headers['user-agent'] = get_pyrel_version(self.generation)
                 request_headers['gi_setup_skipped'] = str(gi_setup_skipped)
                 request_headers['pyrel_program_id'] = debugging.get_program_span_id() or ""
-                response = self._exec_rai_app(
+
+                # Create the transaction
+                result = self._create_v2_txn(
                     database=database,
                     engine=engine,
                     raw_code=raw_code,
                     inputs=inputs,
+                    headers=request_headers,
                     readonly=readonly,
                     nowait_durable=nowait_durable,
-                    request_headers=request_headers,
                     bypass_index=bypass_index,
                     language=language,
                     query_timeout_mins=query_timeout_mins,
                 )
 
-                artifact_info = {}
-                rows = list(iter(response))
-
-                # process the first row since txn_id and state are the same for all rows
-                first_row = rows[0]
-                txn_id = first_row['ID']
-                state = first_row['STATE']
-                filename = first_row['FILENAME']
+                txn_id = result.txn_id
+                state = result.state
 
                 txn_span["txn_id"] = txn_id
                 create_span["txn_id"] = txn_id
                 debugging.event("transaction_created", txn_span, txn_id=txn_id)
 
+            print_txn_progress = should_print_txn_progress(self.config)
+
             # fast path: transaction already finished
             if state in ["COMPLETED", "ABORTED"]:
+                txn_end_time = time.time()
                 if txn_id in self._pending_transactions:
                     self._pending_transactions.remove(txn_id)
 
-                # Process rows to get the rest of the artifacts
-                for row in rows:
-                    filename = row['FILENAME']
-                    artifact_info[filename] = row
+                artifact_info = result.artifact_info
+
+                txn_duration = txn_end_time - txn_start_time
+                if print_txn_progress:
+                    print(
+                        query_complete_message(txn_id, txn_duration, status_header=True)
+                    )
 
             # Slow path: transaction not done yet; start polling
             else:
                 self._pending_transactions.append(txn_id)
+                # Use the interactive poller for transaction status
                 with debugging.span("wait", txn_id=txn_id):
-                    poll_with_specified_overhead(
-                        lambda: self._check_exec_async_status(txn_id, headers=request_headers), 0.1
-                    )
+                    if print_txn_progress:
+                        poller = ExecTxnPoller(resource=self, txn_id=txn_id, headers=request_headers, txn_start_time=txn_start_time)
+                        poller.poll()
+                    else:
+                        poll_with_specified_overhead(
+                            lambda: self._check_exec_async_status(txn_id, headers=request_headers), 0.1
+                        )
                 artifact_info = self._list_exec_async_artifacts(txn_id, headers=request_headers)
 
             with debugging.span("fetch"):
@@ -2408,7 +2515,7 @@ Otherwise, remove it from your '{profile}' configuration profile.
             return None
         return results[0][0]
 
-    # CLI methods (list_warehouses, list_compute_pools, list_roles, list_apps, 
+    # CLI methods (list_warehouses, list_compute_pools, list_roles, list_apps,
     # list_databases, list_sf_schemas, list_tables) are now in CLIResources class
     # schema_info is kept in base Resources class since it's used by SnowflakeSchema._fetch_info()
 

relationalai/clients/resources/snowflake/use_index_poller.py

@@ -189,6 +189,9 @@ class UseIndexPoller:
         # on every 5th iteration we reset the cdc status, so it will be checked again
         self.should_check_cdc = True
 
+        # Flag to only check data stream health once in the first call
+        self.check_data_stream_health = True
+
         self.wait_for_stream_sync = self.res.config.get(
             "wait_for_stream_sync", WAIT_FOR_STREAM_SYNC
         )
@@ -503,6 +506,7 @@ class UseIndexPoller:
                 "init_engine_async": self.init_engine_async,
                 "language": self.language,
                 "data_freshness_mins": self.data_freshness,
+                "check_data_stream_health": self.check_data_stream_health
             })
 
             request_headers = debugging.add_current_propagation_headers(self.headers)
@@ -535,6 +539,7 @@ class UseIndexPoller:
                 errors = use_index_data.get("errors", [])
                 relations = use_index_data.get("relations", {})
                 cdc_enabled = use_index_data.get("cdcEnabled", False)
+                health_checked = use_index_data.get("healthChecked", False)
                 if self.check_ready_count % ERP_CHECK_FREQUENCY == 0 or not cdc_enabled:
                     self.should_check_cdc = True
                 else:
@@ -542,6 +547,9 @@ class UseIndexPoller:
 
                 if engines and self.init_engine_async:
                     self.init_engine_async = False
+                
+                if self.check_data_stream_health and health_checked:
+                    self.check_data_stream_health = False
 
                 break_loop = False
                 has_stream_errors = False

relationalai/errors.py

@@ -2436,6 +2436,24 @@ class QueryTimeoutExceededException(RAIException):
         Consider increasing the 'query_timeout_mins' parameter in your configuration file{f' (stored in {self.config_file_path})' if self.config_file_path else ''} to allow more time for query execution.
         """)
 
+class GuardRailsException(RAIException):
+    def __init__(self, progress: dict[str, Any]={}):
+        self.name = "Guard Rails Violation"
+        self.message = "Transaction aborted due to guard rails violation."
+        self.progress = progress
+        self.content = self.format_message()
+        super().__init__(self.message, self.name, self.content)
+
+    def format_message(self):
+        messages = [] if self.progress else [self.message]
+        for task in self.progress.get("tasks", {}).values():
+            for warning_type, warning_data in task.get("warnings", {}).items():
+                messages.append(textwrap.dedent(f"""
+                Relation Name: [yellow]{task["task_name"]}[/yellow]
+                Warning: {warning_type}
+                Message: {warning_data["message"]}
+                """))
+        return "\n".join(messages)
 
 #--------------------------------------------------
 # Azure Exceptions

relationalai/semantics/lqp/rewrite/extract_keys.py

@@ -424,6 +424,12 @@ class ExtractKeysRewriter(Rewriter):
                             for arg in args[:-1]:
                                 extended_vars.add(arg)
                             there_is_progress = True
+                    elif isinstance(task, ir.Not):
+                        if isinstance(task.task, ir.Logical):
+                            hoisted = helpers.hoisted_vars(task.task.hoisted)
+                            if var in hoisted:
+                                partitions[var].add(task)
+                                there_is_progress = True
                     else:
                         assert False, f"invalid node kind {type(task)}"
 

relationalai/semantics/reasoners/optimization/solvers_pb.py

@@ -1,8 +1,8 @@
-"""Solver model implementation using protobuf format.
+"""Solver model implementation supporting protobuf and CSV formats.
 
 This module provides the SolverModelPB class for defining optimization and
-constraint programming problems that are serialized to protobuf format and
-solved by external solver engines.
+constraint programming problems that are serialized and solved by external
+solver engines. Supports both protobuf (default) and CSV (future) exchange formats.
 
 Note: This protobuf-based implementation will be deprecated in favor of the
       development version (solvers_dev.py) in future releases.
@@ -23,7 +23,6 @@ from relationalai.util.timeout import calc_remaining_timeout_minutes
 
 from .common import make_name
 
-
 # =============================================================================
 # Solver ProtoBuf Format Constants and Helpers
 # =============================================================================
@@ -191,6 +190,7 @@ class SolverModelPB:
         """
         b.define(b.RawSource("rel", textwrap.dedent(install_rel)))
 
+
     # -------------------------------------------------------------------------
     # Variable Handling
     # -------------------------------------------------------------------------
@@ -501,69 +501,218 @@ class SolverModelPB:
     # Solving and Result Handling
     # -------------------------------------------------------------------------
 
-    def solve(
-        self, solver: Solver, log_to_console: bool = False, **kwargs: Any
+    def _export_model_to_csv(
+        self,
+        model_id: str,
+        executor: RelExecutor,
+        prefix_lowercase: str,
+        query_timeout_mins: Optional[int] = None
     ) -> None:
-        """Solve the model.
+        """Export model to CSV files in Snowflake stage.
 
         Args:
-            solver: Solver instance.
-            log_to_console: Whether to show solver output.
-            **kwargs: Solver options and parameters.
+            model_id: Unique model identifier for stage paths.
+            executor: RelExecutor instance.
+            prefix_lowercase: Prefix for relation names.
+            query_timeout_mins: Query timeout in minutes.
         """
-        options = {**kwargs, "version": 1}
+        stage_base_no_txn = f"snowflake://APP_STATE.RAI_INTERNAL_STAGE/SOLVERS/job_{model_id}"
+
+        # Export all model relations using Rel-native export_csv in a single transaction
+        # Transformations (uuid_string, encode_base64) are done inline in the export query
+        export_rel = textwrap.dedent(f"""
+        // Get transaction ID for folder naming - solver service validates ownership
+        // Use uuid_string to get proper UUID format, then replace hyphens with underscores
+        def txn_id_str {{string_replace[uuid_string[current_transaction_id], "-", "_"]}}
+
+        // Define base path with txn_id in folder name: model_{{txn_id}}/
+        def base_path {{"{stage_base_no_txn}/model"}}
+
+        // Export variable_hash.csv - single column: HASH (UUID string)
+        // Transformation: convert Variable UInt128 to UUID string inline
+        def variable_hash_data(:HASH, v, h):
+            {self.Variable._name}(v) and uuid_string(v, h)
+
+        def export[:variable_hash]: {{export_csv[{{
+            (:path, base_path ++ "/variable_hash_" ++ txn_id_str ++ ".csv");
+            (:data, variable_hash_data);
+            (:compression, "gzip")
+        }}]}}
+
+        // Export variable_name.csv - columns: HASH (UUID string), VALUE (name string)
+        // Transformation: convert Variable UInt128 to UUID string inline
+        def variable_name_data(:HASH, v, h):
+            {prefix_lowercase}variable_name(v, _) and uuid_string(v, h)
+        def variable_name_data(:VALUE, v, name):
+            {prefix_lowercase}variable_name(v, name)
+
+        def export[:variable_name]: {{export_csv[{{
+            (:path, base_path ++ "/variable_name_" ++ txn_id_str ++ ".csv");
+            (:data, variable_name_data);
+            (:compression, "gzip")
+        }}]}}
+
+        // Export constraint.csv - single column: VALUE (base64 encoded constraint)
+        // Transformation: encode_base64 done inline
+        def constraint_data(:VALUE, c, e):
+            exists((s) |
+                {self.Constraint._name}(c) and
+                {prefix_lowercase}constraint_serialized(c, s) and
+                encode_base64(s, e))
+
+        def export[:constraint]: {{export_csv[{{
+            (:path, base_path ++ "/constraint_" ++ txn_id_str ++ ".csv");
+            (:data, constraint_data);
+            (:compression, "gzip")
+        }}]}}
+
+        // Export min_objective.csv - columns: HASH (UUID string), VALUE (base64 encoded)
+        // Transformations: uuid_string and encode_base64 done inline
+        def min_objective_data(:HASH, obj, h):
+            {self.MinObjective._name}(obj) and uuid_string(obj, h)
+        def min_objective_data(:VALUE, obj, e):
+            exists((s) |
+                {self.MinObjective._name}(obj) and
+                {prefix_lowercase}minobjective_serialized(obj, s) and
+                encode_base64(s, e))
+
+        def export[:min_objective]: {{export_csv[{{
+            (:path, base_path ++ "/min_objective_" ++ txn_id_str ++ ".csv");
+            (:data, min_objective_data);
+            (:compression, "gzip")
+        }}]}}
+
+        // Export max_objective.csv - columns: HASH (UUID string), VALUE (base64 encoded)
+        // Transformations: uuid_string and encode_base64 done inline
+        def max_objective_data(:HASH, obj, h):
+            {self.MaxObjective._name}(obj) and uuid_string(obj, h)
+        def max_objective_data(:VALUE, obj, e):
+            exists((s) |
+                {self.MaxObjective._name}(obj) and
+                {prefix_lowercase}maxobjective_serialized(obj, s) and
+                encode_base64(s, e))
+
+        def export[:max_objective]: {{export_csv[{{
+            (:path, base_path ++ "/max_objective_" ++ txn_id_str ++ ".csv");
+            (:data, max_objective_data);
+            (:compression, "gzip")
+        }}]}}
+        """)
+
+        executor.execute_raw(export_rel, readonly=False, query_timeout_mins=query_timeout_mins)
+
+    def _import_solver_results_from_csv(
+        self,
+        model_id: str,
+        executor: RelExecutor,
+        prefix_lowercase: str,
+        query_timeout_mins: Optional[int] = None
+    ) -> None:
+        """Import solver results from CSV files in Snowflake stage.
 
-        # Validate solver options
-        for option_key, option_value in options.items():
-            if not isinstance(option_key, str):
-                raise TypeError(
-                    f"Solver option keys must be strings, but got {type(option_key).__name__} for key {option_key!r}."
-                )
-            if not isinstance(option_value, (int, float, str, bool)):
-                raise TypeError(
-                    f"Solver option values must be int, float, str, or bool, "
-                    f"but got {type(option_value).__name__} for option {option_key!r}."
-                )
+        Loads and extracts CSV files in a single transaction to minimize overhead.
 
-        # Three-phase solve process:
-        # 1. Export model to Snowflake as protobuf
-        # 2. Execute solver job (external solver reads from Snowflake)
-        # 3. Extract and load results back into the model
-        input_id = uuid.uuid4()
-        model_uri = f"snowflake://APP_STATE.RAI_INTERNAL_STAGE/job-inputs/solver/{input_id}/model.binpb"
-        sf_input_uri = f"snowflake://job-inputs/solver/{input_id}/model.binpb"
-        payload: dict[str, Any] = {"solver": solver.solver_name.lower()}
-        payload["options"] = options
-        payload["model_uri"] = sf_input_uri
+        Args:
+            model_id: Unique model identifier for stage paths.
+            executor: RelExecutor instance.
+            prefix_lowercase: Prefix for relation names.
+            query_timeout_mins: Query timeout in minutes.
+        """
+        result_stage_base = f"snowflake://APP_STATE.RAI_INTERNAL_STAGE/SOLVERS/job_{model_id}/results"
+        value_parse_fn = "parse_int" if self._num_type == "int" else "parse_float"
+
+        # Single transaction: Load CSV files and extract/map results
+        # Use inline definitions to avoid needing declared relations
+        load_and_extract_rel = textwrap.dedent(f"""
+        // Define CSV loading inline (no declare needed)
+        // Load ancillary.csv - contains solver metadata (NAME, VALUE columns)
+        def ancillary_config[:path]: "{result_stage_base}/ancillary.csv.gz"
+        def ancillary_config[:syntax, :header_row]: 1
+        def ancillary_config[:schema, :NAME]: "string"
+        def ancillary_config[:schema, :VALUE]: "string"
+        def {prefix_lowercase}solver_ancillary_raw {{load_csv[ancillary_config]}}
+
+        // Load objective_values.csv - contains objective values (SOL_INDEX, VALUE columns)
+        def objective_values_config[:path]: "{result_stage_base}/objective_values.csv.gz"
+        def objective_values_config[:syntax, :header_row]: 1
+        def objective_values_config[:schema, :SOL_INDEX]: "string"
+        def objective_values_config[:schema, :VALUE]: "string"
+        def {prefix_lowercase}solver_objective_values_raw {{load_csv[objective_values_config]}}
+
+        // Load points.csv.gz - contains solution points (SOL_INDEX, VAR_HASH, VALUE columns)
+        def points_config[:path]: "{result_stage_base}/points.csv.gz"
+        def points_config[:syntax, :header_row]: 1
+        def points_config[:schema, :SOL_INDEX]: "string"
+        def points_config[:schema, :VAR_HASH]: "string"
+        def points_config[:schema, :VALUE]: "string"
+        def {prefix_lowercase}solver_points_raw {{load_csv[points_config]}}
+
+        // Clear existing result data
+        def delete[:{self.result_info._name}]: {self.result_info._name}
+        def delete[:{self.point._name}]: {self.point._name}
+        def delete[:{self.points._name}]: {self.points._name}
 
-        executor = self._model._to_executor()
-        if not isinstance(executor, RelExecutor):
-            raise ValueError(f"Expected RelExecutor, got {type(executor).__name__}.")
-        prefix_lowercase = f"solvermodel_{self._id}_"
+        // Extract ancillary data (result info) - NAME and VALUE columns
+        def insert(:{self.result_info._name}, key, val): {{
+            exists((row) |
+                {prefix_lowercase}solver_ancillary_raw(:NAME, row, key) and
+                {prefix_lowercase}solver_ancillary_raw(:VALUE, row, val))
+        }}
 
-        query_timeout_mins = kwargs.get("query_timeout_mins", None)
-        config = self._model._config
-        if (
-            query_timeout_mins is None
-            and (
-                timeout_value := config.get(
-                    "query_timeout_mins", DEFAULT_QUERY_TIMEOUT_MINS
-                )
-            )
-            is not None
-        ):
-            query_timeout_mins = int(timeout_value)
-        config_file_path = getattr(config, "file_path", None)
-        start_time = time.monotonic()
-        remaining_timeout_minutes = query_timeout_mins
+        // Extract objective value from objective_values CSV (first solution)
+        def insert(:{self.result_info._name}, "objective_value", val): {{
+            exists((row) |
+                {prefix_lowercase}solver_objective_values_raw(:SOL_INDEX, row, "1") and
+                {prefix_lowercase}solver_objective_values_raw(:VALUE, row, val))
+        }}
 
-        # Step 1: Materialize the model and store it in Snowflake
-        print("export model")
-        # TODO(coey): Weird hack to avoid uninitialized properties error
-        # This forces evaluation of the Variable concept before export
-        b.select(b.count(self.Variable)).to_df()
-        export_model_relation = f"""
-        // TODO maybe only want to pass names if printing - like in old setup
+        // Extract solution points from points.csv.gz into points property
+        // This file has SOL_INDEX, VAR_HASH, VALUE columns
+        // Convert CSV string index to Int128 for points property signature
+        // Convert value to Int128 (for int) or Float64 (for float)
+        def insert(:{self.points._name}, sol_idx_int128, var, val_converted): {{
+            exists((row, sol_idx_str, var_hash_str, val_str, sol_idx_int, val) |
+                {prefix_lowercase}solver_points_raw(:SOL_INDEX, row, sol_idx_str) and
+                {prefix_lowercase}solver_points_raw(:VAR_HASH, row, var_hash_str) and
+                {prefix_lowercase}solver_points_raw(:VALUE, row, val_str) and
+                parse_int(sol_idx_str, sol_idx_int) and
+                parse_uuid(var_hash_str, var) and
+                {value_parse_fn}(val_str, val) and
+                ::std::mirror::convert(std::mirror::typeof[Int128], sol_idx_int, sol_idx_int128) and
+                {'::std::mirror::convert(std::mirror::typeof[Int128], val, val_converted)' if self._num_type == 'int' else '::std::mirror::convert(std::mirror::typeof[Float64], val, val_converted)'})
+        }}
+
+        // Extract first solution into point property (default solution)
+        // Filter to SOL_INDEX = 1
+        def insert(:{self.point._name}, var, val_converted): {{
+            exists((row, var_hash_str, val_str, val) |
+                {prefix_lowercase}solver_points_raw(:SOL_INDEX, row, "1") and
+                {prefix_lowercase}solver_points_raw(:VAR_HASH, row, var_hash_str) and
+                {prefix_lowercase}solver_points_raw(:VALUE, row, val_str) and
+                parse_uuid(var_hash_str, var) and
+                {value_parse_fn}(val_str, val) and
+                {'::std::mirror::convert(std::mirror::typeof[Int128], val, val_converted)' if self._num_type == 'int' else '::std::mirror::convert(std::mirror::typeof[Float64], val, val_converted)'})
+        }}
+        """)
+
+        executor.execute_raw(load_and_extract_rel, readonly=False, query_timeout_mins=query_timeout_mins)
+
+    def _export_model_to_protobuf(
+        self,
+        model_uri: str,
+        executor: RelExecutor,
+        prefix_lowercase: str,
+        query_timeout_mins: Optional[int] = None
+    ) -> None:
+        """Export model to protobuf format in Snowflake stage.
+
+        Args:
+            model_uri: Snowflake URI for the protobuf file.
+            executor: RelExecutor instance.
+            prefix_lowercase: Prefix for relation names.
+            query_timeout_mins: Query timeout in minutes.
+        """
+        export_rel = f"""
         // Collect all model components into a relation for serialization
         def model_relation {{
             (:variable, {self.Variable._name});
@@ -584,31 +733,24 @@ class SolverModelPB:
         def export {{ config }}
         """
         executor.execute_raw(
-            textwrap.dedent(export_model_relation),
-            query_timeout_mins=remaining_timeout_minutes,
+            textwrap.dedent(export_rel),
+            query_timeout_mins=query_timeout_mins
         )
 
-        # Step 2: Execute solver job and wait for completion
-        print("execute solver job")
-        remaining_timeout_minutes = calc_remaining_timeout_minutes(
-            start_time,
-            query_timeout_mins,
-            config_file_path=config_file_path,
-        )
-        job_id = solver._exec_job(
-            payload,
-            log_to_console=log_to_console,
-            query_timeout_mins=remaining_timeout_minutes,
-        )
+    def _import_solver_results_from_protobuf(
+        self,
+        job_id: str,
+        executor: RelExecutor,
+        query_timeout_mins: Optional[int] = None
+    ) -> None:
+        """Import solver results from protobuf format.
 
-        # Step 3: Extract and insert solver results into the model
-        print("extract result")
-        remaining_timeout_minutes = calc_remaining_timeout_minutes(
-            start_time,
-            query_timeout_mins,
-            config_file_path=config_file_path,
-        )
-        extract_results_relation = f"""
+        Args:
+            job_id: Job identifier for result location.
+            executor: RelExecutor instance.
+            query_timeout_mins: Query timeout in minutes.
+        """
+        extract_rel = f"""
         def raw_result {{
             load_binary["snowflake://APP_STATE.RAI_INTERNAL_STAGE/job-results/{job_id}/result.binpb"]
         }}
@@ -625,6 +767,7 @@ class SolverModelPB:
         def insert(:{self.result_info._name}, key, value):
             exists((original_key) | string(extracted[original_key], value) and ::std::mirror::lower(original_key, key))
         """
+
         if self._num_type == "int":
             insert_points_relation = f"""
             def insert(:{self.point._name}, variable, value):
@@ -645,15 +788,123 @@ class SolverModelPB:
                     ::std::mirror::convert(std::mirror::typeof[Int128], float_index, point_index)
                 )
             """
+
         executor.execute_raw(
-            textwrap.dedent(extract_results_relation)
-            + textwrap.dedent(insert_points_relation),
+            textwrap.dedent(extract_rel) + textwrap.dedent(insert_points_relation),
             readonly=False,
-            query_timeout_mins=remaining_timeout_minutes,
+            query_timeout_mins=query_timeout_mins
         )
 
-        print("finished solve")
+    def solve(
+        self, solver: Solver, log_to_console: bool = False, **kwargs: Any
+    ) -> None:
+        """Solve the model.
+
+        Args:
+            solver: Solver instance.
+            log_to_console: Whether to show solver output.
+            **kwargs: Solver options and parameters.
+        """
+
+        use_csv_store = solver.engine_settings.get("store", {})\
+                                        .get("csv", {})\
+                                        .get("enabled", False)
+
+        print(f"Using {'csv' if use_csv_store else 'protobuf'} store...")
+
+        options = {**kwargs, "version": 1}
+
+        # Validate solver options
+        for option_key, option_value in options.items():
+            if not isinstance(option_key, str):
+                raise TypeError(
+                    f"Solver option keys must be strings, but got {type(option_key).__name__} for key {option_key!r}."
+                )
+            if not isinstance(option_value, (int, float, str, bool)):
+                raise TypeError(
+                    f"Solver option values must be int, float, str, or bool, "
+                    f"but got {type(option_value).__name__} for option {option_key!r}."
+                )
+
+        executor = self._model._to_executor()
+        if not isinstance(executor, RelExecutor):
+            raise ValueError(f"Expected RelExecutor, got {type(executor).__name__}.")
+        prefix_lowercase = f"solvermodel_{self._id}_"
+
+        # Initialize timeout from config
+        query_timeout_mins = kwargs.get("query_timeout_mins", None)
+        config = self._model._config
+        if (
+            query_timeout_mins is None
+            and (
+                timeout_value := config.get(
+                    "query_timeout_mins", DEFAULT_QUERY_TIMEOUT_MINS
+                )
+            )
+            is not None
+        ):
+            query_timeout_mins = int(timeout_value)
+        config_file_path = getattr(config, "file_path", None)
+        start_time = time.monotonic()
+
+        # Force evaluation of Variable concept before export
+        b.select(b.count(self.Variable)).to_df()
+
+        # Prepare payload for solver service
+        payload: dict[str, Any] = {"solver": solver.solver_name.lower(), "options": options}
+
+        if use_csv_store:
+            # CSV format: model and results are exchanged via CSV files
+            model_id = str(uuid.uuid4()).upper().replace('-', '_')
+            payload["model_uri"] = f"snowflake://SOLVERS/job_{model_id}/model"
+
+            print("Exporting model to CSV...")
+            remaining_timeout_minutes = calc_remaining_timeout_minutes(
+                start_time, query_timeout_mins, config_file_path=config_file_path
+            )
+            self._export_model_to_csv(model_id, executor, prefix_lowercase, remaining_timeout_minutes)
+            print("Model CSV export completed")
+
+            print("Execute solver job")
+            remaining_timeout_minutes = calc_remaining_timeout_minutes(
+                start_time, query_timeout_mins, config_file_path=config_file_path
+            )
+            solver._exec_job(payload, log_to_console=log_to_console, query_timeout_mins=remaining_timeout_minutes)
+
+            print("Loading and extracting solver results...")
+            remaining_timeout_minutes = calc_remaining_timeout_minutes(
+                start_time, query_timeout_mins, config_file_path=config_file_path
+            )
+            self._import_solver_results_from_csv(model_id, executor, prefix_lowercase, remaining_timeout_minutes)
+
+        else:  # protobuf format
+            # Protobuf format: model and results are exchanged via binary protobuf
+            input_id = uuid.uuid4()
+            model_uri = f"snowflake://APP_STATE.RAI_INTERNAL_STAGE/job-inputs/solver/{input_id}/model.binpb"
+            sf_input_uri = f"snowflake://job-inputs/solver/{input_id}/model.binpb"
+            payload["model_uri"] = sf_input_uri
+
+            print("Export model...")
+            remaining_timeout_minutes = calc_remaining_timeout_minutes(
+                start_time, query_timeout_mins, config_file_path=config_file_path
+            )
+            self._export_model_to_protobuf(model_uri, executor, prefix_lowercase, remaining_timeout_minutes)
+
+            print("Execute solver job...")
+            remaining_timeout_minutes = calc_remaining_timeout_minutes(
+                start_time, query_timeout_mins, config_file_path=config_file_path
+            )
+            job_id = solver._exec_job(payload, log_to_console=log_to_console, query_timeout_mins=remaining_timeout_minutes)
+
+            print("Extract result...")
+            remaining_timeout_minutes = calc_remaining_timeout_minutes(
+                start_time, query_timeout_mins, config_file_path=config_file_path
+            )
+            self._import_solver_results_from_protobuf(job_id, executor, remaining_timeout_minutes)
+
+        print("Finished solve")
         print()
+        return None
 
     def load_point(self, point_index: int) -> None:
         """Load a solution point.

{relationalai-0.13.0.dist-info → relationalai-0.13.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: relationalai
-Version: 0.13.0
+Version: 0.13.1
 Summary: RelationalAI Library and CLI
 Author-email: RelationalAI <support@relational.ai>
 License-File: LICENSE
@@ -11,7 +11,7 @@ Requires-Dist: colorama
 Requires-Dist: cryptography
 Requires-Dist: gravis
 Requires-Dist: inquirerpy
-Requires-Dist: lqp==0.1.19
+Requires-Dist: lqp==0.2.1
 Requires-Dist: nicegui==2.16.1
 Requires-Dist: numpy<2
 Requires-Dist: opentelemetry-api

{relationalai-0.13.0.dist-info → relationalai-0.13.1.dist-info}/RECORD

@@ -4,7 +4,7 @@ relationalai/debugging.py,sha256=wqGly2Yji4ErdV9F_S1Bny35SiGqI3C7S-hap7B8yUs,119
 relationalai/dependencies.py,sha256=tL113efcISkJUiDXYHmRdU_usdD7gmee-VRHA7N4EFA,16574
 relationalai/docutils.py,sha256=1gVv9mk0ytdMB2W7_NvslJefmSQtTOg8LHTCDcGCjyE,1554
 relationalai/dsl.py,sha256=UJr93X8kwnnyUY-kjPzp_jhsp2pYBUnDfu8mhNXPNII,66116
-relationalai/errors.py,sha256=JeYaycLGFeT0jidSMERGkpN02zswW5TGrIV647robmw,95973
+relationalai/errors.py,sha256=PaWzKsfWwjca7J0huILkiR515_PwJXP1Hlq6owOjL0Y,96828
 relationalai/metagen.py,sha256=o10PNvR_myr_61DC8g6lkB093bFo9qXGUkZKgKyfXiE,26821
 relationalai/metamodel.py,sha256=P1hliwHd1nYxbXON4LZeaYZD6T6pZm97HgmFBFrWyCk,32886
 relationalai/rel.py,sha256=ePmAXx4NxOdsPcHNHyGH3Jkp_cB3QzfKu5p_EQSHPh0,38293
@@ -22,6 +22,7 @@ relationalai/clients/__init__.py,sha256=LQ_yHsutRMpoW2mOTmOPGF8mrbP0OiV5E68t8uVw
 relationalai/clients/client.py,sha256=gk_V9KS7_MM2dLL2OCO7EPLHD9dsRwR6R-30SW8lDwU,35759
 relationalai/clients/config.py,sha256=hERaKjc3l4kd-kf0l-NUOHrWunCn8gmFWpuE0j3ScJg,24457
 relationalai/clients/direct_access_client.py,sha256=VGjQ7wzduxCo04BkxSZjlPAgqK-aBc32zIXcMfAzzSU,6436
+relationalai/clients/exec_txn_poller.py,sha256=JbmrTvsWGwxM5fcZVjeJihQDqhMf-ySfHN3Jn0HGwG0,3108
 relationalai/clients/hash_util.py,sha256=pZVR1FX3q4G_19p_r6wpIR2tIM8_WUlfAR7AVZJjIYM,1495
 relationalai/clients/local.py,sha256=vo5ikSWg38l3xQAh9yL--4sMAj_T5Tn7YEZiw7TCH08,23504
 relationalai/clients/profile_polling.py,sha256=pUH7WKH4nYDD0SlQtg3wsWdj0K7qt6nZqUw8jTthCBs,2565
@@ -30,16 +31,16 @@ relationalai/clients/types.py,sha256=eNo6akcMTbnBFbBbHd5IgVeY-zuAgtXlOs8Bo1SWmVU
 relationalai/clients/util.py,sha256=NJC8fnrWHR01NydwESPSetIHRWf7jQJURYpaWJjmDyE,12311
 relationalai/clients/resources/__init__.py,sha256=pymn8gB86Q3C2bVoFei0KAL8pX_U04uDY9TE4TKzTBs,260
 relationalai/clients/resources/azure/azure.py,sha256=TDapfM5rLoHrPrXg5cUe827m3AO0gSqQjNid1VUlUFo,20631
-relationalai/clients/resources/snowflake/__init__.py,sha256=9VR-hSIw4ZSEWisKcWhNEcRVBmBfueXNCTOOfLt-8rs,871
+relationalai/clients/resources/snowflake/__init__.py,sha256=Ofyf1RZu9GLQdvsjpHDUHEQHHVODb9vKYI4hMOxczH4,923
 relationalai/clients/resources/snowflake/cache_store.py,sha256=A-qd11wcwN3TkIqvlN0_iFUU3aEjJal3T2pqFBwkkzQ,3966
 relationalai/clients/resources/snowflake/cli_resources.py,sha256=xTIcCzvgbkxuNAEvzZoRpj0n-js0hZCK30q7IZXztbI,3252
-relationalai/clients/resources/snowflake/direct_access_resources.py,sha256=Xvh1e6TxUW2dTSS-9HadrfWVKrxNQ5GikqM4yjohJkM,29849
+relationalai/clients/resources/snowflake/direct_access_resources.py,sha256=5LZIQCSiX_TT2tVHqbTBzJTBXZzHmQI7cxFOY6CDDBM,28871
 relationalai/clients/resources/snowflake/engine_state_handlers.py,sha256=SQBu4GfbyABU6xrEV-koivC-ubsVrfCBTF0FEQgJM5g,12054
 relationalai/clients/resources/snowflake/error_handlers.py,sha256=581G2xOihUoiPlucC_Z2FOzhKu_swdIc3uORd0yJQuA,8805
 relationalai/clients/resources/snowflake/export_procedure.py.jinja,sha256=00iLO2qmvJoqAeJUWt3bAsFDDnof7Ab2spggzelbwK4,10566
 relationalai/clients/resources/snowflake/resources_factory.py,sha256=4LGd4IQ6z8hGeGlO1TIjSFJEeUNHutaB7j9q1a9rYfQ,3385
-relationalai/clients/resources/snowflake/snowflake.py,sha256=kWs7W_jsf_iy2N0VO5Uqu9VHXsbH2rQghPn1s5-9GAg,133037
-relationalai/clients/resources/snowflake/use_index_poller.py,sha256=gMcfILcpD1-wkJ0C1rGZQSy1oZBq6Gm8x9N7jdgEkhA,48448
+relationalai/clients/resources/snowflake/snowflake.py,sha256=SuFwAdhaFrd3JZFEjVJMVVFuLrC_T3dqXS27ABrXZAw,136973
+relationalai/clients/resources/snowflake/use_index_poller.py,sha256=xF_9XAymO82YZcZOQqz0r9oZZTqvbWW-WjDTZm4-tFM,48859
 relationalai/clients/resources/snowflake/use_index_resources.py,sha256=69PNWHI_uf-Aw_evfwC6j8HLVdjhp84vs8hLkjnhwbg,6462
 relationalai/clients/resources/snowflake/util.py,sha256=BEnm1B1-nqqHdm41RNxblbb-zqXbtqEGGZmTdAYeN_M,13841
 relationalai/early_access/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -334,7 +335,7 @@ relationalai/semantics/lqp/rewrite/__init__.py,sha256=V9ERED9qdh4VvY9Ud_M8Zn8lhV
 relationalai/semantics/lqp/rewrite/annotate_constraints.py,sha256=b_Ly4_80dQpRzWbeLC72JVfxzhwOPBpiCdEqtBiEiwM,2310
 relationalai/semantics/lqp/rewrite/cdc.py,sha256=8By-BrsayAWYxvrOcoqQ9w0-Nj4KQnaRH8haiog0-4o,10400
 relationalai/semantics/lqp/rewrite/extract_common.py,sha256=ZRvmeYHN8JEkU-j3fRx1e0_JK-46n6NqhxtwZe6L10c,14690
-relationalai/semantics/lqp/rewrite/extract_keys.py,sha256=gC-AkDA4DpqpFkXYc3u-LExbK6e79SrSxs5G6-TEb58,21797
+relationalai/semantics/lqp/rewrite/extract_keys.py,sha256=aw5xvsRN3556ncHWyG90jEuykw66c2cWV7n4QRsJh0Y,22150
 relationalai/semantics/lqp/rewrite/function_annotations.py,sha256=9ZzLASvXh_OgQ04eup0AyoMIh2HxWHkoRETLm1-XtWs,4660
 relationalai/semantics/lqp/rewrite/functional_dependencies.py,sha256=4oQcVQtAGDqY850B1bNszigQopf6y9Y_CaUyWx42PtM,12718
 relationalai/semantics/lqp/rewrite/quantify_vars.py,sha256=bOowgQ45zmP0HOhsTlE92WdVBCTXSkszcCYbPMeIibw,12004
@@ -370,7 +371,7 @@ relationalai/semantics/reasoners/graph/tests/README.md,sha256=XbauTzt6VA_YEOcrlZ
 relationalai/semantics/reasoners/optimization/__init__.py,sha256=lpavly1Qa3VKvLgrbpp-tsxY9hcqHL6buxuekgKPakw,2212
 relationalai/semantics/reasoners/optimization/common.py,sha256=V0c9eGHJKI-gt0X-q9o0bIkgBCdWFdjWq2NQITKwmXg,3124
 relationalai/semantics/reasoners/optimization/solvers_dev.py,sha256=lbw3c8Z6PlHRDm7TdAhICPShlGoab9uR_4uacMPvpBw,24493
-relationalai/semantics/reasoners/optimization/solvers_pb.py,sha256=ESwraHU9c4NCEVRZ16tnBZsUCmJg7lUhy-v0-GGq0qo,48000
+relationalai/semantics/reasoners/optimization/solvers_pb.py,sha256=COJuNm5kyyAQULkc1YS-91op4y4doobZEaE_RBLZWwo,60144
 relationalai/semantics/rel/__init__.py,sha256=pMlVTC_TbQ45mP1LpzwFBBgPxpKc0H3uJDvvDXEWzvs,55
 relationalai/semantics/rel/builtins.py,sha256=kQToiELc4NnvCmXyFtu9CsGZNdTQtSzTB-nuyIfQcsM,1562
 relationalai/semantics/rel/compiler.py,sha256=pFkEbuPKVd8AI4tiklcv06LbNnK8KfoV4FwmY9Lrhqo,43044
@@ -443,7 +444,7 @@ relationalai/util/spans_file_handler.py,sha256=a0sDwDPBBvGsM6be2En3mId9sXpuJlXia
 relationalai/util/timeout.py,sha256=2o6BVNFnFc-B2j-i1pEkZcQbMRto9ps2emci0XwiA4I,783
 relationalai/util/tracing_handler.py,sha256=H919ETAxh7Z1tRz9x8m90qP51_264UunHAPw8Sr6x2g,1729
 relationalai_test_util/__init__.py,sha256=Io_9_IQXXnrUlaL7S1Ndv-4YHilNxy36LrL723MI7lw,118
-relationalai_test_util/fixtures.py,sha256=rNOd8HbguWQi0j63QGoh5iFHLFfM1JUBgDyf87dEsKY,9214
+relationalai_test_util/fixtures.py,sha256=WhpPto6LtZTGWQhYV9Tmz9yJGoVAkYlGRrgM_7errm4,9455
 relationalai_test_util/snapshot.py,sha256=FeH2qYBzLxr2-9qs0yElPIgWUjm_SrzawB3Jgn-aSuE,9291
 relationalai_test_util/traceback.py,sha256=lD0qaEmCyO-7xg9CNf6IzwS-Q-sTS8N9YIv8RroAE50,3298
 frontend/debugger/dist/.gitignore,sha256=JAo-DTfS6GthQGP1NH6wLU-ZymwlTea4KHH_jZVTKn0,14
@@ -451,8 +452,8 @@ frontend/debugger/dist/index.html,sha256=0wIQ1Pm7BclVV1wna6Mj8OmgU73B9rSEGPVX-Wo
 frontend/debugger/dist/assets/favicon-Dy0ZgA6N.png,sha256=tPXOEhOrM4tJyZVJQVBO_yFgNAlgooY38ZsjyrFstgg,620
 frontend/debugger/dist/assets/index-Cssla-O7.js,sha256=MxgIGfdKQyBWgufck1xYggQNhW5nj6BPjCF6Wleo-f0,298886
 frontend/debugger/dist/assets/index-DlHsYx1V.css,sha256=21pZtAjKCcHLFjbjfBQTF6y7QmOic-4FYaKNmwdNZVE,60141
-relationalai-0.13.0.dist-info/METADATA,sha256=RsdnclgO5I9Akv8gxCW-DM5-8ohVCYowAwY-ulQmq6A,2562
-relationalai-0.13.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-relationalai-0.13.0.dist-info/entry_points.txt,sha256=fo_oLFJih3PUgYuHXsk7RnCjBm9cqRNR--ab6DgI6-0,88
-relationalai-0.13.0.dist-info/licenses/LICENSE,sha256=pPyTVXFYhirkEW9VsnHIgUjT0Vg8_xsE6olrF5SIgpc,11343
-relationalai-0.13.0.dist-info/RECORD,,
+relationalai-0.13.1.dist-info/METADATA,sha256=zO4iuv-8tLuz3rByVdd1X3MSdh5MU6XZzqzXcAvdf4k,2561
+relationalai-0.13.1.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+relationalai-0.13.1.dist-info/entry_points.txt,sha256=fo_oLFJih3PUgYuHXsk7RnCjBm9cqRNR--ab6DgI6-0,88
+relationalai-0.13.1.dist-info/licenses/LICENSE,sha256=pPyTVXFYhirkEW9VsnHIgUjT0Vg8_xsE6olrF5SIgpc,11343
+relationalai-0.13.1.dist-info/RECORD,,

relationalai_test_util/fixtures.py

@@ -28,7 +28,7 @@ def graph_index_config_fixture():
     yield config
     return
 
-def engine_config_fixture(size, use_direct_access=False, reset_session=False, generation=rai.Generation.V0):
+def engine_config_fixture(size, use_direct_access=False, reset_session=False, query_timeout_mins=None, generation=rai.Generation.V0):
     # Check for an externally provided engine name
     # It is used in GitHub Actions to run tests against a specific engine
     engine_name = os.getenv("ENGINE_NAME")
@@ -44,6 +44,8 @@ def engine_config_fixture(size, use_direct_access=False, reset_session=False, ge
 
         config.set("reuse_model", False)
         config.set("enable_otel_handler", True) 
+        if query_timeout_mins is not None:
+            config.set("query_timeout_mins", query_timeout_mins)
 
         yield config
         return
@@ -54,6 +56,8 @@ def engine_config_fixture(size, use_direct_access=False, reset_session=False, ge
     if config.file_path is not None:
         # Set test defaults
         config.set("reuse_model", False)
+        if query_timeout_mins is not None:
+            config.set("query_timeout_mins", query_timeout_mins)
         # Try to reset the session instead of using active session if reset_session is true
         rai.Resources(config=config, reset_session=reset_session, generation=generation)
         yield config