relationalai 0.12.7__py3-none-any.whl → 0.12.9__py3-none-any.whl

relationalai/clients/snowflake.py

@@ -851,7 +851,17 @@ Otherwise, remove it from your '{profile}' configuration profile.
                 self.generation
             )
             # If cache is valid (data freshness has not expired), skip polling
-            if not poller.cache.is_valid():
+            if poller.cache.is_valid():
+                cached_sources = len(poller.cache.sources)
+                total_sources = len(sources_list)
+                cached_timestamp = poller.cache._metadata.get("cachedIndices", {}).get(poller.cache.key, {}).get("last_use_index_update_on", "")
+
+                message = f"Using cached data for {cached_sources}/{total_sources} data streams"
+                if cached_timestamp:
+                    print(f"\n{message} (cached at {cached_timestamp})\n")
+                else:
+                    print(f"\n{message}\n")
+            else:
                 return poller.poll()
 
     #--------------------------------------------------
@@ -3284,12 +3294,24 @@ class DirectAccessResources(Resources):
             try:
                 response = _send_request()
                 if response.status_code != 200:
+                    # For 404 responses with skip_auto_create=True, return immediately to let caller handle it
+                    # (e.g., get_engine needs to check 404 and return None for auto_create_engine)
+                    # For skip_auto_create=False, continue to auto-creation logic below
+                    if response.status_code == 404 and skip_auto_create:
+                        return response
+
                     try:
                         message = response.json().get("message", "")
                     except requests.exceptions.JSONDecodeError:
-                        raise ResponseStatusException(
-                            f"Failed to parse error response from endpoint {endpoint}.", response
-                        )
+                        # Can't parse JSON response. For skip_auto_create=True (e.g., get_engine),
+                        # this should have been caught by the 404 check above, so this is an error.
+                        # For skip_auto_create=False, we explicitly check status_code below,
+                        # so we don't need to parse the message.
+                        if skip_auto_create:
+                            raise ResponseStatusException(
+                                f"Failed to parse error response from endpoint {endpoint}.", response
+                            )
+                        message = ""  # Not used when we check status_code directly
 
                     # fix engine on engine error and retry
                     # Skip auto-retry if skip_auto_create is True to avoid recursion
@@ -3482,7 +3504,17 @@ class DirectAccessResources(Resources):
                 generation=self.generation,
             )
             # If cache is valid (data freshness has not expired), skip polling
-            if not poller.cache.is_valid():
+            if poller.cache.is_valid():
+                cached_sources = len(poller.cache.sources)
+                total_sources = len(sources_list)
+                cached_timestamp = poller.cache._metadata.get("cachedIndices", {}).get(poller.cache.key, {}).get("last_use_index_update_on", "")
+
+                message = f"Using cached data for {cached_sources}/{total_sources} data streams"
+                if cached_timestamp:
+                    print(f"\n{message} (cached at {cached_timestamp})\n")
+                else:
+                    print(f"\n{message}\n")
+            else:
                 return poller.poll()
 
     def _check_exec_async_status(self, txn_id: str, headers: Dict[str, str] | None = None) -> bool:

relationalai/clients/use_index_poller.py

@@ -250,7 +250,17 @@ class UseIndexPoller:
             # Cache was used - show how many sources were cached
             total_sources = len(self.cache.sources)
             cached_sources = total_sources - len(self.sources)
-            progress.add_sub_task(f"Using cached data for {cached_sources}/{total_sources} data streams", task_id="cache_usage", category=TASK_CATEGORY_CACHE)
+
+            # Get the timestamp when sources were cached
+            entry = self.cache._metadata.get("cachedIndices", {}).get(self.cache.key, {})
+            cached_timestamp = entry.get("last_use_index_update_on", "")
+
+            message = f"Using cached data for {cached_sources}/{total_sources} data streams"
+            # Format the message with timestamp
+            if cached_timestamp:
+                message += f" (cached at {cached_timestamp})"
+
+            progress.add_sub_task(message, task_id="cache_usage", category=TASK_CATEGORY_CACHE)
             # Complete the subtask immediately since it's just informational
             progress.complete_sub_task("cache_usage")
 

relationalai/semantics/internal/internal.py

@@ -40,7 +40,7 @@ _global_id = peekable(itertools.count(0))
 
 # Single context variable with default values
 _overrides = ContextVar("overrides", default = {})
-def overrides(key: str, default: bool | str | dict):
+def overrides(key: str, default: bool | str | dict | datetime | None):
     return _overrides.get().get(key, default)
 
 # Flag that users set in the config or directly on the model, but that can still be
@@ -60,6 +60,13 @@ def with_overrides(**kwargs):
     finally:
         _overrides.reset(token)
 
+# Intrinsic values to override for stable snapshots.
+def get_intrinsic_overrides() -> dict[str, Any]:
+    datetime_now = overrides('datetime_now', None)
+    if datetime_now is not None:
+        return {'datetime_now': datetime_now}
+    return {}
+
 #--------------------------------------------------
 # Root tracking
 #--------------------------------------------------
@@ -953,12 +960,25 @@ class Concept(Producer):
                 self._validate_identifier_relationship(rel)
         self._add_ref_scheme(*args)
 
-    def _add_ref_scheme(self, *args: Relationship|RelationshipReading):
-        self._reference_schemes.append(args)
-        # assumed that all Relationship|RelationshipReading are defined on the identified Concept
-        fields = tuple([arg.__getitem__(0) for arg in args])
-        uc = Unique(*fields, model=self._model)
-        require(uc.to_expressions())
+    def _add_ref_scheme(self, *rels: Relationship|RelationshipReading):
+        # thanks to prior validation we we can safely assume that
+        # * the input types are correct due to prior validation
+        # * all relationships are binary and defined on this concept
+
+        self._reference_schemes.append(rels)
+
+        # for every concept x every field f has at most one value y.
+        # f(x,y): x -> y holds
+        concept_fields = tuple([rel.__getitem__(0) for rel in rels])
+        for field in concept_fields:
+            concept_uc = Unique(field, model=self._model)
+            require(concept_uc.to_expressions())
+
+        # for any combination of field values there is at most one concept x.
+        # f₁(x,y₁) ∧ … ∧ fₙ(x,yₙ): {y₁,…,yₙ} → {x}
+        key_fields = tuple([rel.__getitem__(1) for rel in rels])
+        key_uc = Unique(*key_fields, model=self._model)
+        require(key_uc.to_expressions())
 
     def _validate_identifier_relationship(self, rel:Relationship|RelationshipReading):
         if rel._arity() != 2:
@@ -2603,6 +2623,7 @@ class Model():
         config_overrides = overrides('config', {})
         for k, v in config_overrides.items():
             self._config.set(k, v)
+        self._intrinsic_overrides = get_intrinsic_overrides()
         self._strict = cast(bool, overrides('strict', strict))
         self._use_lqp = overridable_flag('reasoner.rule.use_lqp', self._config, use_lqp, default=not self._use_sql)
         self._enable_otel_handler = overridable_flag('enable_otel_handler', self._config, enable_otel_handler, default=False)
@@ -2644,6 +2665,7 @@ class Model():
                     wide_outputs=self._wide_outputs,
                     connection=self._connection,
                     config=self._config,
+                    intrinsic_overrides=self._intrinsic_overrides,
                 )
             elif self._use_sql:
                 self._executor = SnowflakeExecutor(

relationalai/semantics/lqp/compiler.py

@@ -14,7 +14,7 @@ class Compiler(c.Compiler):
         super().__init__(lqp_passes())
         self.def_names = UniqueNames()
 
-    def do_compile(self, model: ir.Model, options:dict={}) -> tuple[Optional[tuple], lqp.Transaction]:
+    def do_compile(self, model: ir.Model, options:dict={}) -> tuple[Optional[tuple], lqp.Epoch]:
         fragment_id: bytes = options.get("fragment_id", bytes(404))
         # Reset the var context for each compilation
         # TODO: Change to unique var names per lookup

relationalai/semantics/lqp/constructors.py

@@ -59,3 +59,9 @@ def mk_pragma(name: str, terms: list[lqp.Var]) -> lqp.Pragma:
 
 def mk_attribute(name: str, args: list[lqp.Value]) -> lqp.Attribute:
     return lqp.Attribute(name=name, args=args, meta=None)
+
+def mk_transaction(
+    epochs: list[lqp.Epoch],
+    configure: lqp.Configure = lqp.construct_configure({}, None),
+) -> lqp.Transaction:
+    return lqp.Transaction(epochs=epochs, configure=configure, meta=None)

relationalai/semantics/lqp/executor.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 from collections import defaultdict
+from datetime import datetime, timezone
 import atexit
 import re
 
@@ -13,6 +14,7 @@ from relationalai.semantics.lqp import result_helpers
 from relationalai.semantics.metamodel import ir, factory as f, executor as e
 from relationalai.semantics.lqp.compiler import Compiler
 from relationalai.semantics.lqp.intrinsics import mk_intrinsic_datetime_now
+from relationalai.semantics.lqp.constructors import mk_transaction
 from relationalai.semantics.lqp.types import lqp_type_to_sql
 from lqp import print as lqp_print, ir as lqp_ir
 from lqp.parser import construct_configure
@@ -39,6 +41,9 @@ class LQPExecutor(e.Executor):
         wide_outputs: bool = False,
         connection: Session | None = None,
         config: Config | None = None,
+        # In order to facilitate snapshot testing, we allow overriding intrinsic definitions
+        # like the current time, which would otherwise change between runs.
+        intrinsic_overrides: dict = {},
     ) -> None:
         super().__init__()
         self.database = database
@@ -48,6 +53,7 @@ class LQPExecutor(e.Executor):
         self.compiler = Compiler()
         self.connection = connection
         self.config = config or Config()
+        self.intrinsic_overrides = intrinsic_overrides
         self._resources = None
         self._last_model = None
         self._last_sources_version = (-1, None)
@@ -311,19 +317,16 @@ class LQPExecutor(e.Executor):
         with debugging.span("compile_intrinsics") as span:
             span["compile_type"] = "intrinsics"
 
+            now = self.intrinsic_overrides.get('datetime_now', datetime.now(timezone.utc))
+
             debug_info = lqp_ir.DebugInfo(id_to_orig_name={}, meta=None)
             intrinsics_fragment = lqp_ir.Fragment(
                 id = lqp_ir.FragmentId(id=b"__pyrel_lqp_intrinsics", meta=None),
-                declarations = [
-                    mk_intrinsic_datetime_now(),
-                ],
+                declarations = [mk_intrinsic_datetime_now(now)],
                 debug_info = debug_info,
                 meta = None,
             )
 
-
-            span["lqp"] = lqp_print.to_string(intrinsics_fragment, {"print_names": True, "print_debug": False, "print_csv_filename": False})
-
             return lqp_ir.Epoch(
                 writes=[
                     lqp_ir.Write(write_type=lqp_ir.Define(fragment=intrinsics_fragment, meta=None), meta=None)
@@ -354,47 +357,38 @@ class LQPExecutor(e.Executor):
 
     def compile_lqp(self, model: ir.Model, task: ir.Task):
         configure = self._construct_configure()
+        # Merge the epochs into a single transaction. Long term the query bits should all
+        # go into a WhatIf action and the intrinsics could be fused with either of them. But
+        # for now we just use separate epochs.
+        epochs = []
+        epochs.append(self._compile_intrinsics())
 
-        model_txn = None
         if self._last_model != model:
             with debugging.span("compile", metamodel=model) as install_span:
                 install_span["compile_type"] = "model"
-                _, model_txn = self.compiler.compile(model, {"fragment_id": b"model"})
-                model_txn = txn_with_configure(model_txn, configure)
-                install_span["lqp"] = lqp_print.to_string(model_txn, {"print_names": True, "print_debug": False, "print_csv_filename": False})
+                _, model_epoch = self.compiler.compile(model, {"fragment_id": b"model"})
+                epochs.append(model_epoch)
                 self._last_model = model
 
-        with debugging.span("compile", metamodel=task) as compile_span:
-            compile_span["compile_type"] = "query"
+        with debugging.span("compile", metamodel=task) as txn_span:
             query = f.compute_model(f.logical([task]))
             options = {
                 "wide_outputs": self.wide_outputs,
                 "fragment_id": b"query",
             }
             result, final_model = self.compiler.compile_inner(query, options)
-            export_info, query_txn = result
-            query_txn = txn_with_configure(query_txn, configure)
-            compile_span["lqp"] = lqp_print.to_string(query_txn, {"print_names": True, "print_debug": False, "print_csv_filename": False})
+            export_info, query_epoch = result
 
-        # Merge the epochs into a single transactions. Long term the query bits should all
-        # go into a WhatIf action and the intrinsics could be fused with either of them. But
-        # for now we just use separate epochs.
-        epochs = []
+            epochs.append(query_epoch)
+            epochs.append(self._compile_undefine_query(query_epoch))
 
-        epochs.append(self._compile_intrinsics())
+            txn_span["compile_type"] = "query"
+            txn = mk_transaction(epochs=epochs, configure=configure)
+            txn_span["lqp"] = lqp_print.to_string(txn, {"print_names": True, "print_debug": False, "print_csv_filename": False})
 
-        if model_txn is not None:
-            epochs.append(model_txn.epochs[0])
-
-        query_txn_epoch = query_txn.epochs[0]
-        epochs.append(query_txn_epoch)
-        epochs.append(self._compile_undefine_query(query_txn_epoch))
-
-        txn = lqp_ir.Transaction(epochs=epochs, configure=configure, meta=None)
         validate_lqp(txn)
 
         txn_proto = convert_transaction(txn)
-        # TODO (azreika): Should export_info be encoded as part of the txn_proto? [RAI-40312]
         return final_model, export_info, txn_proto
 
     # TODO (azreika): This should probably be split up into exporting and other processing. There are quite a lot of arguments here...
@@ -462,12 +456,3 @@ class LQPExecutor(e.Executor):
             # If processing the results failed, revert to the previous model.
             self._last_model = previous_model
             raise e
-
-def txn_with_configure(txn: lqp_ir.Transaction, configure: lqp_ir.Configure) -> lqp_ir.Transaction:
-    """ Return a new transaction with the given configure. If the transaction already has
-    a configure, it is replaced. """
-    return lqp_ir.Transaction(
-        epochs=txn.epochs,
-        configure=configure,
-        meta=txn.meta,
-    )

relationalai/semantics/lqp/intrinsics.py

@@ -1,15 +1,16 @@
-from datetime import datetime, timezone
+from datetime import datetime
 
 from relationalai.semantics.lqp import ir as lqp
 from relationalai.semantics.lqp.constructors import mk_abstraction, mk_value, mk_type, mk_primitive
 from relationalai.semantics.lqp.utils import lqp_hash
 
-def mk_intrinsic_datetime_now() -> lqp.Def:
+# Constructs a definition of the current datetime.
+def mk_intrinsic_datetime_now(dt: datetime) -> lqp.Def:
     """Constructs a definition of the current datetime."""
     id = lqp_hash("__pyrel_lqp_intrinsic_datetime_now")
     out = lqp.Var(name="out", meta=None)
     out_type = mk_type(lqp.TypeName.DATETIME)
-    now = mk_value(lqp.DateTimeValue(value=datetime.now(timezone.utc), meta=None))
+    now = mk_value(lqp.DateTimeValue(value=dt, meta=None))
     datetime_now = mk_abstraction(
         [(out, out_type)],
         mk_primitive("rel_primitive_eq", [out, now]),

relationalai/semantics/lqp/model2lqp.py

@@ -19,8 +19,8 @@ from warnings import warn
 import re
 import uuid
 
-""" Main access point. Converts the model IR to an LQP transaction. """
-def to_lqp(model: ir.Model, fragment_name: bytes, ctx: TranslationCtx) -> tuple[Optional[tuple], lqp.Transaction]:
+# Main access point for translating metamodel to lqp. Converts the model IR to an LQP epoch.
+def to_lqp(model: ir.Model, fragment_name: bytes, ctx: TranslationCtx) -> tuple[Optional[tuple], lqp.Epoch]:
     assert_valid_input(model)
     decls: list[lqp.Declaration] = []
     reads: list[lqp.Read] = []
@@ -50,16 +50,10 @@ def to_lqp(model: ir.Model, fragment_name: bytes, ctx: TranslationCtx) -> tuple[
     fragment = lqp.Fragment(id=fragment_id, declarations=decls, meta=None, debug_info=debug_info)
     define_op = lqp.Define(fragment=fragment, meta=None)
 
-    txn = lqp.Transaction(
-        epochs=[
-            lqp.Epoch(
-                reads=reads,
-                writes=[lqp.Write(write_type=define_op, meta=None)],
-                meta=None
-            )
-        ],
-        configure=lqp.construct_configure({}, None),
-        meta=None,
+    txn = lqp.Epoch(
+        reads=reads,
+        writes=[lqp.Write(write_type=define_op, meta=None)],
+        meta=None
     )
 
     return (export_info, txn)

relationalai/semantics/lqp/passes.py

@@ -6,8 +6,8 @@ from relationalai.semantics.metamodel.util import FrozenOrderedSet
 
 from relationalai.semantics.metamodel.rewrite import Flatten
 
-from ..metamodel.rewrite import DischargeConstraints, DNFUnionSplitter, ExtractNestedLogicals
-from .rewrite import CDC, ExtractCommon, ExtractKeys, FunctionAnnotations, QuantifyVars, Splinter
+from ..metamodel.rewrite import DischargeConstraints, DNFUnionSplitter, ExtractNestedLogicals, FormatOutputs
+from .rewrite import CDC, ExtractCommon, ExtractKeys, FunctionAnnotations, QuantifyVars, Splinter, SplitMultiCheckRequires
 
 from relationalai.semantics.lqp.utils import output_names
 
@@ -18,6 +18,7 @@ import hashlib
 
 def lqp_passes() -> list[Pass]:
     return [
+        SplitMultiCheckRequires(),
         FunctionAnnotations(),
         DischargeConstraints(),
         Checker(),
@@ -27,6 +28,7 @@ def lqp_passes() -> list[Pass]:
         DNFUnionSplitter(),
         ExtractKeys(),
         ExtractCommon(),
+        FormatOutputs(),
         Flatten(),
         Splinter(), # Splits multi-headed rules into multiple rules
         QuantifyVars(), # Adds missing existentials

relationalai/semantics/lqp/rewrite/__init__.py

@@ -1,7 +1,7 @@
 from .cdc import CDC
 from .extract_common import ExtractCommon
 from .extract_keys import ExtractKeys
-from .function_annotations import FunctionAnnotations
+from .function_annotations import FunctionAnnotations, SplitMultiCheckRequires
 from .quantify_vars import QuantifyVars
 from .splinter import Splinter
 
@@ -12,4 +12,5 @@ __all__ = [
     "FunctionAnnotations",
     "QuantifyVars",
     "Splinter",
+    "SplitMultiCheckRequires",
 ]

relationalai/semantics/lqp/rewrite/function_annotations.py

@@ -1,79 +1,114 @@
 from __future__ import annotations
 
-from dataclasses import dataclass
-from relationalai.semantics.metamodel import ir, compiler as c, visitor as v, builtins
-from relationalai.semantics.metamodel.util import OrderedSet, ordered_set
+from typing import Optional
+from relationalai.semantics.metamodel import builtins
+from relationalai.semantics.metamodel.ir import (
+    Node, Model, Require, Logical, Relation, Annotation, Update
+)
+from relationalai.semantics.metamodel.compiler import Pass
+from relationalai.semantics.metamodel.visitor import Rewriter, Visitor
+from relationalai.semantics.lqp.rewrite.functional_dependencies import (
+    is_valid_unique_constraint, normalized_fd
+)
 
+# In the future iterations of PyRel metamodel, `Require` nodes will have a single `Check`
+# (and no `errors`). Currently, however, the unique constraints may result in multiple
+# `Check` nodes and for simplicity we split them in to separate `Require` nodes. This step
+# will be removed in the future.
+#
+# Note that unique constraints always have an empty `domain` so apply the splitting only
+# to such `Require` nodes.
+class SplitMultiCheckRequires(Pass):
+    """
+    Pass splits unique Require nodes that have empty domain but multiple checks into multiple
+    Require nodes with single check each.
+    """
+
+    def rewrite(self, model: Model, options: dict = {}) -> Model:
+        return SplitMultiCheckRequiresRewriter().walk(model)
+
+
+class SplitMultiCheckRequiresRewriter(Rewriter):
+    """
+    Splits unique Require nodes that have empty domain but multiple checks into multiple
+    Require nodes with single check each.
+    """
+    def handle_require(self, node: Require, parent: Node):
 
-class FunctionAnnotations(c.Pass):
+        if isinstance(node.domain, Logical) and not node.domain.body and len(node.checks) > 1:
+            require_nodes = []
+            for check in node.checks:
+                single_check = self.walk(check, node)
+                require_nodes.append(
+                    node.reconstruct(node.engine, node.domain, (single_check,), node.annotations)
+                )
+            return require_nodes
+
+        return node
+
+
+class FunctionAnnotations(Pass):
     """
-    Pass marks all appropriate relations with `function` annotation.
-    Criteria:
-    - there is a Require node with `unique` builtin (appeared as a result of `require(unique(...))`)
-    - `unique` declared for all the fields in a derived relation expect the last
+    Pass marks all appropriate relations with `function` annotation. Collects functional
+    dependencies from unique Require nodes and uses this information to identify functional
+    relations.
     """
 
-    def rewrite(self, model: ir.Model, options: dict = {}) -> ir.Model:
-        collect_fd = CollectFunctionalRelationsVisitor()
-        new_model = collect_fd.walk(model)
-        # mark relations collected by previous visitor with `@function` annotation
-        return FunctionalAnnotationsVisitor(collect_fd.functional_relations).walk(new_model)
+    def rewrite(self, model: Model, options: dict = {}) -> Model:
+        collect_fds = CollectFDsVisitor()
+        collect_fds.visit_model(model, None)
+        annotated_model = FunctionalAnnotationsRewriter(collect_fds.functional_relations).walk(model)
+        return annotated_model
 
 
-@dataclass
-class CollectFunctionalRelationsVisitor(v.Rewriter):
+class CollectFDsVisitor(Visitor):
     """
-    Visitor collects all relations which should be marked with `functional` annotation.
+    Visitor collects all unique constraints.
     """
 
+    # Currently, only information about k-functional fd is collected.
     def __init__(self):
         super().__init__()
-        self.functional_relations = ordered_set()
-
-    def handle_check(self, node: ir.Check, parent: ir.Node):
-        check = self.walk(node.check, node)
-        assert isinstance(check, ir.Logical)
-        unique_vars = []
-        for item in check.body:
-            # collect vars from `unique` builtin
-            if isinstance(item, ir.Lookup) and item.relation.name == builtins.unique.name:
-                var_set = set()
-                for vargs in item.args:
-                    assert isinstance(vargs, tuple)
-                    var_set.update(vargs)
-                unique_vars.append(var_set)
-        functional_rel = []
-        # mark relations as functional when at least 1 `unique` builtin
-        if len(unique_vars) > 0:
-            for item in check.body:
-                if isinstance(item, ir.Lookup) and not item.relation.name == builtins.unique.name:
-                    for var_set in unique_vars:
-                        # when unique declared for all the fields except the last one in the relation mark it as functional
-                        if var_set == set(item.args[:-1]):
-                            functional_rel.append(item.relation)
-
-        self.functional_relations.update(functional_rel)
-        return node.reconstruct(check, node.error, node.annotations)
-
-
-@dataclass
-class FunctionalAnnotationsVisitor(v.Rewriter):
+        self.functional_relations:dict[Relation, int] = {}
+
+    def visit_require(self, node: Require, parent: Optional[Node]):
+        if is_valid_unique_constraint(node):
+            fd = normalized_fd(node)
+            assert fd is not None
+            if fd.is_structural:
+                relation = fd.structural_relation
+                k = fd.structural_rank
+                current_k = self.functional_relations.get(relation, 0)
+                self.functional_relations[relation] = max(current_k, k)
+
+
+class FunctionalAnnotationsRewriter(Rewriter):
     """
-    This visitor marks functional_relations with `function` annotation.
+    This visitor marks functional_relations with `@function(:checked [, k])` annotation.
     """
 
-    def __init__(self, functional_relations: OrderedSet):
+    def __init__(self, functional_relations: dict[Relation, int]):
         super().__init__()
-        self._functional_relations = functional_relations
+        self.functional_relations = functional_relations
+
+    def get_functional_annotation(self, rel: Relation) -> Optional[Annotation]:
+        k = self.functional_relations.get(rel, None)
+        if k is None:
+            return None
+        if k == 1:
+            return builtins.function_checked_annotation
+        return builtins.function_ranked_checked_annotation(k)
 
-    def handle_relation(self, node: ir.Relation, parent: ir.Node):
-        if node in self._functional_relations:
-            return node.reconstruct(node.name, node.fields, node.requires, node.annotations | [builtins.function_checked_annotation],
-                                    node.overloads)
+    def handle_relation(self, node: Relation, parent: Node):
+        function_annotation = self.get_functional_annotation(node)
+        if function_annotation:
+            return node.reconstruct(node.name, node.fields, node.requires,
+                                    node.annotations | [function_annotation], node.overloads)
         return node.reconstruct(node.name, node.fields, node.requires, node.annotations, node.overloads)
 
-    def handle_update(self, node: ir.Update, parent: ir.Node):
-        if node.relation in self._functional_relations:
+    def handle_update(self, node: Update, parent: Node):
+        function_annotation = self.get_functional_annotation(node.relation)
+        if function_annotation:
             return node.reconstruct(node.engine, node.relation, node.args, node.effect,
-                                    node.annotations | [builtins.function_checked_annotation])
+                                    node.annotations | [function_annotation])
         return node.reconstruct(node.engine, node.relation, node.args, node.effect, node.annotations)