relationalai 0.12.1__py3-none-any.whl → 0.12.3__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,