relationalai 0.11.3__py3-none-any.whl → 0.12.0__py3-none-any.whl

relationalai/clients/config.py

@@ -468,6 +468,13 @@ class Config():
             if not self.file_path:
                 self.file_path = "__inline__"
         self._handle_snowflake_fallback_configurations()
+        # Check if Azure platform is being used without the legacy dependency
+        if self.get("platform", "") == "azure":
+            try:
+                import railib # noqa
+            except ImportError:
+                from relationalai.errors import AzureLegacyDependencyMissingException
+                raise AzureLegacyDependencyMissingException() from None
 
     def fetch(self, profile:str|None=None):
         from relationalai.environments import runtime_env, TerminalEnvironment

relationalai/clients/direct_access_client.py

@@ -0,0 +1,113 @@
+from __future__ import annotations
+
+import requests
+from dataclasses import dataclass
+from urllib.parse import urlencode, quote
+from requests.adapters import HTTPAdapter
+from urllib3.util.retry import Retry
+from typing import Any, Dict, Optional, Tuple
+
+from relationalai.auth.token_handler import TokenHandler
+from relationalai.clients.config import Config
+from relationalai.clients.util import get_pyrel_version
+from relationalai import debugging
+from relationalai.tools.constants import Generation
+from relationalai.environments import runtime_env, SnowbookEnvironment
+
+@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)

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.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
 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"
 
@@ -2957,104 +2966,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 +2979,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
@@ -3140,7 +3058,12 @@ class DirectAccessResources(Resources):
             try:
                 response = _send_request()
                 if response.status_code != 200:
-                    message = response.json().get("message", "")
+                    try:
+                        message = response.json().get("message", "")
+                    except requests.exceptions.JSONDecodeError:
+                        raise ResponseStatusException(
+                            f"Failed to parse error response from endpoint {endpoint}.", response
+                        )
 
                     # fix engine on engine error and retry
                     if _is_engine_issue(message):
@@ -3305,13 +3228,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"