relationalai 0.12.10__tar.gz → 0.12.11__tar.gz

{relationalai-0.12.10 → relationalai-0.12.11}/PKG-INFO

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

{relationalai-0.12.10 → relationalai-0.12.11}/pyproject.toml

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

{relationalai-0.12.10 → relationalai-0.12.11}/src/relationalai/clients/snowflake.py

@@ -441,8 +441,7 @@ class Resources(ResourcesBase):
         code: str,
         params: List[Any] | Any | None = None,
         raw: bool = False,
-        help: bool = True,
-        skip_auto_create: bool = False
+        help: bool = True
     ) -> Any:
         # print(f"\n--- sql---\n{code}\n--- end sql---\n")
         if not self._session:
@@ -459,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")
@@ -467,10 +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
-            # skip creating the engine if the query is a user transaction. exec_async_v2 will handle that case.
-            if _is_engine_issue(orig_message) and not skip_auto_create:
+            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
@@ -1607,7 +1612,6 @@ Otherwise, remove it from your '{profile}' configuration profile.
         response = self._exec(
             sql_string,
             raw_code,
-            skip_auto_create=True,
         )
         if not response:
             raise Exception("Failed to create transaction")
@@ -1625,7 +1629,6 @@ 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 = {}
@@ -1635,8 +1638,6 @@ 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,
@@ -1896,29 +1897,26 @@ Otherwise, remove it from your '{profile}' configuration profile.
     # Exec
     #--------------------------------------------------
 
-    def _exec_with_gi_retry(
+    def exec_lqp(
         self,
         database: str,
         engine: str | 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: bytes,
+        readonly=True,
+        *,
+        inputs: Dict | None = None,
+        nowait_durable=False,
+        headers: Dict | None = None,
+        bypass_index=False,
+        query_timeout_mins: int | None = None,
     ):
-        """Execute with graph index retry logic.
+        raw_code_b64 = base64.b64encode(raw_code).decode("utf-8")
 
-        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, inputs, readonly, nowait_durable,
-                headers=headers, bypass_index=bypass_index, language=language,
-                query_timeout_mins=query_timeout_mins, gi_setup_skipped=True,
+                database, engine, raw_code_b64, inputs, readonly, nowait_durable,
+                headers=headers, bypass_index=bypass_index, language='lqp',
+                query_timeout_mins=query_timeout_mins,
             )
         except Exception as e:
             err_message = str(e).lower()
@@ -1935,32 +1933,13 @@ Otherwise, remove it from your '{profile}' configuration profile.
                 )
 
                 return self._exec_async_v2(
-                    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,
+                    database, engine, raw_code_b64, inputs, readonly, nowait_durable,
+                    headers=headers, bypass_index=bypass_index, language='lqp',
+                    query_timeout_mins=query_timeout_mins,
                 )
             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,
@@ -1976,10 +1955,45 @@ Otherwise, remove it from your '{profile}' configuration profile.
         query_timeout_mins: int | None = None,
     ):
         raw_code = raw_code.replace("'", "\\'")
-        return self._exec_with_gi_retry(
-            database, engine, raw_code, inputs, readonly, nowait_durable,
-            headers, bypass_index, 'rel', query_timeout_mins
-        )
+
+        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
 
 
     def format_results(self, results, task:m.Task|None=None) -> Tuple[DataFrame, List[Any]]:
@@ -3300,10 +3314,19 @@ 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 or to let _exec_async_v2 poll the index.
-                    if _is_engine_issue(message) and not skip_auto_create:
-                        engine = payload.get("engine_name", "") if payload else ""
-                        self.auto_create_engine(engine)
+                    # 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):
+                        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):
@@ -3317,48 +3340,6 @@ 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
-        )
-
-        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
-                )
-            else:
-                raise ResponseStatusException("Failed to create transaction.", response)
-
-        return response
-
     def _exec_async_v2(
         self,
         database: str,
@@ -3371,7 +3352,6 @@ 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:
@@ -3394,16 +3374,13 @@ class DirectAccessResources(Resources):
                     payload["timeout_mins"] = query_timeout_mins
                 query_params={"use_graph_index": str(use_graph_index and not bypass_index)}
 
-                # 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 ""
-
-                response = self._txn_request_with_gi_retry(
-                    payload, headers, query_params, engine
+                response = self.request(
+                    "create_txn", payload=payload, headers=headers, query_params=query_params,
                 )
 
+                if response.status_code != 200:
+                    raise ResponseStatusException("Failed to create transaction.", response)
+
                 artifact_info = {}
                 response_content = response.json()