relationalai 0.12.1__py3-none-any.whl → 0.12.2__py3-none-any.whl

relationalai/clients/direct_access_client.py

@@ -49,6 +49,11 @@ class DirectAccessClient:
             "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"),
+            "get_job": Endpoint(method="GET", endpoint="/v1alpha1/jobs/{job_type}/{job_id}"),
+            "list_jobs": Endpoint(method="GET", endpoint="/v1alpha1/jobs"),
+            "get_job_events": Endpoint(method="GET", endpoint="/v1alpha1/jobs/{job_type}/{job_id}/events/{stream_name}"),
+            "create_job": Endpoint(method="POST", endpoint="/v1alpha1/jobs"),
+            "cancel_job": Endpoint(method="POST", endpoint="/v1alpha1/jobs/{job_type}/{job_id}/cancel"),
         }
         self.http_session = self._create_retry_session()
 

relationalai/clients/snowflake.py

@@ -761,13 +761,13 @@ Otherwise, remove it from your '{profile}' configuration profile.
         with debugging.span("create_model", name=name):
             self._exec(f"call {APP_NAME}.api.create_database('{name}', false, {debugging.gen_current_propagation_headers()});")
 
-    def delete_graph(self, name:str, force=False):
+    def delete_graph(self, name:str, force=False, language:str="rel"):
         prop_hdrs = debugging.gen_current_propagation_headers()
         if self.config.get("use_graph_index", USE_GRAPH_INDEX):
             keep_database = not force and self.config.get("reuse_model", True)
-            with debugging.span("release_index", name=name, keep_database=keep_database):
+            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}));")
+                response = self._exec(f"call {APP_NAME}.api.release_index('{name}', OBJECT_CONSTRUCT('keep_database', {keep_database}, 'language', '{language}'));")
                 if response:
                     result = next(iter(response))
                     obj = json.loads(result["RELEASE_INDEX"])
@@ -795,6 +795,7 @@ Otherwise, remove it from your '{profile}' configuration profile.
         model: str,
         engine_name: str,
         engine_size: str | None = None,
+        language: str = "rel",
         program_span_id: str | None = None,
         headers: Dict | None = None,
     ):
@@ -805,6 +806,7 @@ Otherwise, remove it from your '{profile}' configuration profile.
             model,
             engine_name,
             engine_size,
+            language,
             program_span_id,
             headers,
             self.generation
@@ -2970,7 +2972,16 @@ 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(app_name, self.resources.sources, model, engine_name, engine_size, program_span_id, headers=headers)
+            self.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
+            )
 
         self.last_database_version = len(self.resources.sources)
         self._manage_packages()
@@ -2989,12 +3000,20 @@ class SnowflakeClient(Client):
         model: str,
         engine_name: str,
         engine_size: str | None = None,
+        language: str = "rel",
         program_span_id: str | None = None,
         headers: Dict | None = None,
     ):
         assert isinstance(self.resources, Resources)
         return self.resources.poll_use_index(
-            app_name, sources, model, engine_name, engine_size, program_span_id, headers=headers
+            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
         )
 
 
@@ -3302,6 +3321,7 @@ class DirectAccessResources(Resources):
         model: str,
         engine_name: str,
         engine_size: str = "",
+        language: str = "rel",
         rai_relations: List[str] | None = None,
         pyrel_program_id: str | None  = None,
         skip_pull_relations: bool = False,
@@ -3317,6 +3337,7 @@ class DirectAccessResources(Resources):
             payload = {
                 "model_name": model,
                 "caller_engine_name": engine_name,
+                "language": language,
                 "pyrel_program_id": pyrel_program_id,
                 "skip_pull_relations": skip_pull_relations,
                 "rai_relations": rai_relations or [],
@@ -3342,6 +3363,7 @@ class DirectAccessResources(Resources):
         model: str,
         engine_name: str,
         engine_size: str | None = None,
+        language: str = "rel",
         program_span_id: str | None = None,
         headers: Dict | None = None,
     ):
@@ -3352,6 +3374,7 @@ class DirectAccessResources(Resources):
             model=model,
             engine_name=engine_name,
             engine_size=engine_size,
+            language=language,
             program_span_id=program_span_id,
             headers=headers,
             generation=self.generation,
@@ -3492,14 +3515,14 @@ class DirectAccessResources(Resources):
         with debugging.span("create_model", dbname=name):
             return self._create_database(name,"")
 
-    def delete_graph(self, name:str, force=False):
+    def delete_graph(self, name:str, force=False, language: str = "rel"):
         prop_hdrs = debugging.gen_current_propagation_headers()
         if self.config.get("use_graph_index", USE_GRAPH_INDEX):
             keep_database = not force and self.config.get("reuse_model", True)
-            with debugging.span("release_index", name=name, keep_database=keep_database):
+            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},
+                    payload={"model_name": name, "keep_database": keep_database, "language": language},
                     headers=prop_hdrs,
                 )
                 if (

relationalai/clients/use_index_poller.py

@@ -161,8 +161,9 @@ class UseIndexPoller:
         model: str,
         engine_name: str,
         engine_size: Optional[str],
-        program_span_id: Optional[str],
-        headers: Optional[Dict],
+        language: str = "rel",
+        program_span_id: Optional[str] = None,
+        headers: Optional[Dict] = None,
         generation: Optional[Generation] = None,
     ):
         self.res = resource
@@ -171,6 +172,7 @@ class UseIndexPoller:
         self.model = model
         self.engine_name = engine_name
         self.engine_size = engine_size or self.res.config.get_default_engine_size()
+        self.language = language
         self.program_span_id = program_span_id
         self.headers = headers or {}
         self.counter = 1
@@ -190,8 +192,8 @@ class UseIndexPoller:
         )
         current_user = self.res.get_sf_session().get_current_user()
         assert current_user is not None, "current_user must be set"
-        data_freshness = self.res.config.get_data_freshness_mins()
-        self.cache = GraphIndexCache(current_user, model, data_freshness, self.sources)
+        self.data_freshness = self.res.config.get_data_freshness_mins()
+        self.cache = GraphIndexCache(current_user, model, self.data_freshness, self.sources)
         self.sources = self.cache.choose_sources()
         # execution_id is allowed to group use_index call, which belongs to the same loop iteration
         self.execution_id = str(uuid.uuid4())
@@ -486,6 +488,8 @@ class UseIndexPoller:
                 "wait_for_stream_sync": self.wait_for_stream_sync,
                 "should_check_cdc": self.should_check_cdc,
                 "init_engine_async": self.init_engine_async,
+                "language": self.language,
+                "data_freshness_mins": self.data_freshness,
             })
 
             request_headers = debugging.add_current_propagation_headers(self.headers)
@@ -853,11 +857,23 @@ class DirectUseIndexPoller(UseIndexPoller):
         model: str,
         engine_name: str,
         engine_size: Optional[str],
-        program_span_id: Optional[str],
-        headers: Optional[Dict],
+        language: str = "rel",
+        program_span_id: Optional[str] = None,
+        headers: Optional[Dict] = None,
         generation: Optional[Generation] = None,
     ):
-        super().__init__(resource, app_name, sources, model, engine_name, engine_size, program_span_id, headers, generation)
+        super().__init__(
+            resource=resource,
+            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,
+            generation=generation,
+        )
         from relationalai.clients.snowflake import DirectAccessResources
         self.res: DirectAccessResources = cast(DirectAccessResources, self.res)
 
@@ -878,6 +894,7 @@ class DirectUseIndexPoller(UseIndexPoller):
                         model=self.model,
                         engine_name=self.engine_name,
                         engine_size=self.engine_size,
+                        language=self.language,
                         rai_relations=[],
                         pyrel_program_id=self.program_span_id,
                         skip_pull_relations=True,

relationalai/experimental/solvers.py

@@ -9,11 +9,13 @@ from ..metamodel import Builtins
 from ..tools.cli_controls import Spinner
 from ..tools.constants import DEFAULT_QUERY_TIMEOUT_MINS
 from .. import debugging
+from .. errors import ResponseStatusException
 import uuid
 import relationalai
 import json
 from ..clients.util import poll_with_specified_overhead
 from ..clients.snowflake import Resources as SnowflakeResources
+from ..clients.snowflake import DirectAccessClient, DirectAccessResources
 from ..util.timeout import calc_remaining_timeout_minutes
 
 rel_sv = rel._tagged(Builtins.SingleValued)
@@ -23,7 +25,8 @@ APP_NAME = relationalai.clients.snowflake.APP_NAME
 ENGINE_TYPE_SOLVER = "SOLVER"
 # TODO (dba) The ERP still uses `worker` instead of `engine`. Change
 # this once we fix this in the ERP.
-ENGINE_ERRORS = ["worker is suspended", "create/resume", "worker not found", "no workers found", "worker was deleted"]
+WORKER_ERRORS = ["worker is suspended", "create/resume", "worker not found", "no workers found", "worker was deleted"]
+ENGINE_ERRORS = ["engine is suspended", "create/resume", "engine not found", "no engines found", "engine was deleted"]
 ENGINE_NOT_READY_MSGS = ["worker is in pending", "worker is provisioning", "worker is not ready to accept jobs"]
 
 # --------------------------------------------------
@@ -213,13 +216,6 @@ class SolverModel:
         config_file_path = getattr(rai_config, 'file_path', None)
         start_time = time.monotonic()
         remaining_timeout_minutes = query_timeout_mins
-        # 1. Materialize the model and store it.
-        # TODO(coey) Currently we must run a dummy query to install the pyrel rules in a separate txn
-        # to the solve_output updates. Ideally pyrel would offer an option to flush the rules separately.
-        self.graph.exec_raw("", query_timeout_mins=remaining_timeout_minutes)
-        remaining_timeout_minutes = calc_remaining_timeout_minutes(
-            start_time, query_timeout_mins, config_file_path=config_file_path,
-        )
         response = self.graph.exec_raw(
             textwrap.dedent(f"""
             @inline
@@ -269,7 +265,9 @@ class SolverModel:
             job_id = solver._exec_job(payload, log_to_console=log_to_console, query_timeout_mins=remaining_timeout_minutes)
         except Exception as e:
             err_message = str(e).lower()
-            if any(kw in err_message.lower() for kw in ENGINE_ERRORS + ENGINE_NOT_READY_MSGS):
+            if isinstance(e, ResponseStatusException):
+                err_message = e.response.json().get("message", "")
+            if any(kw in err_message.lower() for kw in ENGINE_ERRORS + WORKER_ERRORS + ENGINE_NOT_READY_MSGS):
                 solver._auto_create_solver_async()
                 remaining_timeout_minutes = calc_remaining_timeout_minutes(
                     start_time, query_timeout_mins, config_file_path=config_file_path
@@ -553,7 +551,11 @@ class Solver:
         # may configure each individual solver.
         self.engine_settings = settings
 
-        return self._auto_create_solver_async()
+        # Optimistically set the engine object to a `READY` engine to
+        # avoid checking the engine status on each execution.
+        self.engine:Optional[dict[str,Any]] = {"name": engine_name, "state": "READY"}
+
+        return None
 
     # --------------------------------------------------
     # Helper
@@ -572,6 +574,7 @@ class Solver:
             assert len(engines) == 1 or len(engines) == 0
             if len(engines) != 0:
                 engine = engines[0]
+
             if engine:
                 # TODO (dba) Logic engines support altering the
                 # auto_suspend_mins setting. Currently, we don't have
@@ -653,31 +656,20 @@ class Solver:
 
             self.engine = engine
 
-    def _exec_job_async(self, payload, query_timeout_mins: Optional[int]=None):
-        payload_json = json.dumps(payload)
-        engine_name = self.engine["name"]
-        if query_timeout_mins is None and (timeout_value := self.rai_config.get("query_timeout_mins", DEFAULT_QUERY_TIMEOUT_MINS)) is not None:
-            query_timeout_mins = int(timeout_value)
-        if query_timeout_mins is not None:
-            sql_string = textwrap.dedent(f"""
-            CALL {APP_NAME}.experimental.exec_job_async('{ENGINE_TYPE_SOLVER}', '{engine_name}', '{payload_json}', null, {query_timeout_mins})
-            """)
-        else:
-            sql_string = textwrap.dedent(f"""
-            CALL {APP_NAME}.experimental.exec_job_async('{ENGINE_TYPE_SOLVER}', '{engine_name}', '{payload_json}')
-            """)
-        res = self.provider.resources._exec(sql_string)
-        return res[0]["ID"]
-
     def _exec_job(self, payload, log_to_console=True, query_timeout_mins: Optional[int]=None):
+        if self.engine is None:
+            raise Exception("Engine not initialized.")
+
         # Make sure the engine is ready.
         if self.engine["state"] != "READY":
             poll_with_specified_overhead(lambda: self._is_solver_ready(), 0.1)
 
         with debugging.span("job") as job_span:
-            job_id = self._exec_job_async(payload, query_timeout_mins=query_timeout_mins)
+            job_id = self.provider.create_job_async(self.engine["name"], payload, query_timeout_mins=query_timeout_mins)
             job_span["job_id"] = job_id
             debugging.event("job_created", job_span, job_id=job_id, engine_name=self.engine["name"], job_type=ENGINE_TYPE_SOLVER)
+            if not isinstance(job_id, str):
+                job_id = ""
             polling_state = PollingState(job_id, "", False, log_to_console)
 
             try:
@@ -693,7 +685,14 @@ class Solver:
             return job_id
 
     def _is_solver_ready(self):
+        if self.engine is None:
+            raise Exception("Engine not initialized.")
+
         result = self.provider.get_solver(self.engine["name"])
+
+        if result is None:
+            raise Exception("No engine available.")
+            
         self.engine = result
         state = result["state"]
         if state != "READY" and state != "PENDING":
@@ -711,20 +710,11 @@ class Solver:
 
         return status == "COMPLETED" or status == "FAILED" or status == "CANCELED"
 
-    def _get_job_events(self, job_id: str, continuation_token: str = ""):
-        results = self.provider.resources._exec(
-            f"SELECT {APP_NAME}.experimental.get_job_events('{ENGINE_TYPE_SOLVER}', '{job_id}', '{continuation_token}');"
-        )
-        if not results:
-            return {"events": [], "continuation_token": None}
-        row = results[0][0]
-        return json.loads(row)
-
     def _print_solver_logs(self, state: PollingState):
         if state.is_done:
             return
 
-        resp = self._get_job_events(state.job_id, state.continuation_token)
+        resp = self.provider.get_job_events(state.job_id, state.continuation_token)
 
         # Print solver logs to stdout.
         for event in resp["events"]:
@@ -754,7 +744,12 @@ class Provider:
             resources = relationalai.Resources()
         if not isinstance(resources, relationalai.clients.snowflake.Resources):
             raise Exception("Solvers are only supported on SPCS.")
+
         self.resources = resources
+        self.direct_access_client: Optional[DirectAccessClient] = None
+
+        if isinstance(self.resources, DirectAccessResources):
+            self.direct_access_client = self.resources.direct_access_client
 
     def create_solver(
         self,
@@ -770,75 +765,285 @@ class Provider:
         engine_config: dict[str, Any] = {"settings": settings}
         if auto_suspend_mins is not None:
             engine_config["auto_suspend_mins"] = auto_suspend_mins
-        self.resources._exec(
-            f"CALL {APP_NAME}.experimental.create_engine('{ENGINE_TYPE_SOLVER}', '{name}', '{size}', {engine_config});"
+        self.resources._exec_sql(
+            f"CALL {APP_NAME}.experimental.create_engine('{ENGINE_TYPE_SOLVER}', '{name}', '{size}', {engine_config});", None
         )
 
     def create_solver_async(
         self,
         name: str,
         size: str | None = None,
-        settings: dict = {},
+        settings: dict | None = None,
         auto_suspend_mins: int | None = None,
     ):
         if size is None:
             size = "HIGHMEM_X64_S"
-        if settings is None:
-            settings = ""
-        engine_config: dict[str, Any] = {"settings": settings}
-        if auto_suspend_mins is not None:
-            engine_config["auto_suspend_mins"] = auto_suspend_mins
-        self.resources._exec(
-            f"CALL {APP_NAME}.experimental.create_engine_async('{ENGINE_TYPE_SOLVER}', '{name}', '{size}', {engine_config});"
-        )
+
+        if self.direct_access_client is not None:
+            payload:dict[str, Any] = {
+                "name": name,
+                "settings": settings,
+               }
+            if auto_suspend_mins is not None:
+                payload["auto_suspend_mins"] = auto_suspend_mins
+            if size is not None:
+                payload["size"] = size
+            response = self.direct_access_client.request(
+                "create_engine",
+                payload=payload,
+                path_params={"engine_type": "solver"},
+               )
+            if response.status_code != 200:
+                raise ResponseStatusException(
+                    f"Failed to create engine {name} with size {size}.", response
+                )
+        else:
+            engine_config: dict[str, Any] = {}
+            if settings is not None:
+                engine_config["settings"] = settings
+            if auto_suspend_mins is not None:
+                engine_config["auto_suspend_mins"] = auto_suspend_mins
+            self.resources._exec_sql(
+                f"CALL {APP_NAME}.experimental.create_engine_async('{ENGINE_TYPE_SOLVER}', '{name}', '{size}', {engine_config});",
+                None
+               )
 
     def delete_solver(self, name: str):
-        self.resources._exec(
-            f"CALL {APP_NAME}.experimental.delete_engine('{ENGINE_TYPE_SOLVER}', '{name}');"
-        )
+        if self.direct_access_client is not None:
+            response = self.direct_access_client.request(
+                "delete_engine", path_params = {"engine_type": ENGINE_TYPE_SOLVER, "engine_name": name}
+            )
+            if response.status_code != 200:
+                raise ResponseStatusException("Failed to delete engine.", response)
+            return None
+        else:
+            self.resources._exec_sql(
+                f"CALL {APP_NAME}.experimental.delete_engine('{ENGINE_TYPE_SOLVER}', '{name}');",
+                None
+            )
 
     def resume_solver_async(self, name: str):
-        self.resources._exec(
-            f"CALL {APP_NAME}.experimental.resume_engine_async('{ENGINE_TYPE_SOLVER}', '{name}');"
-        )
+        if self.direct_access_client is not None:
+            response = self.direct_access_client.request(
+                "resume_engine", path_params = {"engine_type": ENGINE_TYPE_SOLVER, "engine_name": name}
+            )
+            if response.status_code != 200:
+                raise ResponseStatusException("Failed to resume engine.", response)
+            return None
+        else:
+            self.resources._exec_sql(
+                f"CALL {APP_NAME}.experimental.resume_engine_async('{ENGINE_TYPE_SOLVER}', '{name}');",
+                None
+            )
+            return None
 
     def get_solver(self, name: str):
-        results = self.resources._exec(
-            f"CALL {APP_NAME}.experimental.get_engine('{ENGINE_TYPE_SOLVER}', '{name}');"
-        )
-        return solver_list_to_dicts(results)[0]
+        if self.direct_access_client is not None:
+            response = self.direct_access_client.request(
+                "get_engine", path_params = {"engine_type": ENGINE_TYPE_SOLVER, "engine_name": name}
+            )
+            if response.status_code != 200:
+                raise ResponseStatusException("Failed to get engine.", response)
+            solver = response.json()
+            if not solver :
+                return None
+            solver_state = {
+                "name": solver["name"],
+                "id": solver["id"],
+                "size": solver["size"],
+                "state": solver["status"], # callers are expecting 'state'
+                "created_by": solver["created_by"],
+                "created_on": solver["created_on"],
+                "updated_on": solver["updated_on"],
+                "version": solver["version"],
+                "auto_suspend": solver["auto_suspend_mins"],
+                "suspends_at": solver["suspends_at"],
+                "solvers": []
+                if solver["settings"] == ""
+                else [
+                        k
+                        for (k,v) in json.loads(solver["settings"]).items()
+                        if isinstance(v, dict) and v.get("enabled", False)
+                ],
+                }
+            return solver_state
+        else:
+            results = self.resources._exec_sql(
+                f"CALL {APP_NAME}.experimental.get_engine('{ENGINE_TYPE_SOLVER}', '{name}');",
+                None
+               )
+            return solver_list_to_dicts(results)[0]
 
     def list_solvers(self, state: str | None = None):
-        where_clause = f"WHERE TYPE='{ENGINE_TYPE_SOLVER}'"
-        where_clause = (
-            f"{where_clause} AND STATUS = '{state.upper()}'" if state else where_clause
-        )
-        statement = f"SELECT NAME,ID,SIZE,STATUS,CREATED_BY,CREATED_ON,UPDATED_ON,AUTO_SUSPEND_MINS,SETTINGS FROM {APP_NAME}.experimental.engines {where_clause};"
-        results = self.resources._exec(statement)
-        return solver_list_to_dicts(results)
+        if self.direct_access_client is not None:
+            response = self.direct_access_client.request(
+                "list_engines"
+            )
+            if response.status_code != 200:
+                raise ResponseStatusException("Failed to list engines.", response)
+            response_content = response.json()
+            if not response_content:
+                return []
+            engines = [
+                {
+                    "name": engine["name"],
+                    "id": engine["id"],
+                    "size": engine["size"],
+                    "state": engine["status"], # callers are expecting 'state'
+                    "created_by": engine["created_by"],
+                    "created_on": engine["created_on"],
+                    "updated_on": engine["updated_on"],
+                    "auto_suspend_mins": engine["auto_suspend_mins"],
+                    "solvers": []
+                    if engine["settings"] == ""
+                    else [
+                        k
+                        for (k, v) in json.loads(engine["settings"]).items()
+                        if isinstance(v, dict) and v.get("enabled", False)
+                           ],
+                }
+                for engine in response_content.get("engines", [])
+                if (state is None or engine.get("status") == state) and (engine.get("type") == ENGINE_TYPE_SOLVER)
+               ]
+            return sorted(engines, key=lambda x: x["name"])
+        else:
+            where_clause = f"WHERE TYPE='{ENGINE_TYPE_SOLVER}'"
+            where_clause = (
+                f"{where_clause} AND STATUS = '{state.upper()}'" if state else where_clause
+               )
+            statement = f"SELECT NAME,ID,SIZE,STATUS,CREATED_BY,CREATED_ON,UPDATED_ON,AUTO_SUSPEND_MINS,SETTINGS FROM {APP_NAME}.experimental.engines {where_clause};"
+            results = self.resources._exec_sql(statement, None)
+            return solver_list_to_dicts(results)
 
     # --------------------------------------------------
     # Job API
     # --------------------------------------------------
 
+    def create_job_async(self, engine_name, payload, query_timeout_mins: Optional[int]=None):
+        payload_json = json.dumps(payload)
+
+        if query_timeout_mins is None and (timeout_value := self.resources.config.get("query_timeout_mins", DEFAULT_QUERY_TIMEOUT_MINS)) is not None:
+            query_timeout_mins = int(timeout_value)
+            
+        if self.direct_access_client is not None:
+            job = {
+                "job_type":ENGINE_TYPE_SOLVER,
+                "worker_name": engine_name,
+                "timeout_mins": query_timeout_mins,
+                "payload": payload_json,
+            }
+            response = self.direct_access_client.request(
+                "create_job",
+                payload=job,
+            )
+            if response.status_code != 200:
+                raise ResponseStatusException("Failed to create job.", response)
+            response_content = response.json()
+            return response_content["id"]
+        else:
+            if query_timeout_mins is not None:
+                sql_string = textwrap.dedent(f"""
+                CALL {APP_NAME}.experimental.exec_job_async('{ENGINE_TYPE_SOLVER}', '{engine_name}', '{payload_json}', null, {query_timeout_mins})
+                """)
+            else:
+                sql_string = textwrap.dedent(f"""
+                CALL {APP_NAME}.experimental.exec_job_async('{ENGINE_TYPE_SOLVER}', '{engine_name}', '{payload_json}')
+                """)
+            res = self.resources._exec_sql(sql_string, None)
+            return res[0]["ID"]
+
     def list_jobs(self, state=None, limit=None):
-        state_clause = f"AND STATE = '{state.upper()}'" if state else ""
-        limit_clause = f"LIMIT {limit}" if limit else ""
-        results = self.resources._exec(
-            f"SELECT ID,STATE,CREATED_BY,CREATED_ON,FINISHED_AT,DURATION,PAYLOAD,ENGINE_NAME FROM {APP_NAME}.experimental.jobs where type='{ENGINE_TYPE_SOLVER}' {state_clause} ORDER BY created_on DESC {limit_clause};"
-        )
-        return job_list_to_dicts(results)
+        if self.direct_access_client is not None:
+            response = self.direct_access_client.request(
+                "list_jobs"
+            )
+            if response.status_code != 200:
+                raise ResponseStatusException("Failed to list jobs.", response)
+            response_content = response.json()
+            if not response_content:
+                return []
+            jobs = [
+                {
+                    "id": job["id"],
+                    "state": job["state"],
+                    "created_by": job["created_by"],
+                    "created_on": job["created_on"],
+                    "finished_at": job.get("finished_at", None),
+                    "duration": job["duration"] if "duration" in job else 0,
+                    "solver": json.loads(job["payload"]).get("solver", ""),
+                    "engine": job.get("engine_name", job["worker_name"]),
+                }
+                for job in response_content.get("jobs", [])
+                if state is None or job.get("state") == state
+               ]
+            return sorted(jobs, key=lambda x: x["created_on"], reverse=True)
+        else:
+            state_clause = f"AND STATE = '{state.upper()}'" if state else ""
+            limit_clause = f"LIMIT {limit}" if limit else ""
+            results = self.resources._exec_sql(
+                f"SELECT ID,STATE,CREATED_BY,CREATED_ON,FINISHED_AT,DURATION,PAYLOAD,ENGINE_NAME FROM {APP_NAME}.experimental.jobs where type='{ENGINE_TYPE_SOLVER}' {state_clause} ORDER BY created_on DESC {limit_clause};",
+                None
+               )
+            return job_list_to_dicts(results)
 
     def get_job(self, id: str):
-        results = self.resources._exec(
-            f"CALL {APP_NAME}.experimental.get_job('{ENGINE_TYPE_SOLVER}', '{id}');"
-        )
-        return job_list_to_dicts(results)[0]
+        if self.direct_access_client is not None:
+            response = self.direct_access_client.request(
+                "get_job", path_params = {"job_type": ENGINE_TYPE_SOLVER, "job_id": id}
+            )
+            if response.status_code != 200:
+                raise ResponseStatusException("Failed to get job.", response)
+            response_content = response.json()
+            return response_content["job"]
+        else:
+            results = self.resources._exec_sql(
+                f"CALL {APP_NAME}.experimental.get_job('{ENGINE_TYPE_SOLVER}', '{id}');",
+                None
+            )
+            return job_list_to_dicts(results)[0]
+
+    def get_job_events(self, job_id: str, continuation_token: str = ""):
+        if self.direct_access_client is not None:
+            response = self.direct_access_client.request(
+                "get_job_events",
+                path_params = {"job_type": ENGINE_TYPE_SOLVER, "job_id": job_id, "stream_name": "progress"},
+                query_params={"continuation_token": continuation_token},
+            )
+            if response.status_code != 200:
+                raise ResponseStatusException("Failed to get job events.", response)
+            response_content = response.json()
+            if not response_content:
+                return {
+                    "events": [],
+                    "continuation_token": None
+                }
+            return response_content
+        else:
+            results = self.resources._exec_sql(
+                f"SELECT {APP_NAME}.experimental.get_job_events('{ENGINE_TYPE_SOLVER}', '{job_id}', '{continuation_token}');",
+                None
+               )
+            if not results:
+                return {"events": [], "continuation_token": None}
+            row = results[0][0]
+            if not isinstance(row, str):
+                row = ""
+            return json.loads(row)
 
     def cancel_job(self, id: str):
-        self.resources._exec(
-            f"CALL {APP_NAME}.experimental.cancel_job('{ENGINE_TYPE_SOLVER}', '{id}');"
-        )
+        if self.direct_access_client is not None:
+            response = self.direct_access_client.request(
+                "cancel_job", path_params = {"job_type": ENGINE_TYPE_SOLVER, "job_id": id}
+            )
+            if response.status_code != 200:
+                raise ResponseStatusException("Failed to cancel job.", response)
+            return None
+        else:
+            self.resources._exec_sql(
+                f"CALL {APP_NAME}.experimental.cancel_job('{ENGINE_TYPE_SOLVER}', '{id}');",
+                None
+            )
+            return None
 
 
 def solver_list_to_dicts(results):
@@ -865,7 +1070,6 @@ def solver_list_to_dicts(results):
         for row in results
     ]
 
-
 def job_list_to_dicts(results):
     if not results:
         return []

relationalai/semantics/internal/internal.py

@@ -346,7 +346,7 @@ def find_select_keys(item: Any, keys:OrderedSet[Key]|None = None, enable_primiti
 
     if isinstance(item, (list, tuple)):
         for it in item:
-            find_select_keys(it, keys)
+            find_select_keys(it, keys, enable_primitive_key=enable_primitive_key)
 
     elif isinstance(item, (Relationship, RelationshipReading)) and item._parent:
         find_select_keys(item._parent, keys)
@@ -390,7 +390,7 @@ def find_select_keys(item: Any, keys:OrderedSet[Key]|None = None, enable_primiti
         find_select_keys(item._arg, keys)
 
     elif isinstance(item, Alias):
-        find_select_keys(item._thing, keys)
+        find_select_keys(item._thing, keys, enable_primitive_key=enable_primitive_key)
 
     elif isinstance(item, Aggregate):
         keys.update( Key(k, True) for k in item._group )
@@ -2418,21 +2418,21 @@ class Fragment():
 
     def meta(self, **kwargs: Any) -> Fragment:
         """Add metadata to the query.
-        
+
         Metadata can be used for debugging and observability purposes.
-        
+
         Args:
             **kwargs: Metadata key-value pairs
-        
+
         Returns:
             Fragment: Returns self for method chaining
-            
+
         Example:
             select(Person.name).meta(workload_name="test", priority=1, enabled=True)
         """
         if not kwargs:
             raise ValueError("meta() requires at least one argument")
-            
+
         self._meta.update(kwargs)
         return self
 
@@ -2560,7 +2560,7 @@ class Fragment():
         with debugging.span("query", dsl=str(clone), **with_source(clone), meta=clone._meta):
             query_task = qb_model._compiler.fragment(clone)
             qb_model._to_executor().execute(ir_model, query_task, result_cols=result_cols, export_to=table._fqn, update=update, meta=clone._meta)
-    
+
 
 #--------------------------------------------------
 # Select / Where

relationalai/semantics/lqp/executor.py

@@ -60,7 +60,7 @@ class LQPExecutor(e.Executor):
                 if not self.dry_run:
                     self.engine = self._resources.get_default_engine_name()
                     if not self.keep_model:
-                        atexit.register(self._resources.delete_graph, self.database, True)
+                        atexit.register(self._resources.delete_graph, self.database, True, "lqp")
         return self._resources
 
     # Checks the graph index and updates it if necessary
@@ -88,7 +88,15 @@ class LQPExecutor(e.Executor):
         assert self.engine is not None
 
         with debugging.span("poll_use_index", sources=sources, model=model, engine=engine_name):
-            resources.poll_use_index(app_name, sources, model, self.engine, engine_size, program_span_id)
+            resources.poll_use_index(
+                app_name=app_name,
+                sources=sources,
+                model=model,
+                engine_name=self.engine,
+                engine_size=engine_size,
+                language="lqp",
+                program_span_id=program_span_id,
+            )
 
     def report_errors(self, problems: list[dict[str, Any]], abort_on_error=True):
         from relationalai import errors

relationalai/semantics/lqp/model2lqp.py

@@ -580,8 +580,9 @@ def get_relation_id(ctx: TranslationCtx, relation: ir.Relation, projection: list
     if relation.id in ctx.def_names.id_to_name:
         unique_name = ctx.def_names.id_to_name[relation.id]
     else:
-        prefix = helpers.relation_name_prefix(relation)
-        unique_name = ctx.def_names.get_name_by_id(relation.id, prefix + relation.name)
+        name = helpers.relation_name_prefix(relation) + relation.name
+        name = helpers.sanitize(name)
+        unique_name = ctx.def_names.get_name_by_id(relation.id, name)
 
     return utils.gen_rel_id(ctx, unique_name, types)
 

relationalai/semantics/rel/executor.py

@@ -58,7 +58,7 @@ class RelExecutor(e.Executor):
                 if not self.dry_run:
                     self.engine = self._resources.get_default_engine_name()
                     if not self.keep_model:
-                        atexit.register(self._resources.delete_graph, self.database, True)
+                        atexit.register(self._resources.delete_graph, self.database, True, "rel")
         return self._resources
 
     def check_graph_index(self, headers: dict[str, Any] | None = None):
@@ -85,7 +85,16 @@ class RelExecutor(e.Executor):
         assert self.engine is not None
 
         with debugging.span("poll_use_index", sources=sources, model=model, engine=engine_name):
-            resources.poll_use_index(app_name, sources, model, self.engine, engine_size, program_span_id, headers)
+            resources.poll_use_index(
+                app_name=app_name,
+                sources=sources,
+                model=model,
+                engine_name=self.engine,
+                engine_size=engine_size,
+                language="rel",
+                program_span_id=program_span_id,
+                headers=headers,
+            )
 
     def report_errors(self, problems: list[dict[str, Any]], abort_on_error=True):
         from relationalai import errors

{relationalai-0.12.1.dist-info → relationalai-0.12.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: relationalai
-Version: 0.12.1
+Version: 0.12.2
 Summary: RelationalAI Library and CLI
 Author-email: RelationalAI <support@relational.ai>
 License-File: LICENSE

{relationalai-0.12.1.dist-info → relationalai-0.12.2.dist-info}/RECORD

@@ -23,14 +23,14 @@ relationalai/clients/azure.py,sha256=6tYHxKVN2fHMIDV7J_EHZD9WntkYh2IreMRMqlq1Bhg
 relationalai/clients/cache_store.py,sha256=A-qd11wcwN3TkIqvlN0_iFUU3aEjJal3T2pqFBwkkzQ,3966
 relationalai/clients/client.py,sha256=4SSunUwuFEcRFXOPYotpSLDPr0CKuwJ4335W0DesR90,35792
 relationalai/clients/config.py,sha256=hERaKjc3l4kd-kf0l-NUOHrWunCn8gmFWpuE0j3ScJg,24457
-relationalai/clients/direct_access_client.py,sha256=uDML4nFTctvvya4GFWoef2oirlWHiRugn5qBk9fxH9k,5961
+relationalai/clients/direct_access_client.py,sha256=VGjQ7wzduxCo04BkxSZjlPAgqK-aBc32zIXcMfAzzSU,6436
 relationalai/clients/export_procedure.py.jinja,sha256=nhvVcs5hQyWExFDuROQbi1VyYzOCa_ZIRPR2KzZwDtI,10582
 relationalai/clients/hash_util.py,sha256=pZVR1FX3q4G_19p_r6wpIR2tIM8_WUlfAR7AVZJjIYM,1495
 relationalai/clients/profile_polling.py,sha256=pUH7WKH4nYDD0SlQtg3wsWdj0K7qt6nZqUw8jTthCBs,2565
 relationalai/clients/result_helpers.py,sha256=wDSD02Ngx6W-YQqBIGKnpXD4Ju3pA1e9Nz6ORRI6SRI,17808
-relationalai/clients/snowflake.py,sha256=Om3iVea8Ky_P0MMWBJf83xBpG1LMKXrqRRBCjTZ-OcY,158848
+relationalai/clients/snowflake.py,sha256=c5z7sh9wXeMaAm0JiyrOZlHJFK_hbS76ml8Ed6kWxKo,159579
 relationalai/clients/types.py,sha256=eNo6akcMTbnBFbBbHd5IgVeY-zuAgtXlOs8Bo1SWmVU,2890
-relationalai/clients/use_index_poller.py,sha256=J8AUIJE1fFZbB5I6Yrk2pQK9l7yN5SGBn2JQ5TmyqfM,45766
+relationalai/clients/use_index_poller.py,sha256=gxQDK-iaA9iz5-rkNWQjptZSxgdwv8WXIKC4HB44cHo,46291
 relationalai/clients/util.py,sha256=NJC8fnrWHR01NydwESPSetIHRWf7jQJURYpaWJjmDyE,12311
 relationalai/early_access/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 relationalai/early_access/builder/__init__.py,sha256=mrR-aGLPquUGc_e9-DOkVCCCo8QRE5A56GTTtgILNa4,993
@@ -226,7 +226,7 @@ relationalai/experimental/SF.py,sha256=cLIPCgTa6yCkmAU4zw3o2OlSM_Sl8ah3JUpHVoHzr
 relationalai/experimental/__init__.py,sha256=alaOJulJ9YHRl4WGyQ3Kzxhrpc2gP542oNheWXzs7NQ,74
 relationalai/experimental/graphs.py,sha256=vJQ_8ZF_Q2pL-w1xI4O22VGqbnTDimkkqTNqqIkGL7c,3521
 relationalai/experimental/inspect.py,sha256=jpk_xl_bWRcX96xVrHsex-YE3sTqa-W2b5IjGXoxeWk,1635
-relationalai/experimental/solvers.py,sha256=dfgFp_tCoeP6uYHlWzAmuFOtLI8sg4oRc_kM2KozooQ,34410
+relationalai/experimental/solvers.py,sha256=aD2O3ncmKVBgc4R6SKeLmee6FPEpC_2Xaxgc5Lh5t4s,43328
 relationalai/experimental/pathfinder/__init__.py,sha256=wzF7zzl5Cvp2CkFaPpY2VU-Tg7I03wvxGq7kp8FVvU0,5429
 relationalai/experimental/pathfinder/api.py,sha256=HT6PAt5gszgJhbnqsHACzZrr9ps9kwcHYu9Nf2EYtTA,4235
 relationalai/experimental/pathfinder/automaton.py,sha256=Z_TYjf5j89-48rxjxl0jV9Q0BR4VhV9cKRTFES7c1t0,21898
@@ -256,16 +256,16 @@ relationalai/semantics/devtools/compilation_manager.py,sha256=XBqG_nYWtK3s_J6MeC
 relationalai/semantics/devtools/extract_lqp.py,sha256=gxI3EvPUTPAkwgnkCKAkEm2vA6QkLfoM8AXXiVz0c34,3696
 relationalai/semantics/internal/__init__.py,sha256=JXrpFaL-fdZrvKpWTEn1UoLXITOoTGnAYwmgeiglhSk,774
 relationalai/semantics/internal/annotations.py,sha256=P_nRpBm4wLmE_8L0VO3TDORL1p2flXaPOVDC0GG7KsQ,306
-relationalai/semantics/internal/internal.py,sha256=yIrUgMqqjzDqB9abOkj7QVyF8SnJzf-tpd5UbjKlSlo,148938
+relationalai/semantics/internal/internal.py,sha256=Vh4LFatRA8xO6sIBxRSlbAv6HOMJPU65t2c_ZEh9lug,148972
 relationalai/semantics/internal/snowflake.py,sha256=DR6rgbPawen2ZOFegRuPkXt8aQUfciiCYzDI74u1Iwc,13036
 relationalai/semantics/lqp/__init__.py,sha256=XgcQZxK-zz_LqPDVtwREhsIvjTuUIt4BZhIedCeMY-s,48
 relationalai/semantics/lqp/builtins.py,sha256=bRmQ6fdceWU-4xf4l-W-YiuyDxJTPey1s6O4xlyW6iM,540
 relationalai/semantics/lqp/compiler.py,sha256=Nury1gPw_-Oi_mqT1-rhr13L4UmyIP2BGuotbuklQKA,949
 relationalai/semantics/lqp/constructors.py,sha256=8U4eUL8-m1wYRQnws_YWC1coGquTugVH5YC0Zek6VT8,2309
-relationalai/semantics/lqp/executor.py,sha256=0CYFKUyEj4zL-8BhRt_ru9b5hNle0f4PeLMh8AeK-Nw,19430
+relationalai/semantics/lqp/executor.py,sha256=j87LJFnzzNKHrazoYj0X1LK1jyy9ysDCKZ0ExhN1_xs,19643
 relationalai/semantics/lqp/intrinsics.py,sha256=Pb1mLIme7Q-5Y-CVacUOEvapfhKs076bgtRNi3f0ayY,833
 relationalai/semantics/lqp/ir.py,sha256=DUw0ltul0AS9CRjntNlmllWTwXpxMyYg4iJ9t7NFYMA,1791
-relationalai/semantics/lqp/model2lqp.py,sha256=F_4lwWICy1hvU-UhMbxFwEwPxPOj8IOPfr809ssgdwY,31719
+relationalai/semantics/lqp/model2lqp.py,sha256=M---RMVLEollI4aGHFh-2klR4y6C88vPKCHFd6tPNz8,31753
 relationalai/semantics/lqp/passes.py,sha256=nLppoHvIQkGP6VuG56OAZ1oOrYhEqpR_0w91gfJ7t_s,27540
 relationalai/semantics/lqp/pragmas.py,sha256=FzzldrJEAZ1AIcEw6D-FfaVg3CoahRYgPCFo7xHfg1g,375
 relationalai/semantics/lqp/primitives.py,sha256=Gbh6cohoAArhqEJTN_TgIRc7wmtdxXt231NRW0beEj0,10898
@@ -376,7 +376,7 @@ relationalai/semantics/reasoners/optimization/solvers_pb.py,sha256=ryNARpyph3uvr
 relationalai/semantics/rel/__init__.py,sha256=pMlVTC_TbQ45mP1LpzwFBBgPxpKc0H3uJDvvDXEWzvs,55
 relationalai/semantics/rel/builtins.py,sha256=qu4yZvLovn4Vn2x44D4XugqGD6Qo5xLxj_RKA34cpF4,1527
 relationalai/semantics/rel/compiler.py,sha256=hiLIaZzhVU5VMtU6rdo3tH2pmKMEtXvh6AHh94CvRAg,42203
-relationalai/semantics/rel/executor.py,sha256=83PoUBAHmQ4i4XrH7cTOEMY279yAhK1fVEq8o807XZ4,15646
+relationalai/semantics/rel/executor.py,sha256=mRBNyw6qHgmF3aOnM4aTIxBHCn4yaho4nQlMYM08pIA,15883
 relationalai/semantics/rel/rel.py,sha256=9I_V6dQ83QRaLzq04Tt-KjBWhmNxNO3tFzeornBK4zc,15738
 relationalai/semantics/rel/rel_utils.py,sha256=F14Ch8mn45J8SmM7HZnIHUNqDnb3WQLnkEGLme04iBk,9386
 relationalai/semantics/snowflake/__init__.py,sha256=BW_zvPQBWGTAtY6cluG6tDDG-QmU_jRb-F7PeCpDhIU,134
@@ -453,8 +453,8 @@ frontend/debugger/dist/index.html,sha256=0wIQ1Pm7BclVV1wna6Mj8OmgU73B9rSEGPVX-Wo
 frontend/debugger/dist/assets/favicon-Dy0ZgA6N.png,sha256=tPXOEhOrM4tJyZVJQVBO_yFgNAlgooY38ZsjyrFstgg,620
 frontend/debugger/dist/assets/index-Cssla-O7.js,sha256=MxgIGfdKQyBWgufck1xYggQNhW5nj6BPjCF6Wleo-f0,298886
 frontend/debugger/dist/assets/index-DlHsYx1V.css,sha256=21pZtAjKCcHLFjbjfBQTF6y7QmOic-4FYaKNmwdNZVE,60141
-relationalai-0.12.1.dist-info/METADATA,sha256=f-pkqK_i3tihiwfuimhr3XbrxE0YlrjW_VOOr6WaGok,2562
-relationalai-0.12.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-relationalai-0.12.1.dist-info/entry_points.txt,sha256=fo_oLFJih3PUgYuHXsk7RnCjBm9cqRNR--ab6DgI6-0,88
-relationalai-0.12.1.dist-info/licenses/LICENSE,sha256=pPyTVXFYhirkEW9VsnHIgUjT0Vg8_xsE6olrF5SIgpc,11343
-relationalai-0.12.1.dist-info/RECORD,,
+relationalai-0.12.2.dist-info/METADATA,sha256=XPjRexFpG92ZfObHk7CAwEtVSR7Epw2HkPnwnZACiFo,2562
+relationalai-0.12.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+relationalai-0.12.2.dist-info/entry_points.txt,sha256=fo_oLFJih3PUgYuHXsk7RnCjBm9cqRNR--ab6DgI6-0,88
+relationalai-0.12.2.dist-info/licenses/LICENSE,sha256=pPyTVXFYhirkEW9VsnHIgUjT0Vg8_xsE6olrF5SIgpc,11343
+relationalai-0.12.2.dist-info/RECORD,,