relationalai 0.11.0__py3-none-any.whl → 0.11.2__py3-none-any.whl

relationalai/clients/azure.py

@@ -12,7 +12,7 @@ from pandas import DataFrame
 from relationalai import debugging
 from relationalai.clients.util import poll_with_specified_overhead
 
-from ..errors import EngineNotFoundException, RAIException
+from ..errors import EngineNotFoundException, RAIException, AzureUnsupportedQueryTimeoutException
 from ..rel_utils import assert_no_problems
 from ..loaders.loader import emit_delete_import, import_file, list_available_resources
 from .config import Config
@@ -227,10 +227,13 @@ class Resources(ResourcesBase):
     def _exec(self, code:str, params:List[Any]|Any|None = None, raw=False, help=True):
         raise Exception("Azure doesn't support _exec")
 
-    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):
+    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):
         raise Exception("Azure doesn't support exec_lqp")
 
-    def exec_raw(self, database:str, engine:str|None, raw_code:str, readonly=True, *, inputs: Dict | None = None, nowait_durable=False, headers: Dict | None = None, raw_results=True):
+    def exec_raw(self, database:str, engine:str|None, raw_code:str, readonly=True, *, inputs: Dict | None = None, nowait_durable=False, headers: Dict | None = None, raw_results=True, query_timeout_mins: int | None = None):
+        if query_timeout_mins is not None or self.config.get("query_timeout_mins", None) is not None:
+            config_file_path = getattr(self.config, 'file_path', None)
+            raise AzureUnsupportedQueryTimeoutException(config_file_path=config_file_path)
         if engine is None:
             engine = self.get_default_engine_name()
         try:
@@ -289,7 +292,7 @@ class Resources(ResourcesBase):
     # Exec format
     #--------------------------------------------------
 
-    def exec_format(self, database: str, engine: str, raw_code: str, task:m.Task, format:str, inputs: Dict | None = None, readonly: bool = True, nowait_durable=False, skip_invalid_data=False, headers: Dict | None = None) -> Any: # @FIXME: Better type annotation
+    def exec_format(self, database: str, engine: str, raw_code: str, cols:List[str], format:str, inputs: Dict | None = None, readonly: bool = True, nowait_durable=False, skip_invalid_data=False, headers: Dict | None = None, query_timeout_mins: int | None = None) -> Any: # @FIXME: Better type annotation
         raise Exception("Azure doesn't support alternative formats yet")
 
     def to_model_type(self, model:dsl.Graph, name: str, source:str):

relationalai/clients/client.py

@@ -3,7 +3,7 @@ import atexit
 from datetime import datetime, timedelta, timezone
 import re
 from collections import defaultdict
-from typing import Dict, List, Any, Tuple, cast, Callable
+from typing import Dict, List, Any, Optional, Tuple, cast, Callable
 
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
@@ -401,15 +401,15 @@ class ResourcesBase(ABC):
         pass
 
     @abstractmethod
-    def exec_lqp(self, database: str, engine: str | None, raw_code: bytes, readonly: bool = True, *, inputs: Dict | None = None, nowait_durable=False, headers: Dict | None = None) -> Any: # @FIXME: Better type annotation
+    def exec_lqp(self, database: str, engine: str | None, raw_code: bytes, readonly: bool = True, *, inputs: Dict | None = None, nowait_durable=False, headers: Dict | None = None, query_timeout_mins: int | None = None) -> Any: # @FIXME: Better type annotation
         pass
 
     @abstractmethod
-    def exec_raw(self, database: str, engine: str | None, raw_code: str, readonly: bool = True, *, inputs: Dict | None = None, nowait_durable=False, headers: Dict | None = None) -> Any: # @FIXME: Better type annotation
+    def exec_raw(self, database: str, engine: str | None, raw_code: str, readonly: bool = True, *, inputs: Dict | None = None, nowait_durable=False, headers: Dict | None = None, query_timeout_mins: Optional[int]=None) -> Any: # @FIXME: Better type annotation
         pass
 
     @abstractmethod
-    def exec_format(self, database: str, engine: str, raw_code: str, task:m.Task, format:str, inputs: Dict | None = None, readonly: bool = True, nowait_durable=False, skip_invalid_data=False, headers: Dict | None = None) -> Any: # @FIXME: Better type annotation
+    def exec_format(self, database: str, engine: str, raw_code: str, cols:List[str], format:str, inputs: Dict | None = None, readonly: bool = True, nowait_durable=False, skip_invalid_data=False, headers: Dict | None = None, query_timeout_mins: Optional[int]=None) -> Any: # @FIXME: Better type annotation
         pass
 
     @abstractmethod
@@ -695,10 +695,10 @@ class Client():
         code = self.compiler.compile(dsl.build.raw_task(content))
         self._install_batch.set(path, code)
 
-    def exec_raw(self, code:str, readonly=True, raw_results=True, inputs: Dict | None = None, internal=False, nowait_durable=None, abort_on_error=True, headers: Dict | None = None) -> DataFrame|Any:
+    def exec_raw(self, code:str, readonly=True, raw_results=True, inputs: Dict | None = None, internal=False, nowait_durable=None, abort_on_error=True, headers: Dict | None = None, query_timeout_mins: Optional[int]=None) -> DataFrame|Any:
         task = dsl.build.raw_task(code)
         debugging.set_source(task)
-        return self.query(task, read_only=readonly, raw_results=raw_results, inputs=inputs, internal=internal, nowait_durable=nowait_durable, headers=headers, abort_on_error=abort_on_error)
+        return self.query(task, read_only=readonly, raw_results=raw_results, inputs=inputs, internal=internal, nowait_durable=nowait_durable, headers=headers, abort_on_error=abort_on_error, query_timeout_mins=query_timeout_mins)
 
     def exec_control(self, code:str, cb:Callable[[DataFrame]]|None=None):
         self._install_batch.control_items.append((code, cb))
@@ -719,12 +719,12 @@ class Client():
     # Query
     #--------------------------------------------------
 
-    def _query(self, code:str, task:m.Task|None, end_span, readonly=False, inputs: Dict | None = None, nowait_durable=None, headers: Dict | None = None, abort_on_error=True):
+    def _query(self, code:str, task:m.Task|None, end_span, readonly=False, inputs: Dict | None = None, nowait_durable=None, headers: Dict | None = None, abort_on_error=True, query_timeout_mins: Optional[int]=None):
         if nowait_durable is None:
             nowait_durable = self.isolated
 
         try:
-            results = self.resources.exec_raw(self._database, self.get_engine_name(), code, readonly=readonly, inputs=inputs, nowait_durable=nowait_durable, headers=headers)
+            results = self.resources.exec_raw(self._database, self.get_engine_name(), code, readonly=readonly, inputs=inputs, nowait_durable=nowait_durable, headers=headers, query_timeout_mins=query_timeout_mins)
             dataframe, errors = self.resources.format_results(results, task)
             end_span["results"] = dataframe
             end_span["errors"] = errors
@@ -736,13 +736,14 @@ class Client():
                 engine_name = self.get_engine_name()
                 self.resources.resume_engine(engine_name, headers=headers)
                 # invoke _query again to retry the query
-                return self._query(code, task, end_span, readonly=readonly, inputs=inputs, nowait_durable=nowait_durable, headers=headers, abort_on_error=abort_on_error)
+                return self._query(code, task, end_span, readonly=readonly, inputs=inputs, nowait_durable=nowait_durable, headers=headers, abort_on_error=abort_on_error, query_timeout_mins=query_timeout_mins)
             else:
                 raise e
 
 
-    def _query_format(self, code:str, task:m.Task, end_span, format, readonly=False, skip_invalid_data=False, inputs: Dict | None = None):
-        results, raw = self.resources.exec_format(self._database, self.get_engine_name(), code, task, readonly=readonly, inputs=inputs, format=format, skip_invalid_data=skip_invalid_data)
+    def _query_format(self, code:str, task:m.Task, end_span, format, readonly=False, skip_invalid_data=False, inputs: Dict | None = None, query_timeout_mins: Optional[int]=None):
+        cols = task.return_cols(allow_dups=False)
+        results, raw = self.resources.exec_format(self._database, self.get_engine_name(), code, cols, readonly=readonly, inputs=inputs, format=format, skip_invalid_data=skip_invalid_data, query_timeout_mins=query_timeout_mins)
         errors = []
         if raw:
             dataframe, errors = self.resources.format_results(raw, task)
@@ -752,7 +753,7 @@ class Client():
         # return results if raw_results else dataframe
         return results, raw
 
-    def query(self, task:m.Task, rentrant=False, read_only=False, raw_results=False, inputs: Dict | None = None, format="pandas", tag=None, nowait_durable=None, headers: Dict | None = None, internal=False, abort_on_error=True, skip_invalid_data = False) -> DataFrame|Any:
+    def query(self, task:m.Task, rentrant=False, read_only=False, raw_results=False, inputs: Dict | None = None, format="pandas", tag=None, nowait_durable=None, headers: Dict | None = None, internal=False, abort_on_error=True, skip_invalid_data = False, query_timeout_mins: Optional[int]=None) -> DataFrame|Any:
         if not self.dry_run and self.use_graph_index:
             self.create_database(isolated=self.isolated, headers=headers)
 
@@ -793,10 +794,10 @@ class Client():
 
             start = time.perf_counter()
             if format == "pandas":
-                results, raw = self._query(code, task, end_span, readonly=read_only, inputs=inputs, nowait_durable=nowait_durable, headers=headers, abort_on_error=abort_on_error)
+                results, raw = self._query(code, task, end_span, readonly=read_only, inputs=inputs, nowait_durable=nowait_durable, headers=headers, abort_on_error=abort_on_error, query_timeout_mins=query_timeout_mins)
                 debugging.time("query", time.perf_counter() - start, DataFrame() if raw_results else results, internal=internal, source_map=source_map)
             else:
-                results, raw = self._query_format(code, task, end_span, readonly=read_only, inputs=inputs, format=format, skip_invalid_data=skip_invalid_data)
+                results, raw = self._query_format(code, task, end_span, readonly=read_only, inputs=inputs, format=format, skip_invalid_data=skip_invalid_data, query_timeout_mins=query_timeout_mins)
                 debugging.time("query", time.perf_counter() - start, DataFrame(), source_map=source_map, alt_format_results=results)
 
             self._install_batch.clear_dirty()

relationalai/clients/config.py

@@ -28,6 +28,10 @@ PUBLIC_CONFIG_KEYS = [
     "ensure_change_tracking",
     "download_url_type",
     "use_direct_access",
+    # query_timeout_mins allows to specify a timeout in minutes applied to all queries. When
+    # a query execution time exceeds this timeout, the query will be aborted. This is useful
+    # to avoid long-running queries that can incur high costs.
+    "query_timeout_mins",
 ]
 
 #--------------------------------------------------

relationalai/clients/snowflake.py

@@ -21,7 +21,7 @@ from relationalai.clients.use_index_poller import DirectUseIndexPoller, UseIndex
 import snowflake.snowpark
 
 from relationalai.rel_utils import sanitize_identifier, to_fqn_relation_name
-from relationalai.tools.constants import FIELD_PLACEHOLDER, RAI_APP_NAME, SNOWFLAKE_AUTHS, USE_GRAPH_INDEX, USE_DIRECT_ACCESS, WAIT_FOR_STREAM_SYNC, Generation
+from relationalai.tools.constants import FIELD_PLACEHOLDER, RAI_APP_NAME, SNOWFLAKE_AUTHS, USE_GRAPH_INDEX, USE_DIRECT_ACCESS, DEFAULT_QUERY_TIMEOUT_MINS, WAIT_FOR_STREAM_SYNC, Generation
 from .. import std
 from collections import defaultdict
 import requests
@@ -1516,8 +1516,11 @@ Otherwise, remove it from your '{profile}' configuration profile.
         request_headers: Dict | None = None,
         bypass_index=False,
         language: str = "rel",
+        query_timeout_mins: int | None = None,
     ):
         assert language == "rel" or language == "lqp", "Only 'rel' and 'lqp' languages are supported"
+        if query_timeout_mins is None and (timeout_value := self.config.get("query_timeout_mins", DEFAULT_QUERY_TIMEOUT_MINS)) is not None:
+            query_timeout_mins = int(timeout_value)
         # Depending on the shape of the input, the behavior of exec_async_v2 changes.
         # When using the new format (with an object), the function retrieves the
         # 'rai' database by hashing the model and username. In contrast, the
@@ -1526,9 +1529,23 @@ Otherwise, remove it from your '{profile}' configuration profile.
         # graph index to ensure the correct database is utilized.
         use_graph_index = self.config.get("use_graph_index", USE_GRAPH_INDEX)
         if use_graph_index and not bypass_index:
-            sql_string = f"CALL {APP_NAME}.api.exec_async_v2(?, {{'database': '{database}', 'engine': '{engine}', 'inputs': {inputs}, 'readonly': {readonly}, 'nowait_durable': {nowait_durable}, 'language': '{language}', 'headers': {request_headers}}});"
+            payload = {
+                'database': database,
+                'engine': engine,
+                'inputs': inputs,
+                'readonly': readonly,
+                'nowait_durable': nowait_durable,
+                'language': language,
+                'headers': request_headers
+            }
+            if query_timeout_mins is not None:
+                payload["timeout_mins"] = query_timeout_mins
+            sql_string = f"CALL {APP_NAME}.api.exec_async_v2(?, {payload});"
         else:
-            sql_string = f"CALL {APP_NAME}.api.exec_async_v2('{database}','{engine}', ?, {inputs}, {readonly}, {nowait_durable}, '{language}', {request_headers});"
+            if query_timeout_mins is not None:
+                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});"
         response = self._exec(
             sql_string,
             raw_code,
@@ -1548,6 +1565,7 @@ Otherwise, remove it from your '{profile}' configuration profile.
         headers: Dict | None = None,
         bypass_index=False,
         language: str = "rel",
+        query_timeout_mins: int | None = None,
     ):
         if inputs is None:
             inputs = {}
@@ -1567,6 +1585,7 @@ Otherwise, remove it from your '{profile}' configuration profile.
                     request_headers=request_headers,
                     bypass_index=bypass_index,
                     language=language,
+                    query_timeout_mins=query_timeout_mins,
                 )
 
                 artifact_info = {}
@@ -1825,14 +1844,16 @@ Otherwise, remove it from your '{profile}' configuration profile.
         inputs: Dict | None = None,
         nowait_durable=False,
         headers: Dict | None = None,
-        bypass_index=False
+        bypass_index=False,
+        query_timeout_mins: int | None = None,
     ):
         raw_code_b64 = base64.b64encode(raw_code).decode("utf-8")
 
         try:
             return self._exec_async_v2(
                 database, engine, raw_code_b64, inputs, readonly, nowait_durable,
-                headers=headers, bypass_index=bypass_index, language='lqp'
+                headers=headers, bypass_index=bypass_index, language='lqp',
+                query_timeout_mins=query_timeout_mins,
             )
         except Exception as e:
             err_message = str(e).lower()
@@ -1840,7 +1861,8 @@ Otherwise, remove it from your '{profile}' configuration profile.
                 self.auto_create_engine(engine)
                 self._exec_async_v2(
                     database, engine, raw_code_b64, inputs, readonly, nowait_durable,
-                    headers=headers, bypass_index=bypass_index, language='lqp'
+                    headers=headers, bypass_index=bypass_index, language='lqp',
+                    query_timeout_mins=query_timeout_mins,
                 )
             else:
                 raise e
@@ -1856,19 +1878,38 @@ Otherwise, remove it from your '{profile}' configuration profile.
         inputs: Dict | None = None,
         nowait_durable=False,
         headers: Dict | None = None,
-        bypass_index=False
+        bypass_index=False,
+        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
+                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):
                 self.auto_create_engine(engine)
-                return self._exec_async_v2(database, engine, raw_code, inputs, readonly, nowait_durable, headers=headers, bypass_index=bypass_index)
+                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
 
@@ -1885,13 +1926,14 @@ Otherwise, remove it from your '{profile}' configuration profile.
         database: str,
         engine: str,
         raw_code: str,
-        task: m.Task,
+        cols: List[str],
         format: str,
         inputs: Dict | None = None,
         readonly=True,
         nowait_durable=False,
         skip_invalid_data=False,
         headers: Dict | None = None,
+        query_timeout_mins: int | None = None,
     ):
         if inputs is None:
             inputs = {}
@@ -1899,6 +1941,8 @@ Otherwise, remove it from your '{profile}' configuration profile.
             headers = {}
         if 'user-agent' not in headers:
             headers['user-agent'] = get_pyrel_version(self.generation)
+        if query_timeout_mins is None and (timeout_value := self.config.get("query_timeout_mins", DEFAULT_QUERY_TIMEOUT_MINS)) is not None:
+            query_timeout_mins = int(timeout_value)
         # TODO: add headers
         start = time.perf_counter()
         output_table = "out" + str(uuid.uuid4()).replace("-", "_")
@@ -1909,18 +1953,25 @@ Otherwise, remove it from your '{profile}' configuration profile.
         col_names_map = None
         artifacts = None
         assert self._session
-        temp = self._session.createDataFrame([], StructType([StructField(name, StringType()) for name in task.return_cols(allow_dups=False)]))
+        temp = self._session.createDataFrame([], StructType([StructField(name, StringType()) for name in cols]))
         with debugging.span("transaction") as txn_span:
             try:
                 # In the graph index case we need to use the new exec_into_table proc as it obfuscates the db name
                 with debugging.span("exec_format"):
                     if use_graph_index:
-                        res = self._exec(f"call {APP_NAME}.api.exec_into_table(?, ?, ?, ?, ?, ?, ?);", [database, engine, raw_code, output_table, readonly, nowait_durable, skip_invalid_data])
+                        # we do not provide a default value for query_timeout_mins so that we can control the default on app level
+                        if query_timeout_mins is not None:
+                            res = self._exec(f"call {APP_NAME}.api.exec_into_table(?, ?, ?, ?, ?, ?, ?, ?);", [database, engine, raw_code, output_table, readonly, nowait_durable, skip_invalid_data, query_timeout_mins])
+                        else:
+                            res = self._exec(f"call {APP_NAME}.api.exec_into_table(?, ?, ?, ?, ?, ?, ?);", [database, engine, raw_code, output_table, readonly, nowait_durable, skip_invalid_data])
                         txn_id = json.loads(res[0]["EXEC_INTO_TABLE"])["rai_transaction_id"]
                         rejected_rows = json.loads(res[0]["EXEC_INTO_TABLE"]).get("rejected_rows", [])
                         rejected_rows_count = json.loads(res[0]["EXEC_INTO_TABLE"]).get("rejected_rows_count", 0)
                     else:
-                        res = self._exec(f"call {APP_NAME}.api.exec_into(?, ?, ?, ?, ?, {inputs}, ?, {headers}, ?);", [database, engine, raw_code, output_table, readonly, nowait_durable, skip_invalid_data])
+                        if query_timeout_mins is not None:
+                            res = self._exec(f"call {APP_NAME}.api.exec_into(?, ?, ?, ?, ?, {inputs}, ?, {headers}, ?, ?);", [database, engine, raw_code, output_table, readonly, nowait_durable, skip_invalid_data, query_timeout_mins])
+                        else:
+                            res = self._exec(f"call {APP_NAME}.api.exec_into(?, ?, ?, ?, ?, {inputs}, ?, {headers}, ?);", [database, engine, raw_code, output_table, readonly, nowait_durable, skip_invalid_data])
                         txn_id = json.loads(res[0]["EXEC_INTO"])["rai_transaction_id"]
                         rejected_rows = json.loads(res[0]["EXEC_INTO"]).get("rejected_rows", [])
                         rejected_rows_count = json.loads(res[0]["EXEC_INTO"]).get("rejected_rows_count", 0)
@@ -1932,7 +1983,7 @@ Otherwise, remove it from your '{profile}' configuration profile.
                     if out_sample:
                         keys = set([k.lower() for k in out_sample[0].as_dict().keys()])
                         col_names_map = {}
-                        for ix, name in enumerate(task.return_cols(allow_dups=False)):
+                        for ix, name in enumerate(cols):
                             col_key = f"col{ix:03}"
                             if col_key in keys:
                                 col_names_map[col_key] = IdentityParser(name).identity
@@ -3095,6 +3146,7 @@ class DirectAccessResources(Resources):
         headers: Dict[str, str] | None = None,
         bypass_index=False,
         language: str = "rel",
+        query_timeout_mins: int | None = None,
     ):
 
         with debugging.span("transaction") as txn_span:
@@ -3111,6 +3163,10 @@ class DirectAccessResources(Resources):
                     "readonly": readonly,
                     "language": language,
                 }
+                if query_timeout_mins is None and (timeout_value := self.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:
+                    payload["timeout_mins"] = query_timeout_mins
                 query_params={"use_graph_index": str(use_graph_index and not bypass_index)}
 
                 response = self.request(

relationalai/dsl.py

@@ -1736,9 +1736,9 @@ class Graph:
                         file_path = os.path.join(root, file)
                         self._client.load_raw_file(file_path)
 
-    def exec_raw(self, code:str, readonly=False, raw_results=True, abort_on_error=True, inputs={}):
+    def exec_raw(self, code:str, readonly=False, raw_results=True, abort_on_error=True, inputs={}, query_timeout_mins: Optional[int]=None):
         try:
-            return self._client.exec_raw(code, readonly=readonly, raw_results=raw_results, inputs=inputs, abort_on_error=abort_on_error)
+            return self._client.exec_raw(code, readonly=readonly, raw_results=raw_results, inputs=inputs, abort_on_error=abort_on_error, query_timeout_mins=query_timeout_mins)
         except KeyboardInterrupt as e:
             print("Canceling transactions...")
             self.resources.cancel_pending_transactions()

relationalai/early_access/dsl/codegen/weaver.py

@@ -103,8 +103,7 @@ class Weaver:
             # binding for each supertype with a reference scheme.
             #=
             constructor_binding = (
-                self._binder.lookup_constructor_binding(concept, binding.column)
-                if constructed_concept is concept
+                binding if constructed_concept is concept
                 else self._binder.lookup_constructor_binding_by_source(constructed_concept, binding.column.table)
             )
 

relationalai/errors.py

@@ -2395,3 +2395,40 @@ class UnsupportedColumnTypesWarning(RAIWarning):
 
 {note}
 """)
+
+class QueryTimeoutExceededException(RAIException):
+    def __init__(self, timeout_mins: int, config_file_path: str | None = None):
+        self.timeout_mins = timeout_mins
+        self.message = f"Query execution time exceeded the specified timeout of {timeout_mins} minutes."
+        self.name = "Query Timeout Exceeded"
+        self.config_file_path = config_file_path or ""
+        self.content = self.format_message()
+        super().__init__(self.message, self.name, self.content)
+
+    def format_message(self):
+        return textwrap.dedent(f"""
+        {self.message}
+
+        Consider increasing the 'query_timeout_mins' parameter in your configuration file{f' (stored in {self.config_file_path})' if self.config_file_path else ''} to allow more time for query execution.
+        """)
+
+
+#--------------------------------------------------
+# Azure Exceptions
+#--------------------------------------------------
+
+class AzureUnsupportedQueryTimeoutException(RAIException):
+    def __init__(self, config_file_path: str | None = None):
+        self.message = "Query timeouts aren't supported on platform Azure."
+        self.name = "Azure Unsupported Query Timeout Error"
+        self.config_file_path = config_file_path or ""
+        self.content = self.format_message()
+        super().__init__(self.message, self.name, self.content)
+
+    def format_message(self):
+        return textwrap.dedent(f"""
+        {self.message}
+
+        Please remove the 'query_timeout_mins' from your configuration file{f' (stored in {self.config_file_path})' if self.config_file_path else ''} when running on platform Azure.
+        """)
+

relationalai/experimental/solvers.py

@@ -1,17 +1,20 @@
 from __future__ import annotations
-from typing import Any, List
+import time
+from typing import Any, List, Optional
 from dataclasses import dataclass
 import textwrap
 from .. import dsl, std
 from ..std import rel
 from ..metamodel import Builtins
 from ..tools.cli_controls import Spinner
+from ..tools.constants import DEFAULT_QUERY_TIMEOUT_MINS
 from .. import debugging
 import uuid
 import relationalai
 import json
 from ..clients.util import poll_with_specified_overhead
 from ..clients.snowflake import Resources as SnowflakeResources
+from ..util.timeout import calc_remaining_timeout_minutes
 
 rel_sv = rel._tagged(Builtins.SingleValued)
 
@@ -203,10 +206,20 @@ class SolverModel:
         payload["options"] = options
         payload["model_uri"] = sf_input_uri
 
+        rai_config = self.graph._config
+        query_timeout_mins = kwargs.get("query_timeout_mins", None)
+        if query_timeout_mins is None and (timeout_value := rai_config.get("query_timeout_mins", DEFAULT_QUERY_TIMEOUT_MINS)) is not None:
+            query_timeout_mins = int(timeout_value)
+        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("")
+        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
@@ -235,7 +248,8 @@ class SolverModel:
             def config[:envelope, :payload, :data]: {scope}model_string
             def config[:envelope, :payload, :path]: "{model_uri}"
             def export {{ config }}
-            """)
+            """),
+            query_timeout_mins=remaining_timeout_minutes,
         )
         txn = response.transaction or {}
         # The above `exec_raw` will throw an error if the transaction
@@ -248,17 +262,26 @@ class SolverModel:
             raise Exception(f"Transaction that materializes the solver inputs did not complete! ID: `{txn['id']}` State `{txn['state']}`")
 
         # 2. Execute job and wait for completion.
+        remaining_timeout_minutes = calc_remaining_timeout_minutes(
+            start_time, query_timeout_mins, config_file_path=config_file_path
+        )
         try:
-            job_id = solver._exec_job(payload, log_to_console=log_to_console)
+            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):
                 solver._auto_create_solver_async()
-                job_id = solver._exec_job(payload, log_to_console=log_to_console)
+                remaining_timeout_minutes = calc_remaining_timeout_minutes(
+                    start_time, query_timeout_mins, config_file_path=config_file_path
+                )
+                job_id = solver._exec_job(payload, log_to_console=log_to_console, query_timeout_mins=remaining_timeout_minutes)
             else:
                 raise e
 
         # 3. Extract result.
+        remaining_timeout_minutes = calc_remaining_timeout_minutes(
+            start_time, query_timeout_mins, config_file_path=config_file_path
+        )
         res = self.graph.exec_raw(
             textwrap.dedent(f"""
             ic result_not_empty("Solver result is empty.") requires not empty(result)
@@ -273,6 +296,7 @@ class SolverModel:
             def output[:solver_error]: {scope}solve_output[:"i_{self.solve_index}", :error]
             """),
             readonly=False,
+            query_timeout_mins=remaining_timeout_minutes,
         )
         errors = []
         for result in res.results:
@@ -506,10 +530,10 @@ class Solver:
         self.provider = Provider(resources=resources)
         self.solver_name = solver_name.lower()
 
-        rai_config = self.provider.resources.config
+        self.rai_config = self.provider.resources.config
         settings: dict[str, Any] = {}
-        if "experimental" in rai_config:
-            exp_config = rai_config.get("experimental", {})
+        if "experimental" in self.rai_config:
+            exp_config = self.rai_config.get("experimental", {})
             if isinstance(exp_config, dict):
                 if "solvers" in exp_config:
                     settings = exp_config["solvers"].copy()
@@ -629,23 +653,29 @@ class Solver:
 
             self.engine = engine
 
-    def _exec_job_async(self, payload):
+    def _exec_job_async(self, payload, query_timeout_mins: Optional[int]=None):
         payload_json = json.dumps(payload)
         engine_name = self.engine["name"]
-        res = self.provider.resources._exec(
-            textwrap.dedent(f"""
+        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):
+    def _exec_job(self, payload, log_to_console=True, query_timeout_mins: Optional[int]=None):
         # 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)
+            job_id = self._exec_job_async(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)
             polling_state = PollingState(job_id, "", False, log_to_console)

relationalai/semantics/devtools/extract_lqp.py

@@ -9,7 +9,7 @@ from contextlib import contextmanager
 
 from relationalai.clients.snowflake import Resources as snowflake_api
 from relationalai.semantics.internal import internal
-from typing import Dict
+from typing import Dict, Optional
 
 def main():
     parser = argparse.ArgumentParser(description="Extract LQP requests to run locally")
@@ -44,6 +44,7 @@ def instrumented_exec_rai_app(captured_calls, call_counter):
         request_headers=None,
         bypass_index=False,
         language: str = "rel",
+        query_timeout_mins: Optional[int] = None,
     ):
         result = original_exec_rai(
             self,
@@ -56,6 +57,7 @@ def instrumented_exec_rai_app(captured_calls, call_counter):
             request_headers=request_headers,
             bypass_index=bypass_index,
             language=language,
+            query_timeout_mins=query_timeout_mins,
         )
 
         call_counter[0] += 1
@@ -67,6 +69,7 @@ def instrumented_exec_rai_app(captured_calls, call_counter):
             "readonly": readonly,
             "nowait_durable": nowait_durable,
             "language": language,
+            "timeout_mins": query_timeout_mins,
             "raw_code": raw_code,
         }
         captured_calls.append(exec_call_json)