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

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

@@ -13,7 +13,7 @@ from ...config import Config, ConfigStore, ENDPOINT_FILE
 from ...direct_access_client import DirectAccessClient
 from ...types import EngineState
 from ...util import get_pyrel_version, poll_with_specified_overhead, safe_json_loads, ms_to_timestamp
-from ....errors import ResponseStatusException, QueryTimeoutExceededException
+from ....errors import GuardRailsException, ResponseStatusException, QueryTimeoutExceededException
 from snowflake.snowpark import Session
 
 # Import UseIndexResources to enable use_index functionality with direct access
@@ -27,6 +27,7 @@ from typing import Iterable
 
 # Constants
 TXN_ABORT_REASON_TIMEOUT = "transaction timeout"
+TXN_ABORT_REASON_GUARD_RAILS = "guard rail violation"
 
 
 class DirectAccessResources(UseIndexResources):
@@ -355,15 +356,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"

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

@@ -15,6 +15,7 @@ import hashlib
 from dataclasses import dataclass
 
 from ....auth.token_handler import TokenHandler
+from v0.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
@@ -1411,15 +1422,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"
@@ -1704,6 +1718,7 @@ 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:
                 request_headers['user-agent'] = get_pyrel_version(self.generation)
                 request_headers['gi_setup_skipped'] = str(gi_setup_skipped)
@@ -1734,8 +1749,11 @@ Otherwise, remove it from your '{profile}' configuration profile.
                 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)
 
@@ -1744,13 +1762,24 @@ Otherwise, remove it from your '{profile}' configuration profile.
                     filename = row['FILENAME']
                     artifact_info[filename] = row
 
+                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 +2437,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()
 

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

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

v0/relationalai/semantics/lqp/executor.py

@@ -31,7 +31,9 @@ if TYPE_CHECKING:
 
 # Whenever the logic engine introduces a breaking change in behaviour, we bump this version
 # once the client is ready to handle it.
-DEFAULT_LQP_SEMANTICS_VERSION = "0"
+#
+# [2026-01-09] bumping to 1 to opt-into hard validation errors from the engine
+DEFAULT_LQP_SEMANTICS_VERSION = "1"
 
 class LQPExecutor(e.Executor):
     """Executes LQP using the RAI client."""

v0/relationalai/semantics/lqp/rewrite/extract_keys.py

@@ -118,6 +118,17 @@ class ExtractKeys(Pass):
   the same here).
 """
 class ExtractKeysRewriter(Rewriter):
+    def __init__(self):
+        super().__init__()
+        self.compound_keys: dict[Any, ir.Var] = {}
+
+    def _get_compound_key(self, orig_keys: Iterable[ir.Var]) -> ir.Var:
+        if orig_keys in self.compound_keys:
+            return self.compound_keys[orig_keys]
+        compound_key = f.var("compound_key", types.Hash)
+        self.compound_keys[orig_keys] = compound_key
+        return compound_key
+
     def handle_logical(self, node: ir.Logical, parent: ir.Node, ctx:Optional[Any]=None) -> ir.Logical:
         outputs = [x for x in node.body if isinstance(x, ir.Output) and x.keys]
         # We are not in a logical with an output at this level.
@@ -170,7 +181,7 @@ class ExtractKeysRewriter(Rewriter):
         annos = list(output.annotations)
         annos.append(f.annotation(builtins.output_keys, tuple(output_keys)))
         # Create a compound key that will be used in place of the original keys.
-        compound_key = f.var("compound_key", types.Hash)
+        compound_key = self._get_compound_key(output_keys)
 
         for key_combination in combinations:
             missing_keys = OrderedSet.from_iterable(output_keys)
@@ -192,8 +203,13 @@ class ExtractKeysRewriter(Rewriter):
             # handle the construct node in each clone
             values: list[ir.Value] = [compound_key.type]
             for key in output_keys:
-                assert isinstance(key.type, ir.ScalarType)
-                values.append(ir.Literal(types.String, key.type.name))
+                if isinstance(key.type, ir.UnionType):
+                    # the typer can derive union types when multiple distinct entities flow
+                    # into a relation's field, so use AnyEntity as the type marker
+                    values.append(ir.Literal(types.String, "AnyEntity"))
+                else:
+                    assert isinstance(key.type, ir.ScalarType)
+                    values.append(ir.Literal(types.String, key.type.name))
                 if key in key_combination:
                     values.append(key)
             body.add(ir.Construct(None, tuple(values), compound_key, OrderedSet().frozen()))
@@ -408,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)}"