relationalai 0.12.3__py3-none-any.whl → 0.12.6__py3-none-any.whl

relationalai/__init__.py

@@ -1,3 +1,7 @@
+"""
+The RelationalAI Python SDK.
+"""
+
 from __future__ import annotations
 import importlib.metadata
 

relationalai/clients/snowflake.py

@@ -83,6 +83,7 @@ FIELD_MAP = {
 VALID_IMPORT_STATES = ["PENDING", "PROCESSING", "QUARANTINED", "LOADED"]
 ENGINE_ERRORS = ["engine is suspended", "create/resume", "engine not found", "no engines found", "engine was deleted"]
 ENGINE_NOT_READY_MSGS = ["engine is in pending", "engine is provisioning"]
+DATABASE_ERRORS = ["database not found"]
 PYREL_ROOT_DB = 'pyrel_root_db'
 
 TERMINAL_TXN_STATES = ["COMPLETED", "ABORTED"]
@@ -281,6 +282,8 @@ def _sanitize_user_name(user: str) -> str:
 def _is_engine_issue(response_message: str) -> bool:
     return any(kw in response_message.lower() for kw in ENGINE_ERRORS + ENGINE_NOT_READY_MSGS)
 
+def _is_database_issue(response_message: str) -> bool:
+    return any(kw in response_message.lower() for kw in DATABASE_ERRORS)
 
 
 #--------------------------------------------------
@@ -298,6 +301,7 @@ class Resources(ResourcesBase):
         dry_run: bool = False,
         reset_session: bool = False,
         generation: Generation | None = None,
+        language: str = "rel",
     ):
         super().__init__(profile, config=config)
         self._token_handler: TokenHandler | None = None
@@ -315,6 +319,8 @@ class Resources(ResourcesBase):
         # self.sources contains fully qualified Snowflake table/view names
         self.sources: set[str] = set()
         self._sproc_models = None
+        self.database = ""
+        self.language = language
         atexit.register(self.cancel_pending_transactions)
 
     @property
@@ -452,6 +458,7 @@ class Resources(ResourcesBase):
             rai_app = self.config.get("rai_app_name", "")
             current_role = self.config.get("role")
             engine = self.get_default_engine_name()
+            engine_size = self.config.get_default_engine_size()
             assert isinstance(rai_app, str), f"rai_app_name must be a string, not {type(rai_app)}"
             assert isinstance(engine, str), f"engine must be a string, not {type(engine)}"
             print("\n")
@@ -460,9 +467,15 @@ class Resources(ResourcesBase):
             if re.search(f"database '{rai_app}' does not exist or not authorized.".lower(), orig_message):
                 exception = SnowflakeAppMissingException(rai_app, current_role)
                 raise exception from None
-            if any(keyword in orig_message for keyword in ENGINE_ERRORS):
+            if _is_engine_issue(orig_message) or _is_database_issue(orig_message):
                 try:
-                    self.auto_create_engine(engine)
+                    self._poll_use_index(
+                        app_name=self.get_app_name(),
+                        sources=self.sources,
+                        model=self.database,
+                        engine_name=engine,
+                        engine_size=engine_size
+                    )
                     return self._exec(code, params, raw=raw, help=help)
                 except EngineNameValidationException as e:
                     raise EngineNameValidationException(engine) from e
@@ -767,7 +780,7 @@ Otherwise, remove it from your '{profile}' configuration profile.
             keep_database = not force and self.config.get("reuse_model", True)
             with debugging.span("release_index", name=name, keep_database=keep_database, language=language):
                 #TODO add headers to release_index
-                response = self._exec(f"call {APP_NAME}.api.release_index('{name}', OBJECT_CONSTRUCT('keep_database', {keep_database}, 'language', '{language}'));")
+                response = self._exec(f"call {APP_NAME}.api.release_index('{name}', OBJECT_CONSTRUCT('keep_database', {keep_database}, 'language', '{language}', 'user_agent', '{get_pyrel_version(self.generation)}'));")
                 if response:
                     result = next(iter(response))
                     obj = json.loads(result["RELEASE_INDEX"])
@@ -788,14 +801,13 @@ Otherwise, remove it from your '{profile}' configuration profile.
             headers = debugging.gen_current_propagation_headers()
             self._exec(f"call {APP_NAME}.api.clone_database('{target_name}', '{source_name}', {nowait_durable}, {headers});")
 
-    def poll_use_index(
+    def _poll_use_index(
         self,
         app_name: str,
         sources: Iterable[str],
         model: str,
         engine_name: str,
         engine_size: str | None = None,
-        language: str = "rel",
         program_span_id: str | None = None,
         headers: Dict | None = None,
     ):
@@ -806,12 +818,42 @@ Otherwise, remove it from your '{profile}' configuration profile.
             model,
             engine_name,
             engine_size,
-            language,
+            self.language,
             program_span_id,
             headers,
             self.generation
         ).poll()
 
+    def maybe_poll_use_index(
+        self,
+        app_name: str,
+        sources: Iterable[str],
+        model: str,
+        engine_name: str,
+        engine_size: str | None = None,
+        program_span_id: str | None = None,
+        headers: Dict | None = None,
+    ):
+        """Only call poll() if there are sources to process and cache is not valid."""
+        sources_list = list(sources)
+        self.database = model
+        if sources_list:
+            poller = UseIndexPoller(
+                self,
+                app_name,
+                sources_list,
+                model,
+                engine_name,
+                engine_size,
+                self.language,
+                program_span_id,
+                headers,
+                self.generation
+            )
+            # If cache is valid (data freshness has not expired), skip polling
+            if not poller.cache.is_valid():
+                return poller.poll()
+
     #--------------------------------------------------
     # Models
     #--------------------------------------------------
@@ -1868,9 +1910,19 @@ Otherwise, remove it from your '{profile}' configuration profile.
             )
         except Exception as e:
             err_message = str(e).lower()
-            if _is_engine_issue(err_message):
-                self.auto_create_engine(engine, headers=headers)
-                self._exec_async_v2(
+            if _is_engine_issue(err_message) or _is_database_issue(err_message):
+                engine_name = engine or self.get_default_engine_name()
+                engine_size = self.config.get_default_engine_size()
+                self._poll_use_index(
+                    app_name=self.get_app_name(),
+                    sources=self.sources,
+                    model=database,
+                    engine_name=engine_name,
+                    engine_size=engine_size,
+                    headers=headers,
+                )
+
+                return self._exec_async_v2(
                     database, engine, raw_code_b64, inputs, readonly, nowait_durable,
                     headers=headers, bypass_index=bypass_index, language='lqp',
                     query_timeout_mins=query_timeout_mins,
@@ -1908,8 +1960,17 @@ Otherwise, remove it from your '{profile}' configuration profile.
             )
         except Exception as e:
             err_message = str(e).lower()
-            if _is_engine_issue(err_message):
-                self.auto_create_engine(engine, headers=headers)
+            if _is_engine_issue(err_message) or _is_database_issue(err_message):
+                engine_name = engine or self.get_default_engine_name()
+                engine_size = self.config.get_default_engine_size()
+                self._poll_use_index(
+                    app_name=self.get_app_name(),
+                    sources=self.sources,
+                    model=database,
+                    engine_name=engine_name,
+                    engine_size=engine_size,
+                    headers=headers,
+                )
                 return self._exec_async_v2(
                     database,
                     engine,
@@ -2972,13 +3033,12 @@ class SnowflakeClient(Client):
 
         query_attrs_dict = json.loads(headers.get("X-Query-Attributes", "{}")) if headers else {}
         with debugging.span("poll_use_index", sources=self.resources.sources, model=model, engine=engine_name, **query_attrs_dict):
-            self.poll_use_index(
+            self.maybe_poll_use_index(
                 app_name=app_name,
                 sources=self.resources.sources,
                 model=model,
                 engine_name=engine_name,
                 engine_size=engine_size,
-                language="rel",
                 program_span_id=program_span_id,
                 headers=headers
             )
@@ -2989,29 +3049,24 @@ class SnowflakeClient(Client):
         if isolated and not self.keep_model:
             atexit.register(self.delete_database)
 
-    # Polling for use_index
-    # if data is ready, break the loop
-    # if data is not ready, print the status of the tables or engines
-    # if data is not ready and there are errors, collect the errors and raise exceptions
-    def poll_use_index(
+    def maybe_poll_use_index(
         self,
         app_name: str,
         sources: Iterable[str],
         model: str,
         engine_name: str,
         engine_size: str | None = None,
-        language: str = "rel",
         program_span_id: str | None = None,
         headers: Dict | None = None,
     ):
+        """Only call _poll_use_index if there are sources to process."""
         assert isinstance(self.resources, Resources)
-        return self.resources.poll_use_index(
+        return self.resources.maybe_poll_use_index(
             app_name=app_name,
             sources=sources,
             model=model,
             engine_name=engine_name,
             engine_size=engine_size,
-            language=language,
             program_span_id=program_span_id,
             headers=headers
         )
@@ -3136,6 +3191,7 @@ class DirectAccessResources(Resources):
         dry_run: bool = False,
         reset_session: bool = False,
         generation: Optional[Generation] = None,
+        language: str = "rel",
     ):
         super().__init__(
             generation=generation,
@@ -3144,11 +3200,13 @@ class DirectAccessResources(Resources):
             connection=connection,
             reset_session=reset_session,
             dry_run=dry_run,
+            language=language,
         )
         self._endpoint_info = ConfigStore(ENDPOINT_FILE)
         self._service_endpoint = ""
         self._direct_access_client = None
         self.generation = generation
+        self.database = ""
 
     @property
     def service_endpoint(self) -> str:
@@ -3226,9 +3284,18 @@ class DirectAccessResources(Resources):
 
                     # fix engine on engine error and retry
                     # Skip auto-retry if skip_auto_create is True to avoid recursion
-                    if _is_engine_issue(message) and not skip_auto_create:
-                        engine = payload.get("engine_name", "") if payload else ""
-                        self.auto_create_engine(engine)
+                    if (_is_engine_issue(message) and not skip_auto_create) or _is_database_issue(message):
+                        engine_name = payload.get("caller_engine_name", "") if payload else ""
+                        engine_name = engine_name or self.get_default_engine_name()
+                        engine_size = self.config.get_default_engine_size()
+                        self._poll_use_index(
+                            app_name=self.get_app_name(),
+                            sources=self.sources,
+                            model=self.database,
+                            engine_name=engine_name,
+                            engine_size=engine_size,
+                            headers=headers,
+                        )
                         response = _send_request()
             except requests.exceptions.ConnectionError as e:
                 if "NameResolutionError" in str(e):
@@ -3356,14 +3423,13 @@ class DirectAccessResources(Resources):
 
             return response.json()
 
-    def poll_use_index(
+    def _poll_use_index(
         self,
         app_name: str,
         sources: Iterable[str],
         model: str,
         engine_name: str,
         engine_size: str | None = None,
-        language: str = "rel",
         program_span_id: str | None = None,
         headers: Dict | None = None,
     ):
@@ -3374,12 +3440,42 @@ class DirectAccessResources(Resources):
             model=model,
             engine_name=engine_name,
             engine_size=engine_size,
-            language=language,
+            language=self.language,
             program_span_id=program_span_id,
             headers=headers,
             generation=self.generation,
         ).poll()
 
+    def maybe_poll_use_index(
+        self,
+        app_name: str,
+        sources: Iterable[str],
+        model: str,
+        engine_name: str,
+        engine_size: str | None = None,
+        program_span_id: str | None = None,
+        headers: Dict | None = None,
+    ):
+        """Only call poll() if there are sources to process and cache is not valid."""
+        sources_list = list(sources)
+        self.database = model
+        if sources_list:
+            poller = DirectUseIndexPoller(
+                self,
+                app_name=app_name,
+                sources=sources_list,
+                model=model,
+                engine_name=engine_name,
+                engine_size=engine_size,
+                language=self.language,
+                program_span_id=program_span_id,
+                headers=headers,
+                generation=self.generation,
+            )
+            # If cache is valid (data freshness has not expired), skip polling
+            if not poller.cache.is_valid():
+                return poller.poll()
+
     def _check_exec_async_status(self, txn_id: str, headers: Dict[str, str] | None = None) -> bool:
         """Check whether the given transaction has completed."""
 
@@ -3522,7 +3618,12 @@ class DirectAccessResources(Resources):
             with debugging.span("release_index", name=name, keep_database=keep_database, language=language):
                 response = self.request(
                     "release_index",
-                    payload={"model_name": name, "keep_database": keep_database, "language": language},
+                    payload={
+                        "model_name": name,
+                        "keep_database": keep_database,
+                        "language": language,
+                        "user_agent": get_pyrel_version(self.generation),
+                    },
                     headers=prop_hdrs,
                 )
                 if (

relationalai/clients/use_index_poller.py

@@ -527,6 +527,9 @@ class UseIndexPoller:
                 else:
                     self.should_check_cdc = False
 
+                if engines and self.init_engine_async:
+                    self.init_engine_async = False
+
                 break_loop = False
                 has_stream_errors = False
                 has_general_errors = False

relationalai/{semantics/reasoners/graph → experimental}/paths/README.md

@@ -12,8 +12,8 @@ The following is a minimal working example:
 from relationalai.semantics import Model, Integer, select, String, where, define
 from relationalai.semantics.std import strings
 
-from relationalai.semantics.reasoners.graph.paths.graph import Graph
-from relationalai.semantics.reasoners.graph.paths.api import node, edge, path, star, match
+from relationalai.experimental.paths.graph import Graph
+from relationalai.experimental.paths.api import node, edge, path, star, match
 
 model = Model("my_paths")
 

relationalai/experimental/paths/__init__.py

@@ -1,309 +1,14 @@
-import os
-from typing import List, Set, Tuple, Union
-import zipfile
-from enum import Enum
-
-from relationalai import dsl
-from relationalai.dsl import Graph, Instance, Type, Property, create_vars, get_graph, next_id
-from relationalai.metamodel import ActionType, Builtins
-from relationalai.std import rel
-
-class VarLength:
-    def __init__(self, type: Type):
-        self.type = type
-        self.min = 1
-        self.max = 1
-
-    def __getitem__(self, key):
-        if isinstance(key, slice):
-            self.min = key.start
-            self.max = key.stop
-        else:
-            raise ValueError("VarLength only supports slicing")
-        return self
-
-class PathSelection(Enum):
-    ALL = "all" # return all paths between source and target
-    SINGLE = "single" # return a single path between source and target
-    LIMIT = "limit" # return a specified number paths between source and target, as given by parameter k.
-    RANDOM = "random" # return a single randomly-chosen path between source and target, seeded with parameter seed.
-    UNIFORM_RANDOM = "uniform_random" # return a single uniform randomly-chosen path between source and target, seeded with parameter seed.
-
-class PathGroup(Enum):
-    ANY = "any"
-    FOREACH = "for_each"
-
-class Strategy(Enum):
-    FROM_SOURCE = "from_source"
-    TWO_SIDED = "two_sided"
-
-class PathQuery:
-
-    def __init__(
-        self,
-        model: Graph,
-        segments: List[Union[Instance,VarLength]],
-        backwards_edge: str,
-        group: PathGroup = PathGroup.FOREACH,
-        selection: PathSelection = PathSelection.ALL,
-        strategy: Strategy = Strategy.FROM_SOURCE,
-    ):
-        # create types and instance
-        self._path_id = next_id()
-        self._conn_relation_name = f"conn_{self._path_id}"
-        self._model = model
-
-        # create types
-        self._Path = dsl.Type(model, f"Path{self._path_id}", omit_intrinsic_type_in_hash=True)
-        self._PathNode = dsl.Type(model, f"PathNode{self._path_id}", omit_intrinsic_type_in_hash=True)
-        self._PathEdge = dsl.Type(model, f"PathEdge{self._path_id}", omit_intrinsic_type_in_hash=True)
-
-        # declare multivalued attributes
-        self._Path.nodes.has_many()
-        self._Path.edges.has_many()
-
-        # mark these types and their attributes as @no_inline
-        self._add_noinline(self._Path, [])
-        self._add_noinline(self._PathNode, ["path", "index", "node_id"])
-        self._add_noinline(self._PathEdge, ["path", "index", "label"])
-
-        # create instance
-        self._instance = self._Path()
-        self._match_called = False
-
-        # store path query attributes
-        self._segments = segments
-        self._backwards_edge = backwards_edge
-        self._group = group
-        self._selection = selection
-        self._strategy = strategy
-
-    def _add_noinline(self, typ, props):
-        typ._type.parents.append(Builtins.NoInlineAnnotation)
-        for prop in props:
-            getattr(typ, prop)._prop.parents.append(Builtins.NoInlineAnnotation)
-
-    def _match(self):
-        # avoid emitting the Rel multiple times, which throws off Pathfinder
-        if not self._match_called:
-            self._match_called = True
-
-            begin_var, end_var = self._match_inner()
-
-            self._begin_var = begin_var
-            self._end_var = end_var
-
-        return self._conn_relation_name
-
-    def _match_inner(self):
-        begin_var = None
-        cur_end = None
-
-        for idx, segment in enumerate(self._segments):
-            if isinstance(segment, VarLength):
-                var_length_rel_name = self._var_length(idx, segment)
-                seg_begin, seg_end = create_vars(2)
-                getattr(rel, var_length_rel_name)(seg_begin, seg_end)
-
-                if begin_var is None:
-                    begin_var = seg_begin
-                    cur_end = seg_end
-                else:
-                    getattr(seg_begin, self._backwards_edge) == cur_end
-                    cur_end = seg_end
-            else:
-                next_var = segment
-                if begin_var is None:
-                    begin_var = next_var
-                    cur_end = begin_var
-                else:
-                    getattr(next_var, self._backwards_edge) == cur_end
-                    cur_end = next_var
-
-        conn_relation = getattr(rel, self._conn_relation_name)
-        conn_relation._rel.parents.append(Builtins.NoInlineAnnotation)
-        conn_relation.add(begin_var, cur_end)
-
-        return begin_var, cur_end
-
-    def _var_length(self, segment_idx: int, segment: VarLength):
-        # relations for each length
-        prefix = f"conn_{self._path_id}_segment_{segment_idx}"
-
-        lines = []
-
-        for length in range(1, segment.max+1):
-            length_rel_name = f"{prefix}_length_{length}"
-
-            if length == 1:
-                lines.extend([
-                    f"def {length_rel_name}(a, a1):",
-                    f"    {segment.type._type.name}(a) and",
-                    "    a = a1"
-                ])
-                # would just project `(a, a)`, but that gives a
-                # bunch of Rel warnings`
-
-            else:
-                prev_length_relation_name = f"{prefix}_length_{length-1}"
-                lines.extend([
-                    f"def {length_rel_name}(a, c):",
-                    "    exists((b) |",
-                    f"       {prev_length_relation_name}(a, b) and",
-                    f"       {segment.type._type.name}(c) and",
-                    f"       {self._backwards_edge}(c, b)",
-                    "     )"
-                ])
-
-        # union them all together in overall segment relation
-        lines.extend([
-            f"def {prefix}(start, finish):",
-            "    " + " or\n    ".join([
-                f"{prefix}_length_{length}(start, finish)"
-                for length in range(1, segment.max+1)
-            ]),
-        ])
-
-        self._model.install_raw(
-            '\n'.join(lines),
-            name=f"path_{self._path_id}_segment_{segment_idx}",
-        )
-
-        return prefix
-
-
-class PathInstance(Instance):
-
-    def __init__(
-        self,
-        path_query: PathQuery,
-    ):
-        self._path_query = path_query
-
-        self._already_called = False
-        super().__init__(
-            self._path_query._model,
-            ActionType.Get, [self._path_query._instance],
-            named={},
-        )
-
-    # implements producer API
-    def _to_var(self):
-        self._invoke_pathfinder()
-        return super()._to_var()
-
-    def _invoke_pathfinder(self):
-        if self._already_called:
-            return
-        self._already_called = True
-
-        # ensure that the pathfinder Rel library is installed
-        _install_pathfinder(self._path_query._model)
-
-        source_rel = f"source_{self._path_query._path_id}"
-        target_rel = f"target_{self._path_query._path_id}"
-        getattr(rel, source_rel).add(self._path_query._begin_var)
-        getattr(rel, target_rel).add(self._path_query._end_var)
-
-        shortest_paths_rel = f"shortest_paths_{self._path_query._path_id}"
-        path_edge_rel = f"path_edge_{self._path_query._path_id}"
-        path_node_rel = f"path_node_{self._path_query._path_id}"
-
-        invocation = f"""
-        @no_inline
-        def {shortest_paths_rel} {{
-            ::pathfinder::shortest_paths[
-                :{self._path_query._group.value},
-                :{self._path_query._selection.value},
-                :{self._path_query._strategy.value},
-                {self._path_query._conn_relation_name},
-                {source_rel},
-                {target_rel}
-            ]
-        }}
-
-        @no_inline
-        def {path_node_rel}(path, index, id):
-            {shortest_paths_rel}(path, :node, index, id)
-
-        @no_inline
-        def {path_edge_rel}(path, index, label):
-            {shortest_paths_rel}(path, :edge_label, index, label)
-        """
-
-        self._path_query._model.install_raw(invocation)
-
-        # select out nodes
-        with self._path_query._model.rule():
-            path_id, index, node_label = create_vars(3)
-            getattr(rel, path_node_rel)(path_id, index, node_label)
-            path = self._path_query._Path.add(id=path_id)
-            # TODO: avoid perf hit of going in both directions (path<->node)
-            node = self._path_query._PathNode.add(path=path, index=index, value=node_label)
-            path.nodes.add(node)
-
-        # select out edges
-        with self._path_query._model.rule():
-            path_id, index, label = create_vars(3)
-            getattr(rel, path_edge_rel)(path_id, index, label)
-            path = self._path_query._Path.add(id=path_id)
-            # TODO: avoid perf hit of going in both directions (path<->edge)
-            edge = self._path_query._PathEdge.add(path=path, index=index, label=label)
-            path.edges.add(edge)
-
-# TODO: automatically infer this from the path query itself
-def enable_on(edge: Property, filter_attrs: List[Tuple[Type, Set[str]]]):  # noqa: F821
-    edge._prop.parents.append(Builtins.PQEdgeAnnotation)
-    edge._prop.parents.append(Builtins.NoInlineAnnotation)
-    for type, attrs in filter_attrs:
-        type._type.parents.append(Builtins.PQFilterAnnotation)
-        type._type.parents.append(Builtins.NoInlineAnnotation)
-        for attr_name in attrs:
-            prop = getattr(type, attr_name)
-            prop._prop.parents.append(Builtins.PQFilterAnnotation)
-            prop._prop.parents.append(Builtins.NoInlineAnnotation)
-
-# main entry point
-def path(
-    segments: List[Union[Instance,VarLength]],
-    backwards_edge: str,
-    group: PathGroup = PathGroup.FOREACH,
-    selection: PathSelection = PathSelection.ALL,
-    strategy: Strategy = Strategy.FROM_SOURCE,
-):
-    """
-    Find a path consisting of the provided segments, linked by
-    the provided `backwards_edge` property.
-    """
-
-    model = get_graph()
-
-    query = PathQuery(model, segments, backwards_edge, group, selection, strategy)
-
-    # declare filter relations
-    conn_relation_name = query._match()
-
-    # invoke filter
-    a, b = create_vars(2)
-    getattr(rel, conn_relation_name)(a, b)
-
-    # return a special Producer representing the path itself, which
-    # lazily invokes Pathfinder
-    return PathInstance(query)
-
-def _get_pathfinder_source():
-    # check if the current dir is in a zip file
-    current_dir = os.path.dirname(__file__)
-    pathfinder_path = os.path.join(current_dir, "pathfinder.rel")
-    if os.path.exists(pathfinder_path):
-        return open(pathfinder_path).read()
-    # if we're in a zip file, read it within that
-    zip_split = current_dir.split(".zip")
-    zip_path = zip_split[0] + ".zip"
-    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
-        with zip_ref.open('relationalai/experimental/paths/pathfinder.rel') as file:
-            return file.read().decode("utf-8")
-
-def _install_pathfinder(model: Graph):
-    source = _get_pathfinder_source()
-    model.install_raw(source, name="pathfinder", overwrite=True)
+## pathfinder module
+
+from .api import node, edge, path, optional, plus, star, union, match
+
+__all__ = [
+    'node',
+    'edge',
+    'path',
+    'optional',
+    'plus',
+    'star',
+    'union',
+    'match',
+]