relationalai 0.13.2__py3-none-any.whl → 0.13.4__py3-none-any.whl

relationalai/clients/client.py

@@ -614,7 +614,6 @@ class Client():
                     self._timed_query(
                         "update_registry",
                         dependencies.generate_update_registry(),
-                        readonly=False,
                         abort_on_error=False,
                     )
 
@@ -623,7 +622,6 @@ class Client():
                     self._timed_query(
                         "update_packages",
                         dependencies.generate_update_packages(),
-                        readonly=False,
                         abort_on_error=False,
                     )
             else:
@@ -646,10 +644,11 @@ class Client():
         finally:
             self._database = database_name
 
-    def _timed_query(self, span_name:str, code: str, readonly=True, abort_on_error=True):
+    def _timed_query(self, span_name:str, code: str, abort_on_error=True):
         with debugging.span(span_name, model=self._database) as end_span:
             start = time.perf_counter()
-            res, raw = self._query(code, None, end_span, readonly=readonly, abort_on_error=abort_on_error)
+            # NOTE hardcoding to readonly=False, read-only Rel transactions are deprecated.
+            res, raw = self._query(code, None, end_span, readonly=False, abort_on_error=abort_on_error)
             debugging.time(span_name, time.perf_counter() - start, code=code)
             return res, raw
 

relationalai/clients/exec_txn_poller.py

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

relationalai/clients/resources/snowflake/direct_access_resources.py

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

relationalai/clients/resources/snowflake/snowflake.py

@@ -15,7 +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
+from relationalai.clients.exec_txn_poller import ExecTxnPoller
 import snowflake.snowpark
 
 from ....rel_utils import sanitize_identifier, to_fqn_relation_name
@@ -105,6 +105,9 @@ TXN_ABORT_REASON_TIMEOUT = "transaction timeout"
 GUARDRAILS_ABORT_REASON = "guard rail violation"
 
 PRINT_TXN_PROGRESS_FLAG = "print_txn_progress"
+ENABLE_GUARD_RAILS_FLAG = "enable_guard_rails"
+
+ENABLE_GUARD_RAILS_HEADER = "X-RAI-Enable-Guard-Rails"
 
 #--------------------------------------------------
 # Helpers
@@ -113,6 +116,9 @@ PRINT_TXN_PROGRESS_FLAG = "print_txn_progress"
 def should_print_txn_progress(config) -> bool:
     return bool(config.get(PRINT_TXN_PROGRESS_FLAG, False))
 
+def should_enable_guard_rails(config) -> bool:
+    return bool(config.get(ENABLE_GUARD_RAILS_FLAG, False))
+
 #--------------------------------------------------
 # Resources
 #--------------------------------------------------
@@ -1788,65 +1794,62 @@ Otherwise, remove it from your '{profile}' configuration profile.
 
         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 ""
-
-                # 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,
-                    bypass_index=bypass_index,
-                    language=language,
-                    query_timeout_mins=query_timeout_mins,
-                )
+            print_txn_progress = should_print_txn_progress(self.config)
 
-                txn_id = result.txn_id
-                state = result.state
+            with ExecTxnPoller(
+                print_txn_progress=print_txn_progress,
+                resource=self, txn_id=None, headers=request_headers,
+                txn_start_time=txn_start_time
+            ) as poller:
+                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 ""
+                    request_headers[ENABLE_GUARD_RAILS_HEADER] = str(should_enable_guard_rails(self.config))
+
+                    # 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,
+                        bypass_index=bypass_index,
+                        language=language,
+                        query_timeout_mins=query_timeout_mins,
+                    )
 
-                txn_span["txn_id"] = txn_id
-                create_span["txn_id"] = txn_id
-                debugging.event("transaction_created", txn_span, txn_id=txn_id)
+                    txn_id = result.txn_id
+                    state = result.state
 
-            print_txn_progress = should_print_txn_progress(self.config)
+                    txn_span["txn_id"] = txn_id
+                    create_span["txn_id"] = txn_id
+                    debugging.event("transaction_created", txn_span, txn_id=txn_id)
 
-            # 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)
+                # Set the transaction ID now that we have it, to update the progress text
+                poller.txn_id = txn_id
 
-                artifact_info = result.artifact_info
+                # fast path: transaction already finished
+                if state in ["COMPLETED", "ABORTED"]:
+                    if txn_id in self._pending_transactions:
+                        self._pending_transactions.remove(txn_id)
 
-                txn_duration = txn_end_time - txn_start_time
-                if print_txn_progress:
-                    print(
-                        query_complete_message(txn_id, txn_duration, status_header=True)
-                    )
+                    artifact_info = result.artifact_info
 
-            # 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):
-                    if print_txn_progress:
-                        poller = ExecTxnPoller(resource=self, txn_id=txn_id, headers=request_headers, txn_start_time=txn_start_time)
+                # 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):
                         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"):
-                return self._download_results(artifact_info, txn_id, state)
+                    artifact_info = self._list_exec_async_artifacts(txn_id, headers=request_headers)
+
+                with debugging.span("fetch"):
+                    return self._download_results(artifact_info, txn_id, state)
 
     def get_user_based_engine_name(self):
         if not self._session:

relationalai/clients/resources/snowflake/use_index_poller.py

@@ -66,7 +66,7 @@ ERP_CHECK_FREQUENCY = 15
 
 # Polling behavior constants
 POLL_OVERHEAD_RATE = 0.1  # Overhead rate for exponential backoff
-POLL_MAX_DELAY = 2.5  # Maximum delay between polls in seconds
+POLL_MAX_DELAY = 0  # Maximum delay between polls in seconds
 
 # SQL query template for getting stream column hashes
 # This query calculates a hash of column metadata (name, type, precision, scale, nullable)

relationalai/semantics/internal/snowflake.py

@@ -214,6 +214,7 @@ class Table():
         self._col_names = cols
         self._iceberg_config = config
         self._is_iceberg = config is not None
+        self._skip_cdc = False
         info = self._schemas.get((self._database, self._schema))
         if not info:
             info = self._schemas[(self._database, self._schema)] = SchemaInfo(self._database, self._schema)
@@ -303,7 +304,10 @@ class Table():
 
     def _compile_lookup(self, compiler:b.Compiler, ctx:b.CompilerContext):
         self._lazy_init()
-        Table._used_sources.add(self)
+        if not self._skip_cdc:
+            # Don't do CDC if the underlying data has been loaded
+            # directly via `api.load_data`.
+            Table._used_sources.add(self)
         compiler.lookup(self._rel, ctx)
         return compiler.lookup(b.RelationshipFieldRef(None, self._rel, 0), ctx)
 

relationalai/semantics/lqp/algorithms.py

@@ -0,0 +1,173 @@
+from typing import TypeGuard
+from relationalai.semantics.metamodel import ir, factory, types
+from relationalai.semantics.metamodel.visitor import Rewriter, collect_by_type
+from relationalai.semantics.lqp import ir as lqp
+from relationalai.semantics.lqp.types import meta_type_to_lqp
+from relationalai.semantics.lqp.builtins import (
+    has_empty_annotation, has_assign_annotation, has_upsert_annotation,
+    has_monoid_annotation, has_monus_annotation, has_script_annotation,
+    has_algorithm_annotation, has_while_annotation, global_annotation,
+    empty_annotation, assign_annotation, upsert_annotation, monoid_annotation,
+    monus_annotation
+)
+
+# Complex tests for Loopy constructs in the metamodel
+def is_script(task: ir.Task) -> TypeGuard[ir.Sequence]:
+    """ Check if it is a script i.e., a Sequence with @script annotation. """
+    if not isinstance(task, ir.Sequence):
+        return False
+    return has_script_annotation(task)
+
+def is_algorithm_logical(task: ir.Task) -> TypeGuard[ir.Logical]:
+    """ Check if it is an algorithm logical i.e., a Logical task with all subtasks being
+    algorithm scripts. """
+    if not isinstance(task, ir.Logical):
+        return False
+    return all(is_algorithm_script(subtask) for subtask in task.body)
+
+def is_algorithm_script(task: ir.Task) -> TypeGuard[ir.Sequence]:
+    """ Check if it is an algorithm script i.e., a Sequence with @script and @algorithm annotations. """
+    if not isinstance(task, ir.Sequence):
+        return False
+    return is_script(task) and has_algorithm_annotation(task)
+
+def is_while_loop(task: ir.Task) -> TypeGuard[ir.Loop]:
+    """ Check if input is is a while loop i.e., a Loop with @while annotation. """
+    if not isinstance(task, ir.Loop):
+        return False
+    return has_while_annotation(task)
+
+def is_while_script(task: ir.Task) -> TypeGuard[ir.Sequence]:
+    """ Check if input is a while script i.e., a Sequence with @script and @while annotations. """
+    if not isinstance(task, ir.Sequence):
+        return False
+    return is_script(task) and has_while_annotation(task)
+
+# Tools for annotating Loopy constructs
+class LoopyAnnoAdder(Rewriter):
+    """ Rewrites a node by adding the given annotation to all Update nodes. """
+    def __init__(self, anno: ir.Annotation):
+        self.anno = anno
+        super().__init__()
+
+    def handle_update(self, node: ir.Update, parent: ir.Node) -> ir.Update:
+        new_annos = list(node.annotations) + [self.anno]
+        return factory.update(node.relation, node.args, node.effect, new_annos, node.engine)
+
+def mk_global(i: ir.Node):
+    return LoopyAnnoAdder(global_annotation()).walk(i)
+
+def mk_empty(i: ir.Node):
+    return LoopyAnnoAdder(empty_annotation()).walk(i)
+
+def mk_assign(i: ir.Node):
+    return LoopyAnnoAdder(assign_annotation()).walk(i)
+
+def mk_upsert(i: ir.Node, arity: int):
+    return LoopyAnnoAdder(upsert_annotation(arity)).walk(i)
+
+def mk_monoid(i: ir.Node, monoid_type: ir.ScalarType, monoid_op: str, arity: int):
+    return LoopyAnnoAdder(monoid_annotation(monoid_type, monoid_op, arity)).walk(i)
+
+def mk_monus(i: ir.Node, monoid_type: ir.ScalarType, monoid_op: str, arity: int):
+    return LoopyAnnoAdder(monus_annotation(monoid_type, monoid_op, arity)).walk(i)
+
+def construct_monoid(i: ir.Annotation):
+    base_type = None
+    op = None
+    for arg in i.args:
+        if isinstance(arg, ir.ScalarType):
+            base_type = meta_type_to_lqp(arg)
+        elif isinstance(arg, ir.Literal) and arg.type == types.String:
+            op = arg.value
+    assert isinstance(base_type, lqp.Type) and isinstance(op, str), "Failed to get monoid"
+    if op.lower() == "or":
+        return lqp.OrMonoid(meta=None)
+    elif op.lower() == "sum":
+        return lqp.SumMonoid(type=base_type, meta=None)
+    elif op.lower() == "min":
+        return lqp.MinMonoid(type=base_type, meta=None)
+    elif op.lower() == "max":
+        return lqp.MaxMonoid(type=base_type, meta=None)
+    else:
+        assert False, "Failed to get monoid"
+
+# Tools for analyzing Loopy constructs
+def is_logical_instruction(node: ir.Node) -> TypeGuard[ir.Logical]:
+    if not isinstance(node, ir.Logical):
+        return False
+    return any(collect_by_type(ir.Update, node)) and not any(collect_by_type(ir.Sequence, node))
+
+def get_instruction_body_rels(node: ir.Logical) -> set[ir.Relation]:
+    assert is_logical_instruction(node)
+    body: set[ir.Relation] = set()
+    for update in collect_by_type(ir.Lookup, node):
+        body.add(update.relation)
+    return body
+
+def get_instruction_head_rels(node: ir.Logical) -> set[ir.Relation]:
+    assert is_logical_instruction(node)
+    heads: set[ir.Relation] = set()
+    for update in collect_by_type(ir.Update, node):
+        heads.add(update.relation)
+    return heads
+
+# base Loopy instruction: @empty, @assign, @upsert, @monoid, @monus
+def is_instruction(update: ir.Task) -> TypeGuard[ir.Logical]:
+    if not is_logical_instruction(update):
+        return False
+    for u in collect_by_type(ir.Update, update):
+        if (has_empty_annotation(u) or
+            has_assign_annotation(u) or
+            has_upsert_annotation(u) or
+            has_monoid_annotation(u) or
+            has_monus_annotation(u)):
+            return True
+    return False
+
+# update Loopy instruction @upsert, @monoid, @monus
+def is_update_instruction(task: ir.Task) -> TypeGuard[ir.Logical]:
+    if not is_logical_instruction(task):
+        return False
+    for u in collect_by_type(ir.Update, task):
+        if (has_upsert_annotation(u) or
+            has_monoid_annotation(u) or
+            has_monus_annotation(u)):
+            return True
+    return False
+
+def is_empty_instruction(node: ir.Node) -> TypeGuard[ir.Logical]:
+    """ Check if input is an empty Loopy instruction `empty rel = ∅`"""
+    if not is_logical_instruction(node):
+        return False
+    updates = collect_by_type(ir.Update, node)
+    if not any(has_empty_annotation(update) for update in updates):
+        return False
+
+    # At this point, we have the prerequisites for an empty instruction. We check it is
+    # well-formed:
+    # 1. It has only a single @empty Update operation
+    # 2. Has no other operations
+    assert len(updates) == 1, "[Loopy] Empty instruction must have single Update operation"
+    assert len(node.body) == 1, "[Loopy] Empty instruction must have only a single Update operation"
+
+    return True
+
+# Splits a Loopy instruction into its head updates, body lookups, and other body tasks
+def split_instruction(update_logical: ir.Logical) -> tuple[ir.Update,list[ir.Lookup],list[ir.Task]]:
+    assert is_instruction(update_logical)
+    lookups = []
+    update = None
+    others = []
+    for task in update_logical.body:
+        if isinstance(task, ir.Lookup):
+            lookups.append(task)
+        elif isinstance(task, ir.Update):
+            if update is not None:
+                raise AssertionError("[Loopy] Update instruction must have exactly one Update operation")
+            update = task
+        else:
+            others.append(task)
+    assert update is not None, "[Loopy] Update instruction must have exactly one Update operation"
+
+    return update, lookups, others