relationalai 0.13.0__tar.gz → 0.13.1__tar.gz

{relationalai-0.13.0 → relationalai-0.13.1}/PKG-INFO

@@ -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 → relationalai-0.13.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = 'relationalai'
-version = '0.13.0'
+version = '0.13.1'
 description = 'RelationalAI Library and CLI'
 readme="docs/pypi/README.md"
 authors = [
@@ -8,7 +8,7 @@ authors = [
 ]
 requires-python = ">= 3.10"
 dependencies = [
-    "lqp==0.1.19",
+    "lqp==0.2.1",
     "snowflake-connector-python[secure-local-storage]",
     "snowflake-snowpark-python>=1.38.0",
     "typing-extensions",

relationalai-0.13.1/src/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-0.13.0 → relationalai-0.13.1}/src/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-0.13.0 → relationalai-0.13.1}/src/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-0.13.0 → relationalai-0.13.1}/src/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-0.13.0 → relationalai-0.13.1}/src/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-0.13.0 → relationalai-0.13.1}/src/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