relationalai 0.11.1__tar.gz → 0.11.3__tar.gz

{relationalai-0.11.1 → relationalai-0.11.3}/.gitignore

@@ -37,3 +37,4 @@ relationalai*.zip
 logs/
 .envrc
 metadata.json
+dev-notes/

{relationalai-0.11.1 → relationalai-0.11.3}/PKG-INFO

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

{relationalai-0.11.1 → relationalai-0.11.3}/pyproject.toml

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

{relationalai-0.11.1 → relationalai-0.11.3}/src/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-0.11.1 → relationalai-0.11.3}/src/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-0.11.1 → relationalai-0.11.3}/src/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-0.11.1 → relationalai-0.11.3}/src/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
@@ -1997,7 +2048,7 @@ Otherwise, remove it from your '{profile}' configuration profile.
             assert len(parsed) == 4, f"Invalid source: {source}"
             db, schema, entity, identity = parsed
             assert db and schema and entity and identity, f"Invalid source: {source}"
-            source_types[identity] = cast(SourceInfo, {"type": None, "state": ""})
+            source_types[identity] = cast(SourceInfo, {"type": None, "state": "", "columns_hash": None})
             partitioned_sources[db][schema].append(entity)
 
         # TODO: Move to NA layer
@@ -2006,6 +2057,7 @@ Otherwise, remove it from your '{profile}' configuration profile.
                 f"""SELECT
                     inf.FQN,
                     inf.KIND,
+                    inf.COLUMNS_HASH,
                     IFF(DATEDIFF(second, ds.created_at::TIMESTAMP, inf.LAST_DDL::TIMESTAMP) > 0, 'STALE', 'CURRENT') AS STATE
                 FROM (
                     SELECT (SELECT {app_name}.api.normalize_fq_ids(ARRAY_CONSTRUCT(FQ_OBJECT_NAME))[0]:identifier::string) as FQ_OBJECT_NAME,
@@ -2017,26 +2069,45 @@ Otherwise, remove it from your '{profile}' configuration profile.
                         (SELECT {app_name}.api.normalize_fq_ids(
                             ARRAY_CONSTRUCT(
                                 CASE
-                                    WHEN TABLE_CATALOG = UPPER(TABLE_CATALOG) THEN TABLE_CATALOG
-                                    ELSE '"' || TABLE_CATALOG || '"'
+                                    WHEN t.TABLE_CATALOG = UPPER(t.TABLE_CATALOG) THEN t.TABLE_CATALOG
+                                    ELSE '"' || t.TABLE_CATALOG || '"'
                                 END || '.' ||
                                 CASE
-                                    WHEN TABLE_SCHEMA = UPPER(TABLE_SCHEMA) THEN TABLE_SCHEMA
-                                    ELSE '"' || TABLE_SCHEMA || '"'
+                                    WHEN t.TABLE_SCHEMA = UPPER(t.TABLE_SCHEMA) THEN t.TABLE_SCHEMA
+                                    ELSE '"' || t.TABLE_SCHEMA || '"'
                                 END || '.' ||
                                 CASE
-                                    WHEN TABLE_NAME = UPPER(TABLE_NAME) THEN TABLE_NAME
-                                    ELSE '"' || TABLE_NAME || '"'
+                                    WHEN t.TABLE_NAME = UPPER(t.TABLE_NAME) THEN t.TABLE_NAME
+                                    ELSE '"' || t.TABLE_NAME || '"'
                                 END
                             )
                         )[0]:identifier::string) as FQN,
                         CONVERT_TIMEZONE('UTC', LAST_DDL) AS LAST_DDL,
-                        TABLE_TYPE as KIND
-                    FROM {db}.INFORMATION_SCHEMA.TABLES
-                    WHERE TABLE_CATALOG = {IdentityParser.to_sql_value(db)} AND ({" OR ".join(
-                        f"(TABLE_SCHEMA = {IdentityParser.to_sql_value(schema)} AND TABLE_NAME IN ({','.join(f'{IdentityParser.to_sql_value(table)}' for table in tables)}))"
+                        TABLE_TYPE as KIND,
+                        SHA2(LISTAGG(
+                            COLUMN_NAME ||
+                            CASE
+                                WHEN c.NUMERIC_PRECISION IS NOT NULL AND c.NUMERIC_SCALE IS NOT NULL
+                                    THEN c.DATA_TYPE || '(' || c.NUMERIC_PRECISION || ',' || c.NUMERIC_SCALE || ')'
+                                WHEN c.DATETIME_PRECISION IS NOT NULL
+                                    THEN c.DATA_TYPE || '(0,' || c.DATETIME_PRECISION || ')'
+                                WHEN c.CHARACTER_MAXIMUM_LENGTH IS NOT NULL
+                                    THEN c.DATA_TYPE || '(' || c.CHARACTER_MAXIMUM_LENGTH || ')'
+                                ELSE c.DATA_TYPE
+                            END ||
+                            IS_NULLABLE,
+                            ','
+                        ) WITHIN GROUP (ORDER BY COLUMN_NAME), 256) as COLUMNS_HASH
+                    FROM {db}.INFORMATION_SCHEMA.TABLES t
+                    JOIN {db}.INFORMATION_SCHEMA.COLUMNS c
+                        ON t.TABLE_CATALOG = c.TABLE_CATALOG
+                        AND t.TABLE_SCHEMA = c.TABLE_SCHEMA
+                        AND t.TABLE_NAME = c.TABLE_NAME
+                    WHERE t.TABLE_CATALOG = {IdentityParser.to_sql_value(db)} AND ({" OR ".join(
+                        f"(t.TABLE_SCHEMA = {IdentityParser.to_sql_value(schema)} AND t.TABLE_NAME IN ({','.join(f'{IdentityParser.to_sql_value(table)}' for table in tables)}))"
                         for schema, tables in schemas.items()
                     )})
+                    GROUP BY t.TABLE_CATALOG, t.TABLE_SCHEMA, t.TABLE_NAME, t.LAST_DDL, t.TABLE_TYPE
                 ) inf on inf.FQN = ds.FQ_OBJECT_NAME
             """
                 for db, schemas in partitioned_sources.items()
@@ -2051,6 +2122,7 @@ Otherwise, remove it from your '{profile}' configuration profile.
             assert fqn, f"Error parsing returned FQN: {row_fqn}"
 
             source_types[fqn]["type"] = "TABLE" if row["KIND"] == "BASE TABLE" else row["KIND"]
+            source_types[fqn]["columns_hash"] = row["COLUMNS_HASH"]
             source_types[fqn]["state"] = row["STATE"]
 
         return source_types
@@ -3007,10 +3079,13 @@ class DirectAccessResources(Resources):
         return self._retrieve_service_endpoint()
 
     def _retrieve_service_endpoint(self, enforce_update=False) -> str:
+        account = self.config.get("account")
+        app_name = self.config.get("rai_app_name")
+        service_endpoint_key = f"{account}.{app_name}.service_endpoint"
         if self._service_endpoint and not enforce_update:
             return self._service_endpoint
-        if self._endpoint_info.get("service_endpoint", "") and not enforce_update:
-            self._service_endpoint = str(self._endpoint_info.get("service_endpoint", ""))
+        if self._endpoint_info.get(service_endpoint_key, "") and not enforce_update:
+            self._service_endpoint = str(self._endpoint_info.get(service_endpoint_key, ""))
             return self._service_endpoint
 
         is_snowflake_notebook = isinstance(runtime_env, SnowbookEnvironment)
@@ -3022,7 +3097,7 @@ class DirectAccessResources(Resources):
         else:
             self._service_endpoint = f"https://{result[0]['SERVICE_ENDPOINT']}"
 
-        self._endpoint_info.set("service_endpoint", self._service_endpoint)
+        self._endpoint_info.set(service_endpoint_key, self._service_endpoint)
         # save the endpoint to `ENDPOINT_FILE` to avoid calling the endpoint with every
         # pyrel execution
         try:
@@ -3095,6 +3170,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 +3187,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-0.11.1 → relationalai-0.11.3}/src/relationalai/clients/types.py

@@ -41,6 +41,7 @@ class EngineState(TypedDict):
 class SourceInfo(TypedDict):
     type: str|None
     state: str
+    columns_hash: str|None
     source: str