relationalai 0.12.12__tar.gz → 0.12.13__tar.gz

{relationalai-0.12.12 → relationalai-0.12.13}/PKG-INFO

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

{relationalai-0.12.12 → relationalai-0.12.13}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = 'relationalai'
-version = '0.12.12'
+version = '0.12.13'
 description = 'RelationalAI Library and CLI'
 readme="docs/pypi/README.md"
 authors = [

{relationalai-0.12.12 → relationalai-0.12.13}/src/relationalai/clients/snowflake.py

@@ -441,7 +441,8 @@ class Resources(ResourcesBase):
         code: str,
         params: List[Any] | Any | None = None,
         raw: bool = False,
-        help: bool = True
+        help: bool = True,
+        skip_engine_db_error_retry: bool = False
     ) -> Any:
         # print(f"\n--- sql---\n{code}\n--- end sql---\n")
         if not self._session:
@@ -467,7 +468,8 @@ 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 _is_engine_issue(orig_message) or _is_database_issue(orig_message):
+            # skip initializing the index if the query is a user transaction. exec_raw/exec_lqp will handle that case with the correct request headers.
+            if (_is_engine_issue(orig_message) or _is_database_issue(orig_message)) and not skip_engine_db_error_retry:
                 try:
                     self._poll_use_index(
                         app_name=self.get_app_name(),
@@ -1609,9 +1611,11 @@ Otherwise, remove it from your '{profile}' configuration profile.
                 sql_string = f"CALL {APP_NAME}.api.exec_async_v2('{database}','{engine}', ?, {inputs}, {readonly}, {nowait_durable}, '{language}', {query_timeout_mins}, {request_headers});"
             else:
                 sql_string = f"CALL {APP_NAME}.api.exec_async_v2('{database}','{engine}', ?, {inputs}, {readonly}, {nowait_durable}, '{language}', {request_headers});"
+        # Don't let exec setup GI on failure, exec_raw and exec_lqp will do that and add the correct headers.
         response = self._exec(
             sql_string,
             raw_code,
+            skip_engine_db_error_retry=True,
         )
         if not response:
             raise Exception("Failed to create transaction")
@@ -1629,6 +1633,7 @@ Otherwise, remove it from your '{profile}' configuration profile.
         bypass_index=False,
         language: str = "rel",
         query_timeout_mins: int | None = None,
+        gi_setup_skipped: bool = False,
     ):
         if inputs is None:
             inputs = {}
@@ -1638,6 +1643,8 @@ Otherwise, remove it from your '{profile}' configuration profile.
         with debugging.span("transaction", **query_attrs_dict) as txn_span:
             with debugging.span("create_v2", **query_attrs_dict) as create_span:
                 request_headers['user-agent'] = get_pyrel_version(self.generation)
+                request_headers['gi_setup_skipped'] = str(gi_setup_skipped)
+                request_headers['pyrel_program_id'] = debugging.get_program_span_id() or ""
                 response = self._exec_rai_app(
                     database=database,
                     engine=engine,
@@ -1897,26 +1904,29 @@ Otherwise, remove it from your '{profile}' configuration profile.
     # Exec
     #--------------------------------------------------
 
-    def exec_lqp(
+    def _exec_with_gi_retry(
         self,
         database: str,
         engine: str | None,
-        raw_code: bytes,
-        readonly=True,
-        *,
-        inputs: Dict | None = None,
-        nowait_durable=False,
-        headers: Dict | None = None,
-        bypass_index=False,
-        query_timeout_mins: int | None = None,
+        raw_code: str,
+        inputs: Dict | None,
+        readonly: bool,
+        nowait_durable: bool,
+        headers: Dict | None,
+        bypass_index: bool,
+        language: str,
+        query_timeout_mins: int | None,
     ):
-        raw_code_b64 = base64.b64encode(raw_code).decode("utf-8")
+        """Execute with graph index retry logic.
 
+        Attempts execution with gi_setup_skipped=True first. If an engine or database
+        issue occurs, polls use_index and retries with gi_setup_skipped=False.
+        """
         try:
             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,
+                database, engine, raw_code, inputs, readonly, nowait_durable,
+                headers=headers, bypass_index=bypass_index, language=language,
+                query_timeout_mins=query_timeout_mins, gi_setup_skipped=True,
             )
         except Exception as e:
             err_message = str(e).lower()
@@ -1933,13 +1943,32 @@ Otherwise, remove it from your '{profile}' configuration profile.
                 )
 
                 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,
+                    database, engine, raw_code, inputs, readonly, nowait_durable,
+                    headers=headers, bypass_index=bypass_index, language=language,
+                    query_timeout_mins=query_timeout_mins, gi_setup_skipped=False,
                 )
             else:
                 raise e
 
+    def exec_lqp(
+        self,
+        database: str,
+        engine: str | None,
+        raw_code: bytes,
+        readonly=True,
+        *,
+        inputs: Dict | None = None,
+        nowait_durable=False,
+        headers: Dict | None = None,
+        bypass_index=False,
+        query_timeout_mins: int | None = None,
+    ):
+        raw_code_b64 = base64.b64encode(raw_code).decode("utf-8")
+        return self._exec_with_gi_retry(
+            database, engine, raw_code_b64, inputs, readonly, nowait_durable,
+            headers, bypass_index, 'lqp', query_timeout_mins
+        )
+
 
     def exec_raw(
         self,
@@ -1955,45 +1984,10 @@ Otherwise, remove it from your '{profile}' configuration profile.
         query_timeout_mins: int | None = None,
     ):
         raw_code = raw_code.replace("'", "\\'")
-
-        try:
-            return self._exec_async_v2(
-                database,
-                engine,
-                raw_code,
-                inputs,
-                readonly,
-                nowait_durable,
-                headers=headers,
-                bypass_index=bypass_index,
-                query_timeout_mins=query_timeout_mins,
-            )
-        except Exception as e:
-            err_message = str(e).lower()
-            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,
-                    inputs,
-                    readonly,
-                    nowait_durable,
-                    headers=headers,
-                    bypass_index=bypass_index,
-                    query_timeout_mins=query_timeout_mins,
-                )
-            else:
-                raise e
+        return self._exec_with_gi_retry(
+            database, engine, raw_code, inputs, readonly, nowait_durable,
+            headers, bypass_index, 'rel', query_timeout_mins
+        )
 
 
     def format_results(self, results, task:m.Task|None=None) -> Tuple[DataFrame, List[Any]]:
@@ -3281,6 +3275,7 @@ class DirectAccessResources(Resources):
         path_params: Dict[str, str] | None = None,
         query_params: Dict[str, str] | None = None,
         skip_auto_create: bool = False,
+        skip_engine_db_error_retry: bool = False,
     ) -> requests.Response:
         with debugging.span("direct_access_request"):
             def _send_request():
@@ -3314,8 +3309,8 @@ class DirectAccessResources(Resources):
                         message = ""  # Not used when we check status_code directly
 
                     # fix engine on engine error and retry
-                    # Skip auto-retry if skip_auto_create is True to avoid recursion
-                    if (_is_engine_issue(message) and not skip_auto_create) or _is_database_issue(message):
+                    # Skip setting up GI if skip_auto_create is True to avoid recursion or skip_engine_db_error_retry is true to let _exec_async_v2 perform the retry with the correct headers.
+                    if ((_is_engine_issue(message) and not skip_auto_create) or _is_database_issue(message)) and not skip_engine_db_error_retry:
                         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()
@@ -3340,6 +3335,48 @@ class DirectAccessResources(Resources):
                 raise e
             return response
 
+    def _txn_request_with_gi_retry(
+        self,
+        payload: Dict,
+        headers: Dict[str, str],
+        query_params: Dict,
+        engine: Union[str, None],
+    ):
+        """Make request with graph index retry logic.
+
+        Attempts request with gi_setup_skipped=True first. If an engine or database
+        issue occurs, polls use_index and retries with gi_setup_skipped=False.
+        """
+        response = self.request(
+            "create_txn", payload=payload, headers=headers, query_params=query_params, skip_auto_create=True, skip_engine_db_error_retry=True
+        )
+
+        if response.status_code != 200:
+            try:
+                message = response.json().get("message", "")
+            except requests.exceptions.JSONDecodeError:
+                message = ""
+
+            if _is_engine_issue(message) or _is_database_issue(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=self.database,
+                    engine_name=engine_name,
+                    engine_size=engine_size,
+                    headers=headers,
+                )
+                headers['gi_setup_skipped'] = 'False'
+                response = self.request(
+                    "create_txn", payload=payload, headers=headers, query_params=query_params, skip_auto_create=True, skip_engine_db_error_retry=True
+                )
+            else:
+                raise ResponseStatusException("Failed to create transaction.", response)
+
+        return response
+
     def _exec_async_v2(
         self,
         database: str,
@@ -3352,6 +3389,7 @@ class DirectAccessResources(Resources):
         bypass_index=False,
         language: str = "rel",
         query_timeout_mins: int | None = None,
+        gi_setup_skipped: bool = False,
     ):
 
         with debugging.span("transaction") as txn_span:
@@ -3374,12 +3412,15 @@ class DirectAccessResources(Resources):
                     payload["timeout_mins"] = query_timeout_mins
                 query_params={"use_graph_index": str(use_graph_index and not bypass_index)}
 
-                response = self.request(
-                    "create_txn", payload=payload, headers=headers, query_params=query_params,
-                )
+                # Add gi_setup_skipped to headers
+                if headers is None:
+                    headers = {}
+                headers["gi_setup_skipped"] = str(gi_setup_skipped)
+                headers['pyrel_program_id'] = debugging.get_program_span_id() or ""
 
-                if response.status_code != 200:
-                    raise ResponseStatusException("Failed to create transaction.", response)
+                response = self._txn_request_with_gi_retry(
+                    payload, headers, query_params, engine
+                )
 
                 artifact_info = {}
                 response_content = response.json()