relationalai 0.11.4__py3-none-any.whl → 0.12.1__py3-none-any.whl

relationalai/clients/snowflake.py

@@ -28,15 +28,11 @@ import requests
 import snowflake.connector
 import pyarrow as pa
 
-from dataclasses import dataclass
 from snowflake.snowpark import Session
 from snowflake.snowpark.context import get_active_session
 from . import result_helpers
 from .. import debugging
 from typing import Any, Dict, Iterable, Optional, Tuple, List, Literal, Union, cast
-from urllib.parse import urlencode, quote
-from requests.adapters import HTTPAdapter
-from urllib3.util.retry import Retry
 
 from pandas import DataFrame
 
@@ -44,10 +40,11 @@ from ..tools.cli_controls import Spinner
 from ..clients.types import AvailableModel, EngineState, Import, ImportSource, ImportSourceTable, ImportsStatus, SourceInfo, TransactionAsyncResponse
 from ..clients.config import Config, ConfigStore, ENDPOINT_FILE
 from ..clients.client import Client, ExportParams, ProviderBase, ResourcesBase
-from ..clients.util import IdentityParser, escape_for_f_string, get_pyrel_version, get_with_retries, poll_with_specified_overhead, safe_json_loads, sanitize_module_name, scrub_exception, wrap_with_request_id, ms_to_timestamp
+from ..clients.direct_access_client import DirectAccessClient
+from ..clients.util import IdentityParser, escape_for_f_string, get_pyrel_version, get_with_retries, poll_with_specified_overhead, safe_json_loads, sanitize_module_name, scrub_exception, wrap_with_request_id, ms_to_timestamp, normalize_datetime
 from ..environments import runtime_env, HexEnvironment, SnowbookEnvironment
 from .. import dsl, rel, metamodel as m
-from ..errors import DuoSecurityFailed, EngineProvisioningFailed, EngineNameValidationException, EngineNotFoundException, EnginePending, EngineSizeMismatchWarning, EngineResumeFailed, Errors, InvalidAliasError, InvalidEngineSizeError, InvalidSourceTypeWarning, RAIAbortedTransactionError, RAIException, HexSessionException, SnowflakeAppMissingException, SnowflakeChangeTrackingNotEnabledException, SnowflakeDatabaseException, SnowflakeImportMissingException, SnowflakeInvalidSource, SnowflakeMissingConfigValuesException, SnowflakeProxyAPIDeprecationWarning, SnowflakeProxySourceError, SnowflakeRaiAppNotStarted, ModelNotFoundException, UnknownSourceWarning, ResponseStatusException, RowsDroppedFromTargetTableWarning
+from ..errors import DuoSecurityFailed, EngineProvisioningFailed, EngineNameValidationException, EngineNotFoundException, EnginePending, EngineSizeMismatchWarning, EngineResumeFailed, Errors, InvalidAliasError, InvalidEngineSizeError, InvalidSourceTypeWarning, RAIAbortedTransactionError, RAIException, HexSessionException, SnowflakeAppMissingException, SnowflakeChangeTrackingNotEnabledException, SnowflakeDatabaseException, SnowflakeImportMissingException, SnowflakeInvalidSource, SnowflakeMissingConfigValuesException, SnowflakeProxyAPIDeprecationWarning, SnowflakeProxySourceError, SnowflakeRaiAppNotStarted, ModelNotFoundException, UnknownSourceWarning, ResponseStatusException, RowsDroppedFromTargetTableWarning, QueryTimeoutExceededException
 from concurrent.futures import ThreadPoolExecutor
 from datetime import datetime, date, timedelta
 from snowflake.snowpark.types import StringType, StructField, StructType
@@ -92,6 +89,8 @@ TERMINAL_TXN_STATES = ["COMPLETED", "ABORTED"]
 
 DUO_TEXT = "duo security"
 
+TXN_ABORT_REASON_TIMEOUT = "transaction timeout"
+
 #--------------------------------------------------
 # Helpers
 #--------------------------------------------------
@@ -1307,7 +1306,7 @@ Otherwise, remove it from your '{profile}' configuration profile.
             response = self._exec(f"CALL {APP_NAME}.api.get_transaction('{txn_id}',{headers});")
             assert response, f"No results from get_transaction('{txn_id}')"
 
-        response_row = next(iter(response))
+        response_row = next(iter(response)).asDict()
         status: str = response_row['STATE']
 
         # remove the transaction from the pending list if it's completed or aborted
@@ -1315,6 +1314,16 @@ Otherwise, remove it from your '{profile}' configuration profile.
             if txn_id in self._pending_transactions:
                 self._pending_transactions.remove(txn_id)
 
+        if status == "ABORTED" and response_row.get("ABORT_REASON", "") == TXN_ABORT_REASON_TIMEOUT:
+            config_file_path = getattr(self.config, 'file_path', None)
+            # todo: use the timeout returned alongside the transaction as soon as it's exposed
+            timeout_mins = int(self.config.get("query_timeout_mins", DEFAULT_QUERY_TIMEOUT_MINS) or DEFAULT_QUERY_TIMEOUT_MINS)
+            raise QueryTimeoutExceededException(
+                timeout_mins=timeout_mins,
+                query_id=txn_id,
+                config_file_path=config_file_path,
+            )
+
         # @TODO: Find some way to tunnel the ABORT_REASON out. Azure doesn't have this, but it's handy
         return status == "COMPLETED" or status == "ABORTED"
 
@@ -1858,7 +1867,7 @@ Otherwise, remove it from your '{profile}' configuration profile.
         except Exception as e:
             err_message = str(e).lower()
             if _is_engine_issue(err_message):
-                self.auto_create_engine(engine)
+                self.auto_create_engine(engine, headers=headers)
                 self._exec_async_v2(
                     database, engine, raw_code_b64, inputs, readonly, nowait_durable,
                     headers=headers, bypass_index=bypass_index, language='lqp',
@@ -1898,7 +1907,7 @@ Otherwise, remove it from your '{profile}' configuration profile.
         except Exception as e:
             err_message = str(e).lower()
             if _is_engine_issue(err_message):
-                self.auto_create_engine(engine)
+                self.auto_create_engine(engine, headers=headers)
                 return self._exec_async_v2(
                     database,
                     engine,
@@ -1961,9 +1970,9 @@ Otherwise, remove it from your '{profile}' configuration profile.
                     if use_graph_index:
                         # 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])
+                            res = self._exec(f"call {APP_NAME}.api.exec_into_table(?, ?, ?, ?, ?, NULL, ?, {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_table(?, ?, ?, ?, ?, ?, ?);", [database, engine, raw_code, output_table, readonly, nowait_durable, skip_invalid_data])
+                             res = self._exec(f"call {APP_NAME}.api.exec_into_table(?, ?, ?, ?, ?, NULL, ?, {headers}, ?);", [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)
@@ -2038,9 +2047,10 @@ Otherwise, remove it from your '{profile}' configuration profile.
         app_name = self.get_app_name()
 
         source_types = dict[str, SourceInfo]()
-        partitioned_sources: dict[str, dict[str, list[str]]] = defaultdict(
+        partitioned_sources: dict[str, dict[str, list[dict[str, str]]]] = defaultdict(
             lambda: defaultdict(list)
         )
+        fqn_to_parts: dict[str, tuple[str, str, str]] = {}
 
         for source in sources:
             parser = IdentityParser(source, True)
@@ -2048,82 +2058,219 @@ 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": "", "columns_hash": None})
-            partitioned_sources[db][schema].append(entity)
-
-        # TODO: Move to NA layer
-        query = (
-            " UNION ALL ".join(
-                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,
-                        CREATED_AT FROM {app_name}.api.data_streams
-                    WHERE RAI_DATABASE = '{PYREL_ROOT_DB}'
-                ) ds
-                RIGHT JOIN (
+            source_types[identity] = cast(
+                SourceInfo,
+                {
+                    "type": None,
+                    "state": "",
+                    "columns_hash": None,
+                    "table_created_at": None,
+                    "stream_created_at": None,
+                    "last_ddl": None,
+                },
+            )
+            partitioned_sources[db][schema].append({"entity": entity, "identity": identity})
+            fqn_to_parts[identity] = (db, schema, entity)
+
+        if not partitioned_sources:
+            return source_types
+
+        state_queries: list[str] = []
+        for db, schemas in partitioned_sources.items():
+            select_rows: list[str] = []
+            for schema, tables in schemas.items():
+                for table_info in tables:
+                    select_rows.append(
+                        "SELECT "
+                        f"{IdentityParser.to_sql_value(db)} AS catalog_name, "
+                        f"{IdentityParser.to_sql_value(schema)} AS schema_name, "
+                        f"{IdentityParser.to_sql_value(table_info['entity'])} AS table_name"
+                    )
+
+            if not select_rows:
+                continue
+
+            target_entities_clause = "\n                        UNION ALL\n                        ".join(select_rows)
+            # Main query:
+            #   1. Enumerate the target tables via target_entities.
+            #   2. Pull their metadata (last_altered, type) from INFORMATION_SCHEMA.TABLES.
+            #   3. Look up the most recent stream activity for those FQNs only.
+            #   4. Capture creation timestamps and use last_ddl vs created_at to classify each target,
+            #      so we mark tables as stale when they were recreated even if column hashes still match.
+            state_queries.append(
+                f"""WITH target_entities AS (
+                        {target_entities_clause}
+                    ),
+                    table_info AS (
+                        SELECT
+                            {app_name}.api.normalize_fq_ids(
+                                ARRAY_CONSTRUCT(
+                                    CASE
+                                        WHEN t.table_catalog = UPPER(t.table_catalog) THEN t.table_catalog
+                                        ELSE '"' || t.table_catalog || '"'
+                                    END || '.' ||
+                                    CASE
+                                        WHEN t.table_schema = UPPER(t.table_schema) THEN t.table_schema
+                                        ELSE '"' || t.table_schema || '"'
+                                    END || '.' ||
+                                    CASE
+                                        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', t.last_altered) AS last_ddl,
+                            CONVERT_TIMEZONE('UTC', t.created) AS table_created_at,
+                            t.table_type AS kind
+                        FROM {db}.INFORMATION_SCHEMA.tables t
+                        JOIN target_entities te
+                            ON t.table_catalog = te.catalog_name
+                            AND t.table_schema = te.schema_name
+                            AND t.table_name = te.table_name
+                    ),
+                    stream_activity AS (
+                        SELECT
+                            sa.fqn,
+                            MAX(sa.created_at) AS created_at
+                        FROM (
+                            SELECT
+                                {app_name}.api.normalize_fq_ids(ARRAY_CONSTRUCT(fq_object_name))[0]:identifier::string AS fqn,
+                                created_at
+                            FROM {app_name}.api.data_streams
+                            WHERE rai_database = '{PYREL_ROOT_DB}'
+                        ) sa
+                        JOIN table_info ti
+                            ON sa.fqn = ti.fqn
+                        GROUP BY sa.fqn
+                    )
                     SELECT
-                        (SELECT {app_name}.api.normalize_fq_ids(
-                            ARRAY_CONSTRUCT(
-                                CASE
-                                    WHEN t.TABLE_CATALOG = UPPER(t.TABLE_CATALOG) THEN t.TABLE_CATALOG
-                                    ELSE '"' || t.TABLE_CATALOG || '"'
-                                END || '.' ||
-                                CASE
-                                    WHEN t.TABLE_SCHEMA = UPPER(t.TABLE_SCHEMA) THEN t.TABLE_SCHEMA
-                                    ELSE '"' || t.TABLE_SCHEMA || '"'
-                                END || '.' ||
-                                CASE
-                                    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,
-                        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()
+                        ti.fqn,
+                        ti.kind,
+                        ti.last_ddl,
+                        ti.table_created_at,
+                        sa.created_at AS stream_created_at,
+                        IFF(
+                            DATEDIFF(second, sa.created_at::timestamp, ti.last_ddl::timestamp) > 0,
+                            'STALE',
+                            'CURRENT'
+                        ) AS state
+                    FROM table_info ti
+                    LEFT JOIN stream_activity sa
+                        ON sa.fqn = ti.fqn
+                """
             )
-            + ";"
+
+        stale_fqns: list[str] = []
+        for state_query in state_queries:
+            for row in self._exec(state_query):
+                row_dict = row.as_dict() if hasattr(row, "as_dict") else dict(row)
+                row_fqn = row_dict["FQN"]
+                parser = IdentityParser(row_fqn, True)
+                fqn = parser.identity
+                assert fqn, f"Error parsing returned FQN: {row_fqn}"
+
+                source_types[fqn]["type"] = (
+                    "TABLE" if row_dict["KIND"] == "BASE TABLE" else row_dict["KIND"]
+                )
+                source_types[fqn]["state"] = row_dict["STATE"]
+                source_types[fqn]["last_ddl"] = normalize_datetime(row_dict.get("LAST_DDL"))
+                source_types[fqn]["table_created_at"] = normalize_datetime(row_dict.get("TABLE_CREATED_AT"))
+                source_types[fqn]["stream_created_at"] = normalize_datetime(row_dict.get("STREAM_CREATED_AT"))
+                if row_dict["STATE"] == "STALE":
+                    stale_fqns.append(fqn)
+
+        if not stale_fqns:
+            return source_types
+
+        # We batch stale tables by database/schema so each Snowflake query can hash
+        # multiple objects at once instead of issuing one statement per table.
+        stale_partitioned: dict[str, dict[str, list[dict[str, str]]]] = defaultdict(
+            lambda: defaultdict(list)
         )
+        for fqn in stale_fqns:
+            db, schema, table = fqn_to_parts[fqn]
+            stale_partitioned[db][schema].append({"table": table, "identity": fqn})
+
+        # Build one hash query per database, grouping schemas/tables inside so we submit
+        # at most a handful of set-based statements to Snowflake.
+        for db, schemas in stale_partitioned.items():
+            column_select_rows: list[str] = []
+            for schema, tables in schemas.items():
+                for table_info in tables:
+                    # Build the literal rows for this db/schema so we can join back
+                    # against INFORMATION_SCHEMA.COLUMNS in a single statement.
+                    column_select_rows.append(
+                        "SELECT "
+                        f"{IdentityParser.to_sql_value(db)} AS catalog_name, "
+                        f"{IdentityParser.to_sql_value(schema)} AS schema_name, "
+                        f"{IdentityParser.to_sql_value(table_info['table'])} AS table_name"
+                    )
 
-        for row in self._exec(query):
-            row_fqn = row["FQN"]
-            parser = IdentityParser(row_fqn, True)
-            fqn = parser.identity
-            assert fqn, f"Error parsing returned FQN: {row_fqn}"
+            if not column_select_rows:
+                continue
 
-            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"]
+            target_entities_clause = "\n                        UNION ALL\n                        ".join(column_select_rows)
+            # Main query: compute deterministic column hashes for every stale table
+            # in this database/schema batch so we can compare schemas without a round trip per table.
+            column_query = f"""WITH target_entities AS (
+                        {target_entities_clause}
+                    ),
+                    column_info AS (
+                        SELECT
+                            {app_name}.api.normalize_fq_ids(
+                                ARRAY_CONSTRUCT(
+                                    CASE
+                                        WHEN c.table_catalog = UPPER(c.table_catalog) THEN c.table_catalog
+                                        ELSE '"' || c.table_catalog || '"'
+                                    END || '.' ||
+                                    CASE
+                                        WHEN c.table_schema = UPPER(c.table_schema) THEN c.table_schema
+                                        ELSE '"' || c.table_schema || '"'
+                                    END || '.' ||
+                                    CASE
+                                        WHEN c.table_name = UPPER(c.table_name) THEN c.table_name
+                                        ELSE '"' || c.table_name || '"'
+                                    END
+                                )
+                            )[0]:identifier::string AS fqn,
+                            c.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 AS type_signature,
+                            IFF(c.is_nullable = 'YES', 'YES', 'NO') AS nullable_flag
+                        FROM {db}.INFORMATION_SCHEMA.COLUMNS c
+                        JOIN target_entities te
+                            ON c.table_catalog = te.catalog_name
+                            AND c.table_schema = te.schema_name
+                            AND c.table_name = te.table_name
+                    )
+                    SELECT
+                        fqn,
+                        HEX_ENCODE(
+                            HASH_AGG(
+                                HASH(
+                                    column_name,
+                                    type_signature,
+                                    nullable_flag
+                                )
+                            )
+                        ) AS columns_hash
+                    FROM column_info
+                    GROUP BY fqn
+                """
+
+            for row in self._exec(column_query):
+                row_fqn = row["FQN"]
+                parser = IdentityParser(row_fqn, True)
+                fqn = parser.identity
+                assert fqn, f"Error parsing returned FQN: {row_fqn}"
+                source_types[fqn]["columns_hash"] = row["COLUMNS_HASH"]
 
         return source_types
 
@@ -2133,12 +2280,13 @@ Otherwise, remove it from your '{profile}' configuration profile.
         invalid_sources = {}
         source_references = []
         for source, info in source_info.items():
-            if info['type'] is None:
+            source_type = info.get("type")
+            if source_type is None:
                 missing_sources.append(source)
-            elif info['type'] not in ("TABLE", "VIEW"):
-                invalid_sources[source] = info['type']
+            elif source_type not in ("TABLE", "VIEW"):
+                invalid_sources[source] = source_type
             else:
-                source_references.append(f"{app_name}.api.object_reference('{info['type']}', '{source}')")
+                source_references.append(f"{app_name}.api.object_reference('{source_type}', '{source}')")
 
         if missing_sources:
             current_role = self.get_sf_session().get_current_role()
@@ -2957,104 +3105,6 @@ def Graph(
 #--------------------------------------------------
 # Note: All direct access components should live in a separate file
 
-@dataclass
-class Endpoint:
-    method: str
-    endpoint: str
-
-class DirectAccessClient:
-    """
-    DirectAccessClient is a client for direct service access without service function calls.
-    """
-
-    def __init__(self, config: Config, token_handler: TokenHandler, service_endpoint: str, generation: Optional[Generation] = None):
-        self._config: Config = config
-        self._token_handler: TokenHandler = token_handler
-        self.service_endpoint: str = service_endpoint
-        self.generation: Optional[Generation] = generation
-        self._is_snowflake_notebook = isinstance(runtime_env, SnowbookEnvironment)
-        self.endpoints: Dict[str, Endpoint] = {
-            "create_txn": Endpoint(method="POST", endpoint="/v1alpha1/transactions"),
-            "get_txn": Endpoint(method="GET", endpoint="/v1alpha1/transactions/{txn_id}"),
-            "get_txn_artifacts": Endpoint(method="GET", endpoint="/v1alpha1/transactions/{txn_id}/artifacts"),
-            "get_txn_problems": Endpoint(method="GET", endpoint="/v1alpha1/transactions/{txn_id}/problems"),
-            "get_txn_events": Endpoint(method="GET", endpoint="/v1alpha1/transactions/{txn_id}/events/{stream_name}"),
-            "get_package_versions": Endpoint(method="GET", endpoint="/v1alpha1/databases/{db_name}/package_versions"),
-            "get_model_package_versions": Endpoint(method="POST", endpoint="/v1alpha1/models/get_package_versions"),
-            "create_db": Endpoint(method="POST", endpoint="/v1alpha1/databases"),
-            "get_db": Endpoint(method="GET", endpoint="/v1alpha1/databases"),
-            "delete_db": Endpoint(method="DELETE", endpoint="/v1alpha1/databases/{db_name}"),
-            "release_index": Endpoint(method="POST", endpoint="/v1alpha1/index/release"),
-            "list_engines": Endpoint(method="GET", endpoint="/v1alpha1/engines"),
-            "get_engine": Endpoint(method="GET", endpoint="/v1alpha1/engines/{engine_type}/{engine_name}"),
-            "create_engine": Endpoint(method="POST", endpoint="/v1alpha1/engines/{engine_type}"),
-            "delete_engine": Endpoint(method="DELETE", endpoint="/v1alpha1/engines/{engine_type}/{engine_name}"),
-            "suspend_engine": Endpoint(method="POST", endpoint="/v1alpha1/engines/{engine_type}/{engine_name}/suspend"),
-            "resume_engine": Endpoint(method="POST", endpoint="/v1alpha1/engines/{engine_type}/{engine_name}/resume_async"),
-            "prepare_index": Endpoint(method="POST", endpoint="/v1alpha1/index/prepare"),
-        }
-        self.http_session = self._create_retry_session()
-
-    def _create_retry_session(self) -> requests.Session:
-        http_session = requests.Session()
-        retries = Retry(
-            total=3,
-            backoff_factor=0.3,
-            status_forcelist=[500, 502, 503, 504],
-            allowed_methods=frozenset({"GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"}),
-            raise_on_status=False
-        )
-        adapter = HTTPAdapter(max_retries=retries)
-        http_session.mount("http://", adapter)
-        http_session.mount("https://", adapter)
-        http_session.headers.update({"Connection": "keep-alive"})
-        return http_session
-
-    def request(
-        self,
-        endpoint: str,
-        payload: Dict[str, Any] | None = None,
-        headers: Dict[str, str] | None = None,
-        path_params: Dict[str, str] | None = None,
-        query_params: Dict[str, str] | None = None,
-    ) -> requests.Response:
-        """
-        Send a request to the service endpoint.
-        """
-        url, method = self._prepare_url(endpoint, path_params, query_params)
-        request_headers = self._prepare_headers(headers)
-        return self.http_session.request(method, url, json=payload, headers=request_headers)
-
-    def _prepare_url(self, endpoint: str, path_params: Dict[str, str] | None = None, query_params: Dict[str, str] | None = None) -> Tuple[str, str]:
-        try:
-            ep = self.endpoints[endpoint]
-        except KeyError:
-            raise ValueError(f"Invalid endpoint: {endpoint}. Available endpoints: {list(self.endpoints.keys())}")
-        url = f"{self.service_endpoint}{ep.endpoint}"
-        if path_params:
-            escaped_path_params = {k: quote(v, safe='') for k, v in path_params.items()}
-            url = url.format(**escaped_path_params)
-        if query_params:
-            url += '?' + urlencode(query_params)
-        return url, ep.method
-
-    def _prepare_headers(self, headers: Dict[str, str] | None) -> Dict[str, str]:
-        request_headers = {}
-        if headers:
-            request_headers.update(headers)
-        # Authorization tokens are not needed in a snowflake notebook environment
-        if not self._is_snowflake_notebook:
-            request_headers["Authorization"] = f'Snowflake Token="{self._token_handler.get_ingress_token(self.service_endpoint)}"'
-            # needed for oauth, does no harm for other authentication methods
-            request_headers["X-SF-SPCS-Authentication-Method"] = 'OAUTH'
-            request_headers["Content-Type"] = 'application/x-www-form-urlencoded'
-        request_headers["Accept"] = "application/json"
-
-        request_headers["user-agent"] = get_pyrel_version(self.generation)
-        request_headers["pyrel_program_id"] = debugging.get_program_span_id() or ""
-
-        return debugging.add_current_propagation_headers(request_headers)
-
 class DirectAccessResources(Resources):
     """
     Resources class for Direct Service Access avoiding Snowflake service functions.
@@ -3068,7 +3118,14 @@ class DirectAccessResources(Resources):
         reset_session: bool = False,
         generation: Optional[Generation] = None,
     ):
-        super().__init__(generation=generation, profile=profile, config=config, connection=connection, dry_run=dry_run)
+        super().__init__(
+            generation=generation,
+            profile=profile,
+            config=config,
+            connection=connection,
+            reset_session=reset_session,
+            dry_run=dry_run,
+        )
         self._endpoint_info = ConfigStore(ENDPOINT_FILE)
         self._service_endpoint = ""
         self._direct_access_client = None
@@ -3127,6 +3184,7 @@ class DirectAccessResources(Resources):
         headers: Dict[str, str] | None = None,
         path_params: Dict[str, str] | None = None,
         query_params: Dict[str, str] | None = None,
+        skip_auto_create: bool = False,
     ) -> requests.Response:
         with debugging.span("direct_access_request"):
             def _send_request():
@@ -3148,7 +3206,8 @@ class DirectAccessResources(Resources):
                         )
 
                     # fix engine on engine error and retry
-                    if _is_engine_issue(message):
+                    # Skip auto-retry if skip_auto_create is True to avoid recursion
+                    if _is_engine_issue(message) and not skip_auto_create:
                         engine = payload.get("engine_name", "") if payload else ""
                         self.auto_create_engine(engine)
                         response = _send_request()
@@ -3310,13 +3369,24 @@ class DirectAccessResources(Resources):
             assert response, f"No results from get_transaction('{txn_id}')"
 
         response_content = response.json()
-        status: str = response_content["transaction"]['state']
+        transaction = response_content["transaction"]
+        status: str = transaction['state']
 
         # remove the transaction from the pending list if it's completed or aborted
         if status in ["COMPLETED", "ABORTED"]:
             if txn_id in self._pending_transactions:
                 self._pending_transactions.remove(txn_id)
 
+        if status == "ABORTED" and transaction.get("abort_reason", "") == TXN_ABORT_REASON_TIMEOUT:
+            config_file_path = getattr(self.config, 'file_path', None)
+            timeout_ms = int(transaction.get("timeout_ms", 0))
+            timeout_mins = timeout_ms // 60000 if timeout_ms > 0 else int(self.config.get("query_timeout_mins", DEFAULT_QUERY_TIMEOUT_MINS) or DEFAULT_QUERY_TIMEOUT_MINS)
+            raise QueryTimeoutExceededException(
+                timeout_mins=timeout_mins,
+                query_id=txn_id,
+                config_file_path=config_file_path,
+            )
+
         # @TODO: Find some way to tunnel the ABORT_REASON out. Azure doesn't have this, but it's handy
         return status == "COMPLETED" or status == "ABORTED"
 
@@ -3502,7 +3572,7 @@ class DirectAccessResources(Resources):
         return sorted(engines, key=lambda x: x["name"])
 
     def get_engine(self, name: str):
-        response = self.request("get_engine", path_params={"engine_name": name, "engine_type": "logic"})
+        response = self.request("get_engine", path_params={"engine_name": name, "engine_type": "logic"}, skip_auto_create=True)
         if response.status_code == 404: # engine not found return 404
             return None
         elif response.status_code != 200:
@@ -3549,6 +3619,7 @@ class DirectAccessResources(Resources):
             payload=payload,
             path_params={"engine_type": "logic"},
             headers=headers,
+            skip_auto_create=True,
         )
         if response.status_code != 200:
             raise ResponseStatusException(
@@ -3560,6 +3631,7 @@ class DirectAccessResources(Resources):
             "delete_engine",
             path_params={"engine_name": name, "engine_type": "logic"},
             headers=headers,
+            skip_auto_create=True,
         )
         if response.status_code != 200:
             raise ResponseStatusException(
@@ -3570,6 +3642,7 @@ class DirectAccessResources(Resources):
         response = self.request(
             "suspend_engine",
             path_params={"engine_name": name, "engine_type": "logic"},
+            skip_auto_create=True,
         )
         if response.status_code != 200:
             raise ResponseStatusException(
@@ -3581,6 +3654,7 @@ class DirectAccessResources(Resources):
             "resume_engine",
             path_params={"engine_name": name, "engine_type": "logic"},
             headers=headers,
+            skip_auto_create=True,
         )
         if response.status_code != 200:
             raise ResponseStatusException(

relationalai/clients/types.py

@@ -38,10 +38,13 @@ class EngineState(TypedDict):
     auto_suspend: int|None
     suspends_at: datetime|None
 
-class SourceInfo(TypedDict):
+class SourceInfo(TypedDict, total=False):
     type: str|None
     state: str
     columns_hash: str|None
+    table_created_at: datetime|None
+    stream_created_at: datetime|None
+    last_ddl: datetime|None
     source: str