relationalai 0.11.4__py3-none-any.whl → 0.12.0__py3-none-any.whl

relationalai/clients/config.py

@@ -468,6 +468,13 @@ class Config():
             if not self.file_path:
                 self.file_path = "__inline__"
         self._handle_snowflake_fallback_configurations()
+        # Check if Azure platform is being used without the legacy dependency
+        if self.get("platform", "") == "azure":
+            try:
+                import railib # noqa
+            except ImportError:
+                from relationalai.errors import AzureLegacyDependencyMissingException
+                raise AzureLegacyDependencyMissingException() from None
 
     def fetch(self, profile:str|None=None):
         from relationalai.environments import runtime_env, TerminalEnvironment

relationalai/clients/direct_access_client.py

@@ -0,0 +1,113 @@
+from __future__ import annotations
+
+import requests
+from dataclasses import dataclass
+from urllib.parse import urlencode, quote
+from requests.adapters import HTTPAdapter
+from urllib3.util.retry import Retry
+from typing import Any, Dict, Optional, Tuple
+
+from relationalai.auth.token_handler import TokenHandler
+from relationalai.clients.config import Config
+from relationalai.clients.util import get_pyrel_version
+from relationalai import debugging
+from relationalai.tools.constants import Generation
+from relationalai.environments import runtime_env, SnowbookEnvironment
+
+@dataclass
+class Endpoint:
+    method: str
+    endpoint: str
+
+class DirectAccessClient:
+    """
+    DirectAccessClient is a client for direct service access without service function calls.
+    """
+
+    def __init__(self, config: Config, token_handler: TokenHandler, service_endpoint: str, generation: Optional[Generation] = None):
+        self._config: Config = config
+        self._token_handler: TokenHandler = token_handler
+        self.service_endpoint: str = service_endpoint
+        self.generation: Optional[Generation] = generation
+        self._is_snowflake_notebook = isinstance(runtime_env, SnowbookEnvironment)
+        self.endpoints: Dict[str, Endpoint] = {
+            "create_txn": Endpoint(method="POST", endpoint="/v1alpha1/transactions"),
+            "get_txn": Endpoint(method="GET", endpoint="/v1alpha1/transactions/{txn_id}"),
+            "get_txn_artifacts": Endpoint(method="GET", endpoint="/v1alpha1/transactions/{txn_id}/artifacts"),
+            "get_txn_problems": Endpoint(method="GET", endpoint="/v1alpha1/transactions/{txn_id}/problems"),
+            "get_txn_events": Endpoint(method="GET", endpoint="/v1alpha1/transactions/{txn_id}/events/{stream_name}"),
+            "get_package_versions": Endpoint(method="GET", endpoint="/v1alpha1/databases/{db_name}/package_versions"),
+            "get_model_package_versions": Endpoint(method="POST", endpoint="/v1alpha1/models/get_package_versions"),
+            "create_db": Endpoint(method="POST", endpoint="/v1alpha1/databases"),
+            "get_db": Endpoint(method="GET", endpoint="/v1alpha1/databases"),
+            "delete_db": Endpoint(method="DELETE", endpoint="/v1alpha1/databases/{db_name}"),
+            "release_index": Endpoint(method="POST", endpoint="/v1alpha1/index/release"),
+            "list_engines": Endpoint(method="GET", endpoint="/v1alpha1/engines"),
+            "get_engine": Endpoint(method="GET", endpoint="/v1alpha1/engines/{engine_type}/{engine_name}"),
+            "create_engine": Endpoint(method="POST", endpoint="/v1alpha1/engines/{engine_type}"),
+            "delete_engine": Endpoint(method="DELETE", endpoint="/v1alpha1/engines/{engine_type}/{engine_name}"),
+            "suspend_engine": Endpoint(method="POST", endpoint="/v1alpha1/engines/{engine_type}/{engine_name}/suspend"),
+            "resume_engine": Endpoint(method="POST", endpoint="/v1alpha1/engines/{engine_type}/{engine_name}/resume_async"),
+            "prepare_index": Endpoint(method="POST", endpoint="/v1alpha1/index/prepare"),
+        }
+        self.http_session = self._create_retry_session()
+
+    def _create_retry_session(self) -> requests.Session:
+        http_session = requests.Session()
+        retries = Retry(
+            total=3,
+            backoff_factor=0.3,
+            status_forcelist=[500, 502, 503, 504],
+            allowed_methods=frozenset({"GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"}),
+            raise_on_status=False
+        )
+        adapter = HTTPAdapter(max_retries=retries)
+        http_session.mount("http://", adapter)
+        http_session.mount("https://", adapter)
+        http_session.headers.update({"Connection": "keep-alive"})
+        return http_session
+
+    def request(
+        self,
+        endpoint: str,
+        payload: Dict[str, Any] | None = None,
+        headers: Dict[str, str] | None = None,
+        path_params: Dict[str, str] | None = None,
+        query_params: Dict[str, str] | None = None,
+    ) -> requests.Response:
+        """
+        Send a request to the service endpoint.
+        """
+        url, method = self._prepare_url(endpoint, path_params, query_params)
+        request_headers = self._prepare_headers(headers)
+        return self.http_session.request(method, url, json=payload, headers=request_headers)
+
+    def _prepare_url(self, endpoint: str, path_params: Dict[str, str] | None = None, query_params: Dict[str, str] | None = None) -> Tuple[str, str]:
+        try:
+            ep = self.endpoints[endpoint]
+        except KeyError:
+            raise ValueError(f"Invalid endpoint: {endpoint}. Available endpoints: {list(self.endpoints.keys())}")
+        url = f"{self.service_endpoint}{ep.endpoint}"
+        if path_params:
+            escaped_path_params = {k: quote(v, safe='') for k, v in path_params.items()}
+            url = url.format(**escaped_path_params)
+        if query_params:
+            url += '?' + urlencode(query_params)
+        return url, ep.method
+
+    def _prepare_headers(self, headers: Dict[str, str] | None) -> Dict[str, str]:
+        request_headers = {}
+        if headers:
+            request_headers.update(headers)
+        # Authorization tokens are not needed in a snowflake notebook environment
+        if not self._is_snowflake_notebook:
+            request_headers["Authorization"] = f'Snowflake Token="{self._token_handler.get_ingress_token(self.service_endpoint)}"'
+            # needed for oauth, does no harm for other authentication methods
+            request_headers["X-SF-SPCS-Authentication-Method"] = 'OAUTH'
+            request_headers["Content-Type"] = 'application/x-www-form-urlencoded'
+        request_headers["Accept"] = "application/json"
+
+        request_headers["user-agent"] = get_pyrel_version(self.generation)
+        request_headers["pyrel_program_id"] = debugging.get_program_span_id() or ""
+
+        return debugging.add_current_propagation_headers(request_headers)

relationalai/clients/snowflake.py

@@ -28,15 +28,11 @@ import requests
 import snowflake.connector
 import pyarrow as pa
 
-from dataclasses import dataclass
 from snowflake.snowpark import Session
 from snowflake.snowpark.context import get_active_session
 from . import result_helpers
 from .. import debugging
 from typing import Any, Dict, Iterable, Optional, Tuple, List, Literal, Union, cast
-from urllib.parse import urlencode, quote
-from requests.adapters import HTTPAdapter
-from urllib3.util.retry import Retry
 
 from pandas import DataFrame
 
@@ -44,10 +40,11 @@ from ..tools.cli_controls import Spinner
 from ..clients.types import AvailableModel, EngineState, Import, ImportSource, ImportSourceTable, ImportsStatus, SourceInfo, TransactionAsyncResponse
 from ..clients.config import Config, ConfigStore, ENDPOINT_FILE
 from ..clients.client import Client, ExportParams, ProviderBase, ResourcesBase
+from ..clients.direct_access_client import DirectAccessClient
 from ..clients.util import IdentityParser, escape_for_f_string, get_pyrel_version, get_with_retries, poll_with_specified_overhead, safe_json_loads, sanitize_module_name, scrub_exception, wrap_with_request_id, ms_to_timestamp
 from ..environments import runtime_env, HexEnvironment, SnowbookEnvironment
 from .. import dsl, rel, metamodel as m
-from ..errors import DuoSecurityFailed, EngineProvisioningFailed, EngineNameValidationException, EngineNotFoundException, EnginePending, EngineSizeMismatchWarning, EngineResumeFailed, Errors, InvalidAliasError, InvalidEngineSizeError, InvalidSourceTypeWarning, RAIAbortedTransactionError, RAIException, HexSessionException, SnowflakeAppMissingException, SnowflakeChangeTrackingNotEnabledException, SnowflakeDatabaseException, SnowflakeImportMissingException, SnowflakeInvalidSource, SnowflakeMissingConfigValuesException, SnowflakeProxyAPIDeprecationWarning, SnowflakeProxySourceError, SnowflakeRaiAppNotStarted, ModelNotFoundException, UnknownSourceWarning, ResponseStatusException, RowsDroppedFromTargetTableWarning
+from ..errors import DuoSecurityFailed, EngineProvisioningFailed, EngineNameValidationException, EngineNotFoundException, EnginePending, EngineSizeMismatchWarning, EngineResumeFailed, Errors, InvalidAliasError, InvalidEngineSizeError, InvalidSourceTypeWarning, RAIAbortedTransactionError, RAIException, HexSessionException, SnowflakeAppMissingException, SnowflakeChangeTrackingNotEnabledException, SnowflakeDatabaseException, SnowflakeImportMissingException, SnowflakeInvalidSource, SnowflakeMissingConfigValuesException, SnowflakeProxyAPIDeprecationWarning, SnowflakeProxySourceError, SnowflakeRaiAppNotStarted, ModelNotFoundException, UnknownSourceWarning, ResponseStatusException, RowsDroppedFromTargetTableWarning, QueryTimeoutExceededException
 from concurrent.futures import ThreadPoolExecutor
 from datetime import datetime, date, timedelta
 from snowflake.snowpark.types import StringType, StructField, StructType
@@ -92,6 +89,8 @@ TERMINAL_TXN_STATES = ["COMPLETED", "ABORTED"]
 
 DUO_TEXT = "duo security"
 
+TXN_ABORT_REASON_TIMEOUT = "transaction timeout"
+
 #--------------------------------------------------
 # Helpers
 #--------------------------------------------------
@@ -1307,7 +1306,7 @@ Otherwise, remove it from your '{profile}' configuration profile.
             response = self._exec(f"CALL {APP_NAME}.api.get_transaction('{txn_id}',{headers});")
             assert response, f"No results from get_transaction('{txn_id}')"
 
-        response_row = next(iter(response))
+        response_row = next(iter(response)).asDict()
         status: str = response_row['STATE']
 
         # remove the transaction from the pending list if it's completed or aborted
@@ -1315,6 +1314,16 @@ 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,
+            )
+
         # @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"
 
@@ -2957,104 +2966,6 @@ def Graph(
 #--------------------------------------------------
 # Note: All direct access components should live in a separate file
 
-@dataclass
-class Endpoint:
-    method: str
-    endpoint: str
-
-class DirectAccessClient:
-    """
-    DirectAccessClient is a client for direct service access without service function calls.
-    """
-
-    def __init__(self, config: Config, token_handler: TokenHandler, service_endpoint: str, generation: Optional[Generation] = None):
-        self._config: Config = config
-        self._token_handler: TokenHandler = token_handler
-        self.service_endpoint: str = service_endpoint
-        self.generation: Optional[Generation] = generation
-        self._is_snowflake_notebook = isinstance(runtime_env, SnowbookEnvironment)
-        self.endpoints: Dict[str, Endpoint] = {
-            "create_txn": Endpoint(method="POST", endpoint="/v1alpha1/transactions"),
-            "get_txn": Endpoint(method="GET", endpoint="/v1alpha1/transactions/{txn_id}"),
-            "get_txn_artifacts": Endpoint(method="GET", endpoint="/v1alpha1/transactions/{txn_id}/artifacts"),
-            "get_txn_problems": Endpoint(method="GET", endpoint="/v1alpha1/transactions/{txn_id}/problems"),
-            "get_txn_events": Endpoint(method="GET", endpoint="/v1alpha1/transactions/{txn_id}/events/{stream_name}"),
-            "get_package_versions": Endpoint(method="GET", endpoint="/v1alpha1/databases/{db_name}/package_versions"),
-            "get_model_package_versions": Endpoint(method="POST", endpoint="/v1alpha1/models/get_package_versions"),
-            "create_db": Endpoint(method="POST", endpoint="/v1alpha1/databases"),
-            "get_db": Endpoint(method="GET", endpoint="/v1alpha1/databases"),
-            "delete_db": Endpoint(method="DELETE", endpoint="/v1alpha1/databases/{db_name}"),
-            "release_index": Endpoint(method="POST", endpoint="/v1alpha1/index/release"),
-            "list_engines": Endpoint(method="GET", endpoint="/v1alpha1/engines"),
-            "get_engine": Endpoint(method="GET", endpoint="/v1alpha1/engines/{engine_type}/{engine_name}"),
-            "create_engine": Endpoint(method="POST", endpoint="/v1alpha1/engines/{engine_type}"),
-            "delete_engine": Endpoint(method="DELETE", endpoint="/v1alpha1/engines/{engine_type}/{engine_name}"),
-            "suspend_engine": Endpoint(method="POST", endpoint="/v1alpha1/engines/{engine_type}/{engine_name}/suspend"),
-            "resume_engine": Endpoint(method="POST", endpoint="/v1alpha1/engines/{engine_type}/{engine_name}/resume_async"),
-            "prepare_index": Endpoint(method="POST", endpoint="/v1alpha1/index/prepare"),
-        }
-        self.http_session = self._create_retry_session()
-
-    def _create_retry_session(self) -> requests.Session:
-        http_session = requests.Session()
-        retries = Retry(
-            total=3,
-            backoff_factor=0.3,
-            status_forcelist=[500, 502, 503, 504],
-            allowed_methods=frozenset({"GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"}),
-            raise_on_status=False
-        )
-        adapter = HTTPAdapter(max_retries=retries)
-        http_session.mount("http://", adapter)
-        http_session.mount("https://", adapter)
-        http_session.headers.update({"Connection": "keep-alive"})
-        return http_session
-
-    def request(
-        self,
-        endpoint: str,
-        payload: Dict[str, Any] | None = None,
-        headers: Dict[str, str] | None = None,
-        path_params: Dict[str, str] | None = None,
-        query_params: Dict[str, str] | None = None,
-    ) -> requests.Response:
-        """
-        Send a request to the service endpoint.
-        """
-        url, method = self._prepare_url(endpoint, path_params, query_params)
-        request_headers = self._prepare_headers(headers)
-        return self.http_session.request(method, url, json=payload, headers=request_headers)
-
-    def _prepare_url(self, endpoint: str, path_params: Dict[str, str] | None = None, query_params: Dict[str, str] | None = None) -> Tuple[str, str]:
-        try:
-            ep = self.endpoints[endpoint]
-        except KeyError:
-            raise ValueError(f"Invalid endpoint: {endpoint}. Available endpoints: {list(self.endpoints.keys())}")
-        url = f"{self.service_endpoint}{ep.endpoint}"
-        if path_params:
-            escaped_path_params = {k: quote(v, safe='') for k, v in path_params.items()}
-            url = url.format(**escaped_path_params)
-        if query_params:
-            url += '?' + urlencode(query_params)
-        return url, ep.method
-
-    def _prepare_headers(self, headers: Dict[str, str] | None) -> Dict[str, str]:
-        request_headers = {}
-        if headers:
-            request_headers.update(headers)
-        # Authorization tokens are not needed in a snowflake notebook environment
-        if not self._is_snowflake_notebook:
-            request_headers["Authorization"] = f'Snowflake Token="{self._token_handler.get_ingress_token(self.service_endpoint)}"'
-            # needed for oauth, does no harm for other authentication methods
-            request_headers["X-SF-SPCS-Authentication-Method"] = 'OAUTH'
-            request_headers["Content-Type"] = 'application/x-www-form-urlencoded'
-        request_headers["Accept"] = "application/json"
-
-        request_headers["user-agent"] = get_pyrel_version(self.generation)
-        request_headers["pyrel_program_id"] = debugging.get_program_span_id() or ""
-
-        return debugging.add_current_propagation_headers(request_headers)
-
 class DirectAccessResources(Resources):
     """
     Resources class for Direct Service Access avoiding Snowflake service functions.
@@ -3068,7 +2979,14 @@ class DirectAccessResources(Resources):
         reset_session: bool = False,
         generation: Optional[Generation] = None,
     ):
-        super().__init__(generation=generation, profile=profile, config=config, connection=connection, dry_run=dry_run)
+        super().__init__(
+            generation=generation,
+            profile=profile,
+            config=config,
+            connection=connection,
+            reset_session=reset_session,
+            dry_run=dry_run,
+        )
         self._endpoint_info = ConfigStore(ENDPOINT_FILE)
         self._service_endpoint = ""
         self._direct_access_client = None
@@ -3310,13 +3228,24 @@ class DirectAccessResources(Resources):
             assert response, f"No results from get_transaction('{txn_id}')"
 
         response_content = response.json()
-        status: str = response_content["transaction"]['state']
+        transaction = response_content["transaction"]
+        status: str = transaction['state']
 
         # remove the transaction from the pending list if it's completed or aborted
         if status in ["COMPLETED", "ABORTED"]:
             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,
+            )
+
         # @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/early_access/metamodel/rewrite/__init__.py

@@ -1,5 +1,7 @@
-from relationalai.semantics.metamodel.rewrite import Splinter, RewriteListTypes, GarbageCollectNodes, Flatten, \
-    DNFUnionSplitter, ExtractKeys, ExtractNestedLogicals, FDConstraints, flatten
+from relationalai.semantics.metamodel.rewrite import Flatten, \
+    DNFUnionSplitter, ExtractNestedLogicals, flatten
+from relationalai.semantics.lqp.rewrite import Splinter,  \
+    ExtractKeys, FDConstraints
 
-__all__ = ["Splinter", "RewriteListTypes", "GarbageCollectNodes", "Flatten", "DNFUnionSplitter", "ExtractKeys",
+__all__ = ["Splinter", "Flatten", "DNFUnionSplitter", "ExtractKeys",
            "ExtractNestedLogicals", "FDConstraints", "flatten"]

relationalai/early_access/rel/rewrite/__init__.py

@@ -1,4 +1,4 @@
-from relationalai.semantics.rel.rewrite import CDC, ExtractCommon, QuantifyVars
+from relationalai.semantics.lqp.rewrite import CDC, ExtractCommon, QuantifyVars
 
 __all__ = [
     "CDC",

relationalai/errors.py

@@ -2397,17 +2397,18 @@ class UnsupportedColumnTypesWarning(RAIWarning):
 """)
 
 class QueryTimeoutExceededException(RAIException):
-    def __init__(self, timeout_mins: int, config_file_path: str | None = None):
+    def __init__(self, timeout_mins: int, query_id: str | None = None, config_file_path: str | None = None):
         self.timeout_mins = timeout_mins
-        self.message = f"Query execution time exceeded the specified timeout of {timeout_mins} minutes."
         self.name = "Query Timeout Exceeded"
+        self.message = f"Query execution time exceeded the specified timeout of {self.timeout_mins} minutes."
+        self.query_id = query_id or ""
         self.config_file_path = config_file_path or ""
         self.content = self.format_message()
         super().__init__(self.message, self.name, self.content)
 
     def format_message(self):
         return textwrap.dedent(f"""
-        {self.message}
+        Query execution time exceeded the specified timeout of {self.timeout_mins} minutes{f' for query with ID: {self.query_id}' if self.query_id else ''}.
 
         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.
         """)
@@ -2432,3 +2433,23 @@ class AzureUnsupportedQueryTimeoutException(RAIException):
         Please remove the 'query_timeout_mins' from your configuration file{f' (stored in {self.config_file_path})' if self.config_file_path else ''} when running on platform Azure.
         """)
 
+class AzureLegacyDependencyMissingException(RAIException):
+    def __init__(self):
+        self.message = "The Azure platform requires the 'legacy' extras to be installed."
+        self.name = "Azure Legacy Dependency Missing"
+        self.content = self.format_message()
+        super().__init__(self.message, self.name, self.content)
+
+    def format_message(self):
+        return textwrap.dedent("""
+        The Azure platform requires the 'rai-sdk' package, which is not installed.
+
+        To use the Azure platform, please install the legacy extras:
+
+            pip install relationalai[legacy]
+
+        Or if upgrading an existing installation:
+
+            pip install --upgrade relationalai[legacy]
+        """)
+

relationalai/semantics/internal/annotations.py

@@ -5,3 +5,4 @@ external = Relationship.builtins["external"]
 concept_population = Relationship.builtins["concept_population"]
 function = Relationship.builtins["function"]
 from_cdc = Relationship.builtins["from_cdc"]
+track = Relationship.builtins["track"]

relationalai/semantics/lqp/builtins.py

@@ -11,4 +11,5 @@ adhoc_annotation = f.annotation(adhoc, [])
 annotations_to_emit = FrozenOrderedSet([
     adhoc.name,
     builtins.function.name,
+    builtins.track.name,
 ])

relationalai/semantics/lqp/passes.py

@@ -2,13 +2,12 @@ from relationalai.semantics.metamodel.compiler import Pass
 from relationalai.semantics.metamodel import ir, builtins as rel_builtins, factory as f, visitor
 from relationalai.semantics.metamodel.typer import Checker, InferTypes, typer
 from relationalai.semantics.metamodel import helpers, types
-from relationalai.semantics.metamodel.rewrite import (Splinter, ExtractNestedLogicals, ExtractKeys, FDConstraints,
-                                                      DNFUnionSplitter, DischargeConstraints)
 from relationalai.semantics.metamodel.util import FrozenOrderedSet
 
 from relationalai.semantics.metamodel.rewrite import Flatten
-# TODO: Move this into metamodel.rewrite
-from relationalai.semantics.rel.rewrite import QuantifyVars, CDC, ExtractCommon
+
+from ..metamodel.rewrite import DischargeConstraints, DNFUnionSplitter, ExtractNestedLogicals
+from .rewrite import CDC, ExtractCommon, ExtractKeys, FDConstraints, QuantifyVars, Splinter
 
 from relationalai.semantics.lqp.utils import output_names
 

relationalai/semantics/{rel → lqp}/rewrite/__init__.py

@@ -1,9 +1,15 @@
 from .cdc import CDC
 from .extract_common import ExtractCommon
+from .extract_keys import ExtractKeys
+from .fd_constraints import FDConstraints
 from .quantify_vars import QuantifyVars
+from .splinter import Splinter
 
 __all__ = [
     "CDC",
     "ExtractCommon",
+    "ExtractKeys",
+    "FDConstraints",
     "QuantifyVars",
+    "Splinter",
 ]

relationalai/semantics/metamodel/builtins.py

@@ -496,6 +496,17 @@ function = f.relation("function", [f.input_field("code", types.Symbol)])
 function_checked_annotation = f.annotation(function, [f.lit("checked")])
 function_annotation = f.annotation(function, [])
 
+# Indicates this relation should be tracked in telemetry. Only supported for Relationships.
+# `RAI_BackIR.with_relation_tracking` produces log messages at the start and end of each
+# SCC evaluation, if any declarations bear the `track` annotation.
+track = f.relation("track", [
+    # BackIR evaluation expects 2 parameters on the track annotation: the tracking
+    # library name and tracking relation name, which appear as log metadata fields.
+    f.input_field("library", types.Symbol),
+    f.input_field("relation", types.Symbol)
+])
+track_annotation = f.annotation(track, [])
+
 # All ir nodes marked by this annotation will be removed from the final metamodel before compilation.
 # Specifically it happens in `Flatten` pass when rewrites for `require` happen
 discharged = f.relation("discharged", [])
@@ -672,7 +683,7 @@ def _compute_builtin_overloads() -> list[ir.Relation]:
     return overloads
 
 # manually maintain the list of relations that are actually annotations
-builtin_annotations = [external, export, concept_population, from_cdc, from_cast]
+builtin_annotations = [external, export, concept_population, from_cdc, from_cast, track]
 builtin_annotations_by_name = dict((r.name, r) for r in builtin_annotations)
 
 builtin_relations = _compute_builtin_relations()

relationalai/semantics/metamodel/rewrite/__init__.py

@@ -1,12 +1,6 @@
-from .splinter import Splinter
-from .list_types import RewriteListTypes
-from .gc_nodes import GarbageCollectNodes
-from .flatten import Flatten
+from .discharge_constraints import DischargeConstraints
 from .dnf_union_splitter import DNFUnionSplitter
-from .extract_keys import ExtractKeys
 from .extract_nested_logicals import ExtractNestedLogicals
-from .fd_constraints import FDConstraints
-from .discharge_constraints import DischargeConstraints
+from .flatten import Flatten
 
-__all__ = ["Splinter", "RewriteListTypes", "GarbageCollectNodes", "Flatten", "DNFUnionSplitter", "ExtractKeys",
-           "ExtractNestedLogicals", "FDConstraints", "DischargeConstraints"]
+__all__ = ["DischargeConstraints", "DNFUnionSplitter", "ExtractNestedLogicals", "Flatten"]