relationalai 0.12.12__py3-none-any.whl → 0.13.0__py3-none-any.whl

relationalai/clients/{snowflake.py → resources/snowflake/snowflake.py}

@@ -1,10 +1,8 @@
 # pyright: reportUnusedExpression=false
 from __future__ import annotations
 import base64
-import decimal
 import importlib.resources
 import io
-from numbers import Number
 import re
 import json
 import time
@@ -14,15 +12,14 @@ import uuid
 import warnings
 import atexit
 import hashlib
+from dataclasses import dataclass
 
-
-from relationalai.auth.token_handler import TokenHandler
-from relationalai.clients.use_index_poller import DirectUseIndexPoller, UseIndexPoller
+from ....auth.token_handler import TokenHandler
 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, DEFAULT_QUERY_TIMEOUT_MINS, WAIT_FOR_STREAM_SYNC, Generation
-from .. import std
+from ....rel_utils import sanitize_identifier, to_fqn_relation_name
+from ....tools.constants import FIELD_PLACEHOLDER, SNOWFLAKE_AUTHS, USE_GRAPH_INDEX, DEFAULT_QUERY_TIMEOUT_MINS, WAIT_FOR_STREAM_SYNC, Generation
+from .... import std
 from collections import defaultdict
 import requests
 import snowflake.connector
@@ -30,33 +27,63 @@ import pyarrow as pa
 
 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 ... import result_helpers
+from .... import debugging
+from typing import Any, Dict, Iterable, Tuple, List, Literal, cast
 
 from pandas import DataFrame
 
-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, 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, QueryTimeoutExceededException
+from ....tools.cli_controls import Spinner
+from ...types import AvailableModel, EngineState, Import, ImportSource, ImportSourceTable, ImportsStatus, SourceInfo, TransactionAsyncResponse
+from ...config import Config
+from ...client import Client, ExportParams, ProviderBase, ResourcesBase
+from ...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, normalize_datetime
+from .util import (
+    collect_error_messages,
+    process_jinja_template,
+    type_to_sql,
+    type_to_snowpark,
+    sanitize_user_name as _sanitize_user_name,
+    normalize_params,
+    format_sproc_name,
+    is_azure_url,
+    is_container_runtime,
+    imports_to_dicts,
+    txn_list_to_dicts,
+    decrypt_artifact,
+)
+from ....environments import runtime_env, HexEnvironment, SnowbookEnvironment
+from .... import dsl, rel, metamodel as m
+from ....errors import EngineProvisioningFailed, EngineNameValidationException, Errors, InvalidAliasError, InvalidEngineSizeError, InvalidSourceTypeWarning, RAIException, HexSessionException, SnowflakeChangeTrackingNotEnabledException, SnowflakeDatabaseException, SnowflakeImportMissingException, SnowflakeInvalidSource, SnowflakeMissingConfigValuesException, SnowflakeProxyAPIDeprecationWarning, SnowflakeProxySourceError, ModelNotFoundException, UnknownSourceWarning, RowsDroppedFromTargetTableWarning, QueryTimeoutExceededException
 from concurrent.futures import ThreadPoolExecutor
-from datetime import datetime, date, timedelta
+from datetime import datetime, timedelta
 from snowflake.snowpark.types import StringType, StructField, StructType
+# Import error handlers and constants
+from .error_handlers import (
+    ErrorHandler,
+    DuoSecurityErrorHandler,
+    AppMissingErrorHandler,
+    DatabaseErrorsHandler,
+    EngineErrorsHandler,
+    ServiceNotStartedErrorHandler,
+    TransactionAbortedErrorHandler,
+)
+# Import engine state handlers
+from .engine_state_handlers import (
+    EngineStateHandler,
+    EngineContext,
+    SyncPendingStateHandler,
+    SyncSuspendedStateHandler,
+    SyncReadyStateHandler,
+    SyncGoneStateHandler,
+    SyncMissingEngineHandler,
+    AsyncPendingStateHandler,
+    AsyncSuspendedStateHandler,
+    AsyncReadyStateHandler,
+    AsyncGoneStateHandler,
+    AsyncMissingEngineHandler,
+)
 
-# warehouse-based snowflake notebooks currently don't have hazmat
-crypto_disabled = False
-try:
-    from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
-    from cryptography.hazmat.backends import default_backend
-    from cryptography.hazmat.primitives import padding
-except (ModuleNotFoundError, ImportError):
-    crypto_disabled = True
 
 #--------------------------------------------------
 # Constants
@@ -66,232 +93,44 @@ VALID_POOL_STATUS = ["ACTIVE", "IDLE", "SUSPENDED"]
 # transaction list and get return different fields (duration vs timings)
 LIST_TXN_SQL_FIELDS = ["id", "database_name", "engine_name", "state", "abort_reason", "read_only","created_by", "created_on", "finished_at", "duration"]
 GET_TXN_SQL_FIELDS = ["id", "database", "engine", "state", "abort_reason", "read_only","created_by", "created_on", "finished_at", "timings"]
-IMPORT_STREAM_FIELDS = ["ID", "CREATED_AT", "CREATED_BY", "STATUS", "REFERENCE_NAME", "REFERENCE_ALIAS", "FQ_OBJECT_NAME", "RAI_DATABASE",
-                        "RAI_RELATION", "DATA_SYNC_STATUS", "PENDING_BATCHES_COUNT", "NEXT_BATCH_STATUS", "NEXT_BATCH_UNLOADED_TIMESTAMP",
-                        "NEXT_BATCH_DETAILS", "LAST_BATCH_DETAILS", "LAST_BATCH_UNLOADED_TIMESTAMP", "CDC_STATUS"]
 VALID_ENGINE_STATES = ["READY", "PENDING"]
 
 # Cloud-specific engine sizes
 INTERNAL_ENGINE_SIZES = ["XS", "S", "M", "L"]
 ENGINE_SIZES_AWS = ["HIGHMEM_X64_S", "HIGHMEM_X64_M", "HIGHMEM_X64_L"]
 ENGINE_SIZES_AZURE = ["HIGHMEM_X64_S", "HIGHMEM_X64_M", "HIGHMEM_X64_SL"]
-
-FIELD_MAP = {
-    "database_name": "database",
-    "engine_name": "engine",
-}
-VALID_IMPORT_STATES = ["PENDING", "PROCESSING", "QUARANTINED", "LOADED"]
-ENGINE_ERRORS = ["engine is suspended", "create/resume", "engine not found", "no engines found", "engine was deleted"]
-ENGINE_NOT_READY_MSGS = ["engine is in pending", "engine is provisioning"]
-DATABASE_ERRORS = ["database not found"]
+# Note: ENGINE_ERRORS, ENGINE_NOT_READY_MSGS, DATABASE_ERRORS moved to util.py
 PYREL_ROOT_DB = 'pyrel_root_db'
 
 TERMINAL_TXN_STATES = ["COMPLETED", "ABORTED"]
 
-DUO_TEXT = "duo security"
-
 TXN_ABORT_REASON_TIMEOUT = "transaction timeout"
 
-#--------------------------------------------------
-# Helpers
-#--------------------------------------------------
-
-def process_jinja_template(template: str, indent_spaces = 0, **substitutions) -> str:
-    """Process a Jinja-like template.
-
-    Supports:
-    - Variable substitution {{ var }}
-    - Conditional blocks {% if condition %} ... {% endif %}
-    - For loops {% for item in items %} ... {% endfor %}
-    - Comments {# ... #}
-    - Whitespace control with {%- and -%}
-
-    Args:
-        template: The template string
-        indent_spaces: Number of spaces to indent the result
-        **substitutions: Variable substitutions
-    """
-
-    def evaluate_condition(condition: str, context: dict) -> bool:
-        """Safely evaluate a condition string using the context."""
-        # Replace variables with their values
-        for k, v in context.items():
-            if isinstance(v, str):
-                condition = condition.replace(k, f"'{v}'")
-            else:
-                condition = condition.replace(k, str(v))
-        try:
-            return bool(eval(condition, {"__builtins__": {}}, {}))
-        except Exception:
-            return False
-
-    def process_expression(expr: str, context: dict) -> str:
-        """Process a {{ expression }} block."""
-        expr = expr.strip()
-        if expr in context:
-            return str(context[expr])
-        return ""
-
-    def process_block(lines: List[str], context: dict, indent: int = 0) -> List[str]:
-        """Process a block of template lines recursively."""
-        result = []
-        i = 0
-        while i < len(lines):
-            line = lines[i]
-
-            # Handle comments
-            line = re.sub(r'{#.*?#}', '', line)
-
-            # Handle if blocks
-            if_match = re.search(r'{%\s*if\s+(.+?)\s*%}', line)
-            if if_match:
-                condition = if_match.group(1)
-                if_block = []
-                else_block = []
-                i += 1
-                nesting = 1
-                in_else_block = False
-                while i < len(lines) and nesting > 0:
-                    if re.search(r'{%\s*if\s+', lines[i]):
-                        nesting += 1
-                    elif re.search(r'{%\s*endif\s*%}', lines[i]):
-                        nesting -= 1
-                    elif nesting == 1 and re.search(r'{%\s*else\s*%}', lines[i]):
-                        in_else_block = True
-                        i += 1
-                        continue
-
-                    if nesting > 0:
-                        if in_else_block:
-                            else_block.append(lines[i])
-                        else:
-                            if_block.append(lines[i])
-                    i += 1
-                if evaluate_condition(condition, context):
-                    result.extend(process_block(if_block, context, indent))
-                else:
-                    result.extend(process_block(else_block, context, indent))
-                continue
-
-            # Handle for loops
-            for_match = re.search(r'{%\s*for\s+(\w+)\s+in\s+(\w+)\s*%}', line)
-            if for_match:
-                var_name, iterable_name = for_match.groups()
-                for_block = []
-                i += 1
-                nesting = 1
-                while i < len(lines) and nesting > 0:
-                    if re.search(r'{%\s*for\s+', lines[i]):
-                        nesting += 1
-                    elif re.search(r'{%\s*endfor\s*%}', lines[i]):
-                        nesting -= 1
-                    if nesting > 0:
-                        for_block.append(lines[i])
-                    i += 1
-                if iterable_name in context and isinstance(context[iterable_name], (list, tuple)):
-                    for item in context[iterable_name]:
-                        loop_context = dict(context)
-                        loop_context[var_name] = item
-                        result.extend(process_block(for_block, loop_context, indent))
-                continue
-
-            # Handle variable substitution
-            line = re.sub(r'{{\s*(\w+)\s*}}', lambda m: process_expression(m.group(1), context), line)
-
-            # Handle whitespace control
-            line = re.sub(r'{%-', '{%', line)
-            line = re.sub(r'-%}', '%}', line)
-
-            # Add line with proper indentation, preserving blank lines
-            if line.strip():
-                result.append(" " * (indent_spaces + indent) + line)
-            else:
-                result.append("")
-
-            i += 1
-
-        return result
-
-    # Split template into lines and process
-    lines = template.split('\n')
-    processed_lines = process_block(lines, substitutions)
-
-    return '\n'.join(processed_lines)
-
-def type_to_sql(type) -> str:
-    if type is str:
-        return "VARCHAR"
-    if type is int:
-        return "NUMBER"
-    if type is Number:
-        return "DECIMAL(38, 15)"
-    if type is float:
-        return "FLOAT"
-    if type is decimal.Decimal:
-        return "DECIMAL(38, 15)"
-    if type is bool:
-        return "BOOLEAN"
-    if type is dict:
-        return "VARIANT"
-    if type is list:
-        return "ARRAY"
-    if type is bytes:
-        return "BINARY"
-    if type is datetime:
-        return "TIMESTAMP"
-    if type is date:
-        return "DATE"
-    if isinstance(type, dsl.Type):
-        return "VARCHAR"
-    raise ValueError(f"Unknown type {type}")
-
-def type_to_snowpark(type) -> str:
-    if type is str:
-        return "StringType()"
-    if type is int:
-        return "IntegerType()"
-    if type is float:
-        return "FloatType()"
-    if type is Number:
-        return "DecimalType(38, 15)"
-    if type is decimal.Decimal:
-        return "DecimalType(38, 15)"
-    if type is bool:
-        return "BooleanType()"
-    if type is dict:
-        return "MapType()"
-    if type is list:
-        return "ArrayType()"
-    if type is bytes:
-        return "BinaryType()"
-    if type is datetime:
-        return "TimestampType()"
-    if type is date:
-        return "DateType()"
-    if isinstance(type, dsl.Type):
-        return "StringType()"
-    raise ValueError(f"Unknown type {type}")
-
-def _sanitize_user_name(user: str) -> str:
-    # Extract the part before the '@'
-    sanitized_user = user.split('@')[0]
-    # Replace any character that is not a letter, number, or underscore with '_'
-    sanitized_user = re.sub(r'[^a-zA-Z0-9_]', '_', sanitized_user)
-    return sanitized_user
-
-def _is_engine_issue(response_message: str) -> bool:
-    return any(kw in response_message.lower() for kw in ENGINE_ERRORS + ENGINE_NOT_READY_MSGS)
-
-def _is_database_issue(response_message: str) -> bool:
-    return any(kw in response_message.lower() for kw in DATABASE_ERRORS)
-
-
 #--------------------------------------------------
 # Resources
 #--------------------------------------------------
 
 APP_NAME = "___RAI_APP___"
 
+@dataclass
+class ExecContext:
+    """Execution context for SQL queries, containing all parameters needed for execution and retry."""
+    code: str
+    params: List[Any] | None = None
+    raw: bool = False
+    help: bool = True
+    skip_engine_db_error_retry: bool = False
+
+    def re_execute(self, resources: 'Resources') -> Any:
+        """Re-execute this context's query using the provided resources instance."""
+        return resources._exec(
+            code=self.code,
+            params=self.params,
+            raw=self.raw,
+            help=self.help,
+            skip_engine_db_error_retry=self.skip_engine_db_error_retry
+        )
+
 class Resources(ResourcesBase):
     def __init__(
         self,
@@ -301,7 +140,7 @@ class Resources(ResourcesBase):
         dry_run: bool = False,
         reset_session: bool = False,
         generation: Generation | None = None,
-        language: str = "rel",
+        language: str = "rel",  # Accepted for backward compatibility, but not stored in base class
     ):
         super().__init__(profile, config=config)
         self._token_handler: TokenHandler | None = None
@@ -319,16 +158,97 @@ class Resources(ResourcesBase):
         # self.sources contains fully qualified Snowflake table/view names
         self.sources: set[str] = set()
         self._sproc_models = None
-        self.database = ""
+        # Store language for backward compatibility (used by child classes for use_index polling)
         self.language = language
+        # Register error and state handlers
+        self._register_handlers()
+        # Register atexit callback to cancel pending transactions
         atexit.register(self.cancel_pending_transactions)
 
+    #--------------------------------------------------
+    # Initialization & Properties
+    #--------------------------------------------------
+
+    def _register_handlers(self) -> None:
+        """Register error and engine state handlers for processing."""
+        # Register base handlers using getter methods that subclasses can override
+        # Use defensive copying to ensure each instance has its own handler lists
+        # and prevent cross-instance contamination from subclass mutations
+        self._error_handlers = list(self._get_error_handlers())
+        self._sync_engine_state_handlers = list(self._get_engine_state_handlers(is_async=False))
+        self._async_engine_state_handlers = list(self._get_engine_state_handlers(is_async=True))
+
+    def _get_error_handlers(self) -> list[ErrorHandler]:
+        """Get list of error handlers. Subclasses can override to add custom handlers.
+
+        Returns:
+            List of error handlers for standard error processing using Strategy Pattern.
+
+        Example:
+            def _get_error_handlers(self) -> list[ErrorHandler]:
+                # Get base handlers
+                handlers = super()._get_error_handlers()
+                # Add custom handler
+                handlers.append(MyCustomErrorHandler())
+                return handlers
+        """
+        return [
+            DuoSecurityErrorHandler(),
+            AppMissingErrorHandler(),
+            DatabaseErrorsHandler(),
+            EngineErrorsHandler(),
+            ServiceNotStartedErrorHandler(),
+            TransactionAbortedErrorHandler(),
+        ]
+
+    def _get_engine_state_handlers(self, is_async: bool = False) -> list[EngineStateHandler]:
+        """Get list of engine state handlers. Subclasses can override.
+
+        Args:
+            is_async: If True, returns async handlers; if False, returns sync handlers.
+
+        Returns:
+            List of engine state handlers for processing engine states.
+
+        Example:
+            def _get_engine_state_handlers(self, is_async: bool = False) -> list[EngineStateHandler]:
+                # Get base handlers
+                handlers = super()._get_engine_state_handlers(is_async)
+                # Add custom handler
+                handlers.append(MyCustomStateHandler())
+                return handlers
+        """
+        if is_async:
+            return [
+                AsyncPendingStateHandler(),
+                AsyncSuspendedStateHandler(),
+                AsyncReadyStateHandler(),
+                AsyncGoneStateHandler(),
+                AsyncMissingEngineHandler(),
+            ]
+        else:
+            return [
+                SyncPendingStateHandler(),
+                SyncSuspendedStateHandler(),
+                SyncReadyStateHandler(),
+                SyncGoneStateHandler(),
+                SyncMissingEngineHandler(),
+            ]
+
     @property
     def token_handler(self) -> TokenHandler:
         if not self._token_handler:
             self._token_handler = TokenHandler.from_config(self.config)
         return self._token_handler
 
+    def reset(self):
+        """Reset the session."""
+        self._session = None
+
+    #--------------------------------------------------
+    # Session Management
+    #--------------------------------------------------
+
     def is_erp_running(self, app_name: str) -> bool:
         """Check if the ERP is running. The app.service_status() returns single row/column containing an array of JSON service status objects."""
         query = f"CALL {app_name}.app.service_status();"
@@ -426,7 +346,28 @@ class Resources(ResourcesBase):
         except Exception as e:
             raise e
 
+    #--------------------------------------------------
+    # Core Execution Methods
+    #--------------------------------------------------
+
     def _exec_sql(self, code: str, params: List[Any] | None, raw=False):
+        """
+        Lowest-level SQL execution method.
+
+        Directly executes SQL via the Snowflake session. This is the foundation
+        for all other execution methods. It:
+        - Replaces APP_NAME placeholder with actual app name
+        - Executes SQL with optional parameters
+        - Returns either raw session results or collected results
+
+        Args:
+            code: SQL code to execute (may contain APP_NAME placeholder)
+            params: Optional SQL parameters
+            raw: If True, return raw session results; if False, collect results
+
+        Returns:
+            Raw session results if raw=True, otherwise collected results
+        """
         assert self._session is not None
         sess_results = self._session.sql(
             code.replace(APP_NAME, self.get_app_name()),
@@ -441,86 +382,91 @@ class Resources(ResourcesBase):
         code: str,
         params: List[Any] | Any | None = None,
         raw: bool = False,
-        help: bool = True
+        help: bool = True,
+        skip_engine_db_error_retry: bool = False
     ) -> Any:
+        """
+        Mid-level SQL execution method with error handling.
+
+        This is the primary method for executing SQL queries. It wraps _exec_sql
+        with comprehensive error handling and parameter normalization. Used
+        extensively throughout the codebase for direct SQL operations like:
+        - SHOW commands (warehouses, databases, etc.)
+        - CALL statements to RAI app stored procedures
+        - Transaction management queries
+
+        The error handling flow:
+        1. Normalizes parameters and creates execution context
+        2. Calls _exec_sql to execute the query
+        3. On error, uses standard error handling (Strategy Pattern), which subclasses
+            can influence via `_get_error_handlers()` or by overriding `_handle_standard_exec_errors()`
+
+        Args:
+            code: SQL code to execute
+            params: Optional SQL parameters (normalized to list if needed)
+            raw: If True, return raw session results; if False, collect results
+            help: If True, enable error handling; if False, raise errors immediately
+            skip_engine_db_error_retry: If True, skip use_index retry logic in error handlers
+
+        Returns:
+            Query results (collected or raw depending on 'raw' parameter)
+        """
         # print(f"\n--- sql---\n{code}\n--- end sql---\n")
+        # Ensure session is initialized
         if not self._session:
             self._session = self.get_sf_session()
 
+        # Normalize parameters
+        normalized_params = normalize_params(params)
+
+        # Create execution context
+        ctx = ExecContext(
+            code=code,
+            params=normalized_params,
+            raw=raw,
+            help=help,
+            skip_engine_db_error_retry=skip_engine_db_error_retry
+        )
+
+        # Execute SQL
         try:
-            if params is not None and not isinstance(params, list):
-                params = cast(List[Any], [params])
-            return self._exec_sql(code, params, raw=raw)
+            return self._exec_sql(ctx.code, ctx.params, raw=ctx.raw)
         except Exception as e:
-            if not help:
+            if not ctx.help:
                 raise e
-            orig_message = str(e).lower()
-            rai_app = self.config.get("rai_app_name", "")
-            current_role = self.config.get("role")
-            engine = self.get_default_engine_name()
-            engine_size = self.config.get_default_engine_size()
-            assert isinstance(rai_app, str), f"rai_app_name must be a string, not {type(rai_app)}"
-            assert isinstance(engine, str), f"engine must be a string, not {type(engine)}"
-            print("\n")
-            if DUO_TEXT in orig_message:
-                raise DuoSecurityFailed(e)
-            if re.search(f"database '{rai_app}' does not exist or not authorized.".lower(), orig_message):
-                exception = SnowflakeAppMissingException(rai_app, current_role)
-                raise exception from None
-            if _is_engine_issue(orig_message) or _is_database_issue(orig_message):
-                try:
-                    self._poll_use_index(
-                        app_name=self.get_app_name(),
-                        sources=self.sources,
-                        model=self.database,
-                        engine_name=engine,
-                        engine_size=engine_size
-                    )
-                    return self._exec(code, params, raw=raw, help=help)
-                except EngineNameValidationException as e:
-                    raise EngineNameValidationException(engine) from e
-                except Exception as e:
-                    raise EngineProvisioningFailed(engine, e) from e
-            elif re.search(r"javascript execution error", orig_message):
-                match = re.search(r"\"message\":\"(.*)\"", orig_message)
-                if match:
-                    message = match.group(1)
-                    if "engine is in pending" in message or "engine is provisioning" in message:
-                        raise EnginePending(engine)
-                    else:
-                        raise RAIException(message) from None
-
-            if re.search(r"the relationalai service has not been started.", orig_message):
-                app_name = self.config.get("rai_app_name", "")
-                assert isinstance(app_name, str), f"rai_app_name must be a string, not {type(app_name)}"
-                raise SnowflakeRaiAppNotStarted(app_name)
-
-            if re.search(r"state:\s*aborted", orig_message):
-                txn_id_match = re.search(r"id:\s*([0-9a-f\-]+)", orig_message)
-                if txn_id_match:
-                    txn_id = txn_id_match.group(1)
-                    problems = self.get_transaction_problems(txn_id)
-                    if problems:
-                        for problem in problems:
-                            if isinstance(problem, dict):
-                                type_field = problem.get('TYPE')
-                                message_field = problem.get('MESSAGE')
-                                report_field = problem.get('REPORT')
-                            else:
-                                type_field = problem.TYPE
-                                message_field = problem.MESSAGE
-                                report_field = problem.REPORT
 
-                            raise RAIAbortedTransactionError(type_field, message_field, report_field)
-                raise RAIException(str(e))
-            raise RAIException(str(e))
+            # Handle standard errors
+            result = self._handle_standard_exec_errors(e, ctx)
+            if result is not None:
+                return result
+
+    #--------------------------------------------------
+    # Error Handling
+    #--------------------------------------------------
 
+    def _handle_standard_exec_errors(self, e: Exception, ctx: ExecContext) -> Any | None:
+        """
+        Handle standard Snowflake/RAI errors using Strategy Pattern.
 
-    def reset(self):
-        self._session = None
+        Each error type has a dedicated handler class that encapsulates
+        the detection logic and exception creation. Handlers are processed
+        in order until one matches and handles the error.
+        """
+        message = str(e).lower()
+
+        # Try each handler in order until one matches
+        for handler in self._error_handlers:
+            if handler.matches(e, message, ctx, self):
+                result = handler.handle(e, ctx, self)
+                if result is not None:
+                    return result
+                return  # Handler raised exception, we're done
+
+        # Fallback: transform to RAIException
+        raise RAIException(str(e))
 
     #--------------------------------------------------
-    # Check direct access is enabled
+    # Feature Detection & Configuration
     #--------------------------------------------------
 
     def is_direct_access_enabled(self) -> bool:
@@ -542,9 +488,6 @@ class Resources(ResourcesBase):
         except Exception as e:
             raise Exception(f"Unable to determine if direct access is enabled. Details error: {e}") from e
 
-    #--------------------------------------------------
-    # Snowflake Account Flags
-    #--------------------------------------------------
 
     def is_account_flag_set(self, flag: str) -> bool:
         results = self._exec(
@@ -564,7 +507,8 @@ class Resources(ResourcesBase):
                 f"call {APP_NAME}.api.get_database('{database}');"
             )
         except Exception as e:
-            if "Database does not exist" in str(e):
+            messages = collect_error_messages(e)
+            if any("database does not exist" in msg for msg in messages):
                 return None
             raise e
 
@@ -588,10 +532,11 @@ class Resources(ResourcesBase):
         try:
             results = self._exec(query)
         except Exception as e:
-            if "Database does not exist" in str(e):
+            messages = collect_error_messages(e)
+            if any("database does not exist" in msg for msg in messages):
                 return None
             # fallback to None for old sql-lib versions
-            if "Unknown user-defined function" in str(e):
+            if any("unknown user-defined function" in msg for msg in messages):
                 return None
             raise e
 
@@ -608,6 +553,139 @@ class Resources(ResourcesBase):
     # Engines
     #--------------------------------------------------
 
+    def _prepare_engine_params(
+        self,
+        name: str | None,
+        size: str | None,
+        use_default_size: bool = False
+    ) -> tuple[str, str | None]:
+        """
+        Prepare engine parameters by resolving and validating name and size.
+
+        Args:
+            name: Engine name (None to use default)
+            size: Engine size (None to use config or default)
+            use_default_size: If True and size is None, use get_default_engine_size()
+
+        Returns:
+            Tuple of (engine_name, engine_size)
+
+        Raises:
+            EngineNameValidationException: If engine name is invalid
+            Exception: If engine size is invalid
+        """
+        from relationalai.tools.cli_helpers import validate_engine_name
+
+        engine_name = name or self.get_default_engine_name()
+
+        # Resolve engine size
+        if size:
+            engine_size = size
+        else:
+            if use_default_size:
+                engine_size = self.config.get_default_engine_size()
+            else:
+                engine_size = self.config.get("engine_size", None)
+
+        # Validate engine size
+        if engine_size:
+            is_size_valid, sizes = self.validate_engine_size(engine_size)
+            if not is_size_valid:
+                error_msg = f"Invalid engine size '{engine_size}'. Valid sizes are: {', '.join(sizes)}"
+                if use_default_size:
+                    error_msg = f"Invalid engine size in config: '{engine_size}'. Valid sizes are: {', '.join(sizes)}"
+                raise Exception(error_msg)
+
+        # Validate engine name
+        is_name_valid, _ = validate_engine_name(engine_name)
+        if not is_name_valid:
+            raise EngineNameValidationException(engine_name)
+
+        return engine_name, engine_size
+
+    def _get_state_handler(self, state: str | None, handlers: list[EngineStateHandler]) -> EngineStateHandler:
+        """Find the appropriate state handler for the given state."""
+        for handler in handlers:
+            if handler.handles_state(state):
+                return handler
+        # Fallback to missing engine handler if no match
+        return handlers[-1]  # Last handler should be MissingEngineHandler
+
+    def _process_engine_state(
+        self,
+        engine: EngineState | Dict[str, Any] | None,
+        context: EngineContext,
+        handlers: list[EngineStateHandler],
+        set_active_on_success: bool = False
+    ) -> EngineState | Dict[str, Any] | None:
+        """
+        Process engine state using appropriate state handler.
+
+        Args:
+            engine: Current engine state (or None if missing)
+            context: Engine context for state handling
+            handlers: List of state handlers to use (sync or async)
+            set_active_on_success: If True, set engine as active when handler returns engine
+
+        Returns:
+            Engine state after processing, or None if engine needs to be created
+        """
+        # Find and execute appropriate state handler
+        state = engine["state"] if engine else None
+        handler = self._get_state_handler(state, handlers)
+        engine = handler.handle(engine, context, self)
+
+        # If handler returned None and we didn't start with None state, engine needs to be created
+        # (e.g., GONE state deleted the engine, so we need to create a new one)
+        if not engine and state is not None:
+            handler = self._get_state_handler(None, handlers)
+            handler.handle(None, context, self)
+        elif set_active_on_success:
+            # Cast to EngineState for type safety (handlers return EngineDict which is compatible)
+            self._set_active_engine(cast(EngineState, engine))
+
+        return engine
+
+    def _handle_engine_creation_errors(self, error: Exception, engine_name: str, preserve_rai_exception: bool = False) -> None:
+        """
+        Handle errors during engine creation using error handlers.
+
+        Args:
+            error: The exception that occurred
+            engine_name: Name of the engine being created
+            preserve_rai_exception: If True, re-raise RAIException without wrapping
+
+        Raises:
+            RAIException: If preserve_rai_exception is True and error is RAIException
+            EngineProvisioningFailed: If error is not handled by error handlers
+        """
+        # Preserve RAIException passthrough if requested (for async mode)
+        if preserve_rai_exception and isinstance(error, RAIException):
+            raise error
+
+        # Check if this is a known error type that should be handled by error handlers
+        message = str(error).lower()
+        handled = False
+        # Engine creation isn't tied to a specific SQL ExecContext; pass a context that
+        # disables use_index retry behavior (and any future ctx-dependent handlers).
+        ctx = ExecContext(code="", help=True, skip_engine_db_error_retry=True)
+        for handler in self._error_handlers:
+            if handler.matches(error, message, ctx, self):
+                handler.handle(error, ctx, self)
+                handled = True
+                break  # Handler raised exception, we're done
+
+        # If not handled by error handlers, wrap in EngineProvisioningFailed
+        if not handled:
+            raise EngineProvisioningFailed(engine_name, error) from error
+
+    def validate_engine_size(self, size: str) -> Tuple[bool, List[str]]:
+        if size is not None:
+            sizes = self.get_engine_sizes()
+            if size not in sizes:
+                return False, sizes
+        return True, []
+
     def get_engine_sizes(self, cloud_provider: str|None=None):
         sizes = []
         if cloud_provider is None:
@@ -810,19 +888,17 @@ Otherwise, remove it from your '{profile}' configuration profile.
         engine_size: str | None = None,
         program_span_id: str | None = None,
         headers: Dict | None = None,
-    ):
-        return UseIndexPoller(
-            self,
-            app_name,
-            sources,
-            model,
-            engine_name,
-            engine_size,
-            self.language,
-            program_span_id,
-            headers,
-            self.generation
-        ).poll()
+    ) -> None:
+        """
+        Poll use_index to prepare indices for the given sources.
+
+        This is an optional interface method. Base Resources provides a no-op implementation.
+        UseIndexResources and DirectAccessResources override this to provide actual polling.
+
+        Returns:
+            None for base implementation. Child classes may return poller results.
+        """
+        return None
 
     def maybe_poll_use_index(
         self,
@@ -833,36 +909,17 @@ Otherwise, remove it from your '{profile}' configuration profile.
         engine_size: str | None = None,
         program_span_id: str | None = None,
         headers: Dict | None = None,
-    ):
-        """Only call poll() if there are sources to process and cache is not valid."""
-        sources_list = list(sources)
-        self.database = model
-        if sources_list:
-            poller = UseIndexPoller(
-                self,
-                app_name,
-                sources_list,
-                model,
-                engine_name,
-                engine_size,
-                self.language,
-                program_span_id,
-                headers,
-                self.generation
-            )
-            # If cache is valid (data freshness has not expired), skip polling
-            if poller.cache.is_valid():
-                cached_sources = len(poller.cache.sources)
-                total_sources = len(sources_list)
-                cached_timestamp = poller.cache._metadata.get("cachedIndices", {}).get(poller.cache.key, {}).get("last_use_index_update_on", "")
-
-                message = f"Using cached data for {cached_sources}/{total_sources} data streams"
-                if cached_timestamp:
-                    print(f"\n{message} (cached at {cached_timestamp})\n")
-                else:
-                    print(f"\n{message}\n")
-            else:
-                return poller.poll()
+    ) -> None:
+        """
+        Only call _poll_use_index if there are sources to process.
+
+        This is an optional interface method. Base Resources provides a no-op implementation.
+        UseIndexResources and DirectAccessResources override this to provide actual polling with caching.
+
+        Returns:
+            None for base implementation. Child classes may return poller results.
+        """
+        return None
 
     #--------------------------------------------------
     # Models
@@ -900,11 +957,6 @@ Otherwise, remove it from your '{profile}' configuration profile.
     def list_exports(self, database: str, engine: str):
         return []
 
-    def format_sproc_name(self, name: str, type:Any) -> str:
-        if type is datetime:
-            return f"{name}.astimezone(ZoneInfo('UTC')).isoformat(timespec='milliseconds')"
-        else:
-            return name
 
     def get_export_code(self, params: ExportParams, all_installs):
         sql_inputs = ", ".join([f"{name} {type_to_sql(type)}" for (name, _, type) in params.inputs])
@@ -926,15 +978,14 @@ Otherwise, remove it from your '{profile}' configuration profile.
                 clean_inputs.append(f"{name} = '\"' + escape({name}) + '\"'")
             # Replace `var` with `name` and keep the following non-word character unchanged
             pattern = re.compile(re.escape(var) + r'(\W)')
-            value = self.format_sproc_name(name, type)
+            value = format_sproc_name(name, type)
             safe_rel = re.sub(pattern, rf"{{{value}}}\1", safe_rel)
         if py_inputs:
             py_inputs = f", {py_inputs}"
         clean_inputs = ("\n").join(clean_inputs)
-        assert __package__ is not None, "Package name must be set"
         file = "export_procedure.py.jinja"
         with importlib.resources.open_text(
-            __package__, file
+            "relationalai.clients.resources.snowflake", file
         ) as f:
             template = f.read()
         def quote(s: str, f = False) -> str:
@@ -1092,15 +1143,6 @@ Otherwise, remove it from your '{profile}' configuration profile.
     # Imports
     #--------------------------------------------------
 
-    def is_valid_import_state(self, state:str):
-        return state in VALID_IMPORT_STATES
-
-    def imports_to_dicts(self, results):
-        parsed_results = [
-            {field.lower(): row[field] for field in IMPORT_STREAM_FIELDS}
-            for row in results
-        ]
-        return parsed_results
 
     def change_stream_status(self, stream_id: str, model:str, suspend: bool):
         if stream_id and model:
@@ -1272,7 +1314,7 @@ Otherwise, remove it from your '{profile}' configuration profile.
         results = self._exec(f"CALL {APP_NAME}.api.get_data_stream('{name}', '{model}');")
         if not results:
             return None
-        return self.imports_to_dicts(results)
+        return imports_to_dicts(results)
 
     def create_import_stream(self, source:ImportSource, model:str, rate = 1, options: dict|None = None):
         assert isinstance(source, ImportSourceTable), "Snowflake integration only supports loading from SF Tables. Try loading your data as a table via the Snowflake interface first."
@@ -1307,7 +1349,8 @@ Otherwise, remove it from your '{profile}' configuration profile.
             try:
                 self._exec(command)
             except Exception as e:
-                if "ensure that CHANGE_TRACKING is enabled on the source object" in str(e):
+                messages = collect_error_messages(e)
+                if any("ensure that change_tracking is enabled on the source object" in msg for msg in messages):
                     if self.config.get("ensure_change_tracking", False) and not tracking_just_changed:
                         try:
                             self._exec(f"ALTER {kind} {object} SET CHANGE_TRACKING = TRUE;")
@@ -1318,7 +1361,7 @@ Otherwise, remove it from your '{profile}' configuration profile.
                         print("\n")
                         exception = SnowflakeChangeTrackingNotEnabledException((object, kind))
                         raise exception from None
-                elif "Database does not exist" in str(e):
+                elif any("database does not exist" in msg for msg in messages):
                     print("\n")
                     raise ModelNotFoundException(model) from None
                 raise e
@@ -1381,42 +1424,6 @@ Otherwise, remove it from your '{profile}' configuration profile.
         # @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"
 
-    def decrypt_stream(self, key: bytes, iv: bytes, src: bytes) -> bytes:
-        """Decrypt the provided stream with PKCS#5 padding handling."""
-
-        if crypto_disabled:
-            if isinstance(runtime_env, SnowbookEnvironment) and runtime_env.runner == "warehouse":
-                raise Exception("Please open the navigation-bar dropdown labeled *Packages* and select `cryptography` under the *Anaconda Packages* section, and then re-run your query.")
-            else:
-                raise Exception("library `cryptography.hazmat` missing; please install")
-
-        # `type:ignore`s are because of the conditional import, which
-        # we have because warehouse-based snowflake notebooks don't support
-        # the crypto library we're using.
-        cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend()) # type: ignore
-        decryptor = cipher.decryptor()
-
-        # Decrypt the data
-        decrypted_padded_data = decryptor.update(src) + decryptor.finalize()
-
-        # Unpad the decrypted data using PKCS#5
-        unpadder = padding.PKCS7(128).unpadder()  # type: ignore # Use 128 directly for AES
-        unpadded_data = unpadder.update(decrypted_padded_data) + unpadder.finalize()
-
-        return unpadded_data
-
-    def _decrypt_artifact(self, data: bytes, encryption_material: str) -> bytes:
-        """Decrypts the artifact data using provided encryption material."""
-        encryption_material_parts = encryption_material.split("|")
-        assert len(encryption_material_parts) == 3, "Invalid encryption material"
-
-        algorithm, key_base64, iv_base64 = encryption_material_parts
-        assert algorithm == "AES_128_CBC", f"Unsupported encryption algorithm {algorithm}"
-
-        key = base64.standard_b64decode(key_base64)
-        iv = base64.standard_b64decode(iv_base64)
-
-        return self.decrypt_stream(key, iv, data)
 
     def _list_exec_async_artifacts(self, txn_id: str, headers: Dict | None = None) -> Dict[str, Dict]:
         """Grab the list of artifacts produced in the transaction and the URLs to retrieve their contents."""
@@ -1468,7 +1475,7 @@ Otherwise, remove it from your '{profile}' configuration profile.
             if len(content) == 0:
                 return b""
 
-            return self._decrypt_artifact(content, encryption_material)
+            return decrypt_artifact(content, encryption_material)
 
         # otherwise, return content directly
         return content
@@ -1548,7 +1555,7 @@ Otherwise, remove it from your '{profile}' configuration profile.
     def get_url_key(self, metadata) -> str:
         # In Azure, there is only one type of URL, which is used for both internal and
         # external access; always use that one
-        if self.is_azure(metadata['PRESIGNED_URL']):
+        if is_azure_url(metadata['PRESIGNED_URL']):
             return 'PRESIGNED_URL'
 
         configured = self.config.get("download_url_type", None)
@@ -1557,17 +1564,11 @@ Otherwise, remove it from your '{profile}' configuration profile.
         elif configured == "external":
             return "PRESIGNED_URL"
 
-        if self.is_container_runtime():
+        if is_container_runtime():
             return 'PRESIGNED_URL_AP'
 
         return 'PRESIGNED_URL'
 
-    def is_azure(self, url) -> bool:
-        return "blob.core.windows.net" in url
-
-    def is_container_runtime(self) -> bool:
-        return isinstance(runtime_env, SnowbookEnvironment) and runtime_env.runner == "container"
-
     def _exec_rai_app(
         self,
         database: str,
@@ -1581,6 +1582,40 @@ Otherwise, remove it from your '{profile}' configuration profile.
         language: str = "rel",
         query_timeout_mins: int | None = None,
     ):
+        """
+        High-level method to execute RAI app stored procedures.
+
+        Builds and executes SQL to call the RAI app's exec_async_v2 stored procedure.
+        This method handles the SQL string construction for two different formats:
+        1. New format (with graph index): Uses object payload with parameterized query
+        2. Legacy format: Uses positional parameters
+
+        The choice between formats depends on the use_graph_index configuration.
+        The new format allows the stored procedure to hash the model and username
+        to determine the database, while the legacy format uses the passed database directly.
+
+        This method is called by _exec_async_v2 to create transactions. It skips
+        use_index retry logic (skip_engine_db_error_retry=True) because that
+        is handled at a higher level by exec_raw/exec_lqp.
+
+        Args:
+            database: Database/model name
+            engine: Engine name (optional)
+            raw_code: Code to execute (REL, LQP, or SQL)
+            inputs: Input parameters for the query
+            readonly: Whether the transaction is read-only
+            nowait_durable: Whether to wait for durable writes
+            request_headers: Optional HTTP headers
+            bypass_index: Whether to bypass graph index setup
+            language: Query language ("rel" or "lqp")
+            query_timeout_mins: Optional query timeout in minutes
+
+        Returns:
+            Response from the stored procedure call (transaction creation result)
+
+        Raises:
+            Exception: If transaction creation fails
+        """
         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)
@@ -1609,9 +1644,11 @@ Otherwise, remove it from your '{profile}' configuration profile.
                 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});"
+        # Don't let exec setup GI on failure, exec_raw and exec_lqp will do that and add the correct headers.
         response = self._exec(
             sql_string,
             raw_code,
+            skip_engine_db_error_retry=True,
         )
         if not response:
             raise Exception("Failed to create transaction")
@@ -1629,7 +1666,38 @@ Otherwise, remove it from your '{profile}' configuration profile.
         bypass_index=False,
         language: str = "rel",
         query_timeout_mins: int | None = None,
+        gi_setup_skipped: bool = False,
     ):
+        """
+        High-level async execution method with transaction polling and artifact management.
+
+        This is the core method for executing queries asynchronously. It:
+        1. Creates a transaction by calling _exec_rai_app
+        2. Handles two execution paths:
+            - Fast path: Transaction completes immediately (COMPLETED/ABORTED)
+            - Slow path: Transaction is pending, requires polling until completion
+        3. Manages pending transactions list
+        4. Downloads and returns query results/artifacts
+
+        This method is called by _execute_code (base implementation) and can be
+        overridden by child classes (e.g., DirectAccessResources uses HTTP instead).
+
+        Args:
+            database: Database/model name
+            engine: Engine name (optional)
+            raw_code: Code to execute (REL, LQP, or SQL)
+            inputs: Input parameters for the query
+            readonly: Whether the transaction is read-only
+            nowait_durable: Whether to wait for durable writes
+            headers: Optional HTTP headers
+            bypass_index: Whether to bypass graph index setup
+            language: Query language ("rel" or "lqp")
+            query_timeout_mins: Optional query timeout in minutes
+            gi_setup_skipped: Whether graph index setup was skipped (for retry logic)
+
+        Returns:
+            Query results (downloaded artifacts)
+        """
         if inputs is None:
             inputs = {}
         request_headers = debugging.add_current_propagation_headers(headers)
@@ -1638,6 +1706,8 @@ Otherwise, remove it from your '{profile}' configuration profile.
         with debugging.span("transaction", **query_attrs_dict) as txn_span:
             with debugging.span("create_v2", **query_attrs_dict) as create_span:
                 request_headers['user-agent'] = get_pyrel_version(self.generation)
+                request_headers['gi_setup_skipped'] = str(gi_setup_skipped)
+                request_headers['pyrel_program_id'] = debugging.get_program_span_id() or ""
                 response = self._exec_rai_app(
                     database=database,
                     engine=engine,
@@ -1699,204 +1769,127 @@ Otherwise, remove it from your '{profile}' configuration profile.
         return engine and engine["state"] == "READY"
 
     def auto_create_engine(self, name: str | None = None, size: str | None = None, headers: Dict | None = None):
-        from relationalai.tools.cli_helpers import validate_engine_name
+        """Synchronously create/ensure an engine is ready, blocking until ready."""
         with debugging.span("auto_create_engine", active=self._active_engine) as span:
             active = self._get_active_engine()
             if active:
                 return active
 
-            engine_name = name or self.get_default_engine_name()
-
-            # Use the provided size or fall back to the config
-            if size:
-                engine_size = size
-            else:
-                engine_size = self.config.get("engine_size", None)
-
-            # Validate engine size
-            if engine_size:
-                is_size_valid, sizes = self.validate_engine_size(engine_size)
-                if not is_size_valid:
-                    raise Exception(f"Invalid engine size '{engine_size}'. Valid sizes are: {', '.join(sizes)}")
-
-            # Validate engine name
-            is_name_valid, _ = validate_engine_name(engine_name)
-            if not is_name_valid:
-                raise EngineNameValidationException(engine_name)
+            # Resolve and validate parameters
+            engine_name, engine_size = self._prepare_engine_params(name, size)
 
             try:
+                # Get current engine state
                 engine = self.get_engine(engine_name)
                 if engine:
                     span.update(cast(dict, engine))
 
-                # if engine is in the pending state, poll until its status changes
-                # if engine is gone, delete it and create new one
-                # if engine is in the ready state, return engine name
-                if engine:
-                    if engine["state"] == "PENDING":
-                        # if the user explicitly specified a size, warn if the pending engine size doesn't match it
-                        if size is not None and engine["size"] != size:
-                            EngineSizeMismatchWarning(engine_name, engine["size"], size)
-                        # poll until engine is ready
-                        with Spinner(
-                            "Waiting for engine to be initialized",
-                            "Engine ready",
-                        ):
-                            poll_with_specified_overhead(lambda: self.is_engine_ready(engine_name), overhead_rate=0.1, max_delay=0.5, timeout=900)
-
-                    elif engine["state"] == "SUSPENDED":
-                        with Spinner(f"Resuming engine '{engine_name}'", f"Engine '{engine_name}' resumed", f"Failed to resume engine '{engine_name}'"):
-                            try:
-                                self.resume_engine_async(engine_name, headers=headers)
-                                poll_with_specified_overhead(lambda: self.is_engine_ready(engine_name), overhead_rate=0.1, max_delay=0.5, timeout=900)
-                            except Exception:
-                                raise EngineResumeFailed(engine_name)
-                    elif engine["state"] == "READY":
-                        # if the user explicitly specified a size, warn if the ready engine size doesn't match it
-                        if size is not None and engine["size"] != size:
-                            EngineSizeMismatchWarning(engine_name, engine["size"], size)
-                        self._set_active_engine(engine)
-                        return engine_name
-                    elif engine["state"] == "GONE":
-                        try:
-                            # "Gone" is abnormal condition when metadata and SF service don't match
-                            # Therefore, we have to delete the engine and create a new one
-                            # it could be case that engine is already deleted, so we have to catch the exception
-                            self.delete_engine(engine_name, headers=headers)
-                            # After deleting the engine, set it to None so that we can create a new engine
-                            engine = None
-                        except Exception as e:
-                            # if engine is already deleted, we will get an exception
-                            # we can ignore this exception and create a new engine
-                            if isinstance(e, EngineNotFoundException):
-                                engine = None
-                                pass
-                            else:
-                                raise EngineProvisioningFailed(engine_name, e) from e
-
-                if not engine:
-                    with Spinner(
-                        f"Auto-creating engine {engine_name}",
-                        f"Auto-created engine {engine_name}",
-                        "Engine creation failed",
-                    ):
-                        self.create_engine(engine_name, size=engine_size, headers=headers)
+                # Create context for state handling
+                context = EngineContext(
+                    engine_name=engine_name,
+                    engine_size=engine_size,
+                    headers=headers,
+                    requested_size=size,
+                    span=span,
+                )
+
+                # Process engine state using sync handlers
+                self._process_engine_state(engine, context, self._sync_engine_state_handlers)
+
             except Exception as e:
-                print(e)
-                if DUO_TEXT in str(e).lower():
-                    raise DuoSecurityFailed(e)
-                raise EngineProvisioningFailed(engine_name, e) from e
+                self._handle_engine_creation_errors(e, engine_name)
+
             return engine_name
 
     def auto_create_engine_async(self, name: str | None = None):
+        """Asynchronously create/ensure an engine, returns immediately."""
         active = self._get_active_engine()
         if active and (active == name or name is None):
-            return # @NOTE: This method weirdly doesn't return engine name even though all the other ones do?
+            return active
 
         with Spinner(
             "Checking engine status",
             leading_newline=True,
         ) as spinner:
-            from relationalai.tools.cli_helpers import validate_engine_name
             with debugging.span("auto_create_engine_async", active=self._active_engine):
-                engine_name = name or self.get_default_engine_name()
-                engine_size = self.config.get("engine_size", None)
-                if engine_size:
-                    is_size_valid, sizes = self.validate_engine_size(engine_size)
-                    if not is_size_valid:
-                        raise Exception(f"Invalid engine size in config: '{engine_size}'. Valid sizes are: {', '.join(sizes)}")
-                else:
-                    engine_size = self.config.get_default_engine_size()
+                # Resolve and validate parameters (use_default_size=True for async)
+                engine_name, engine_size = self._prepare_engine_params(name, None, use_default_size=True)
 
-                is_name_valid, _ = validate_engine_name(engine_name)
-                if not is_name_valid:
-                    raise EngineNameValidationException(engine_name)
                 try:
+                    # Get current engine state
                     engine = self.get_engine(engine_name)
-                    # if engine is gone, delete it and create new one
-                    # in case of pending state, do nothing, it is use_index responsibility to wait for engine to be ready
-                    if engine:
-                        if engine["state"] == "PENDING":
-                            spinner.update_messages(
-                                {
-                                    "finished_message": f"Starting engine {engine_name}",
-                                }
-                            )
-                            pass
-                        elif engine["state"] == "SUSPENDED":
-                            spinner.update_messages(
-                                {
-                                    "finished_message": f"Resuming engine {engine_name}",
-                                }
-                            )
-                            try:
-                                self.resume_engine_async(engine_name)
-                            except Exception:
-                                raise EngineResumeFailed(engine_name)
-                        elif engine["state"] == "READY":
-                            spinner.update_messages(
-                                {
-                                    "finished_message": f"Engine {engine_name} initialized",
-                                }
-                            )
-                            pass
-                        elif engine["state"] == "GONE":
-                            spinner.update_messages(
-                                {
-                                    "message": f"Restarting engine {engine_name}",
-                                }
-                            )
-                            try:
-                                # "Gone" is abnormal condition when metadata and SF service don't match
-                                # Therefore, we have to delete the engine and create a new one
-                                # it could be case that engine is already deleted, so we have to catch the exception
-                                # set it to None so that we can create a new engine
-                                engine = None
-                                self.delete_engine(engine_name)
-                            except Exception as e:
-                                # if engine is already deleted, we will get an exception
-                                # we can ignore this exception and create a new engine asynchronously
-                                if isinstance(e, EngineNotFoundException):
-                                    engine = None
-                                    pass
-                                else:
-                                    print(e)
-                                    raise EngineProvisioningFailed(engine_name, e) from e
-
-                    if not engine:
-                        self.create_engine_async(engine_name, size=self.config.get("engine_size", None))
-                        spinner.update_messages(
-                            {
-                                "finished_message": f"Starting engine {engine_name}...",
-                            }
-                        )
-                    else:
-                        self._set_active_engine(engine)
 
-                except Exception as e:
-                    spinner.update_messages(
-                        {
-                            "finished_message": f"Failed to create engine {engine_name}",
-                        }
+                    # Create context for state handling
+                    context = EngineContext(
+                        engine_name=engine_name,
+                        engine_size=engine_size,
+                        headers=None,
+                        requested_size=None,
+                        spinner=spinner,
                     )
-                    if DUO_TEXT in str(e).lower():
-                        raise DuoSecurityFailed(e)
-                    if isinstance(e, RAIException):
-                        raise e
-                    print(e)
-                    raise EngineProvisioningFailed(engine_name, e) from e
 
-    def validate_engine_size(self, size: str) -> Tuple[bool, List[str]]:
-        if size is not None:
-            sizes = self.get_engine_sizes()
-            if size not in sizes:
-                return False, sizes
-        return True, []
+                    # Process engine state using async handlers
+                    self._process_engine_state(engine, context, self._async_engine_state_handlers, set_active_on_success=True)
+
+                except Exception as e:
+                    spinner.update_messages({
+                        "finished_message": f"Failed to create engine {engine_name}",
+                    })
+                    self._handle_engine_creation_errors(e, engine_name, preserve_rai_exception=True)
+
+                return engine_name
 
     #--------------------------------------------------
     # Exec
     #--------------------------------------------------
 
+    def _execute_code(
+        self,
+        database: str,
+        engine: str | None,
+        raw_code: str,
+        inputs: Dict | None,
+        readonly: bool,
+        nowait_durable: bool,
+        headers: Dict | None,
+        bypass_index: bool,
+        language: str,
+        query_timeout_mins: int | None,
+    ) -> Any:
+        """
+        Template method for code execution - can be overridden by child classes.
+
+        This is a template method that provides a hook for child classes to add
+        execution logic (like retry mechanisms). The base implementation simply
+        calls _exec_async_v2 directly.
+
+        UseIndexResources overrides this method to use _exec_with_gi_retry, which
+        adds automatic use_index polling on engine/database errors.
+
+        This method is called by exec_lqp() and exec_raw() to provide a single
+        execution point that can be customized per resource class.
+
+        Args:
+            database: Database/model name
+            engine: Engine name (optional)
+            raw_code: Code to execute (already processed/encoded)
+            inputs: Input parameters for the query
+            readonly: Whether the transaction is read-only
+            nowait_durable: Whether to wait for durable writes
+            headers: Optional HTTP headers
+            bypass_index: Whether to bypass graph index setup
+            language: Query language ("rel" or "lqp")
+            query_timeout_mins: Optional query timeout in minutes
+
+        Returns:
+            Query results
+        """
+        return self._exec_async_v2(
+            database, engine, raw_code, inputs, readonly, nowait_durable,
+            headers=headers, bypass_index=bypass_index, language=language,
+            query_timeout_mins=query_timeout_mins, gi_setup_skipped=True,
+        )
+
     def exec_lqp(
         self,
         database: str,
@@ -1910,36 +1903,12 @@ Otherwise, remove it from your '{profile}' configuration profile.
         bypass_index=False,
         query_timeout_mins: int | None = None,
     ):
+        """Execute LQP code."""
         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',
-                query_timeout_mins=query_timeout_mins,
-            )
-        except Exception as e:
-            err_message = str(e).lower()
-            if _is_engine_issue(err_message) or _is_database_issue(err_message):
-                engine_name = engine or self.get_default_engine_name()
-                engine_size = self.config.get_default_engine_size()
-                self._poll_use_index(
-                    app_name=self.get_app_name(),
-                    sources=self.sources,
-                    model=database,
-                    engine_name=engine_name,
-                    engine_size=engine_size,
-                    headers=headers,
-                )
-
-                return self._exec_async_v2(
-                    database, engine, raw_code_b64, inputs, readonly, nowait_durable,
-                    headers=headers, bypass_index=bypass_index, language='lqp',
-                    query_timeout_mins=query_timeout_mins,
-                )
-            else:
-                raise e
-
+        return self._execute_code(
+            database, engine, raw_code_b64, inputs, readonly, nowait_durable,
+            headers, bypass_index, 'lqp', query_timeout_mins
+        )
 
     def exec_raw(
         self,
@@ -1954,46 +1923,12 @@ Otherwise, remove it from your '{profile}' configuration profile.
         bypass_index=False,
         query_timeout_mins: int | None = None,
     ):
+        """Execute raw code."""
         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,
-                query_timeout_mins=query_timeout_mins,
-            )
-        except Exception as e:
-            err_message = str(e).lower()
-            if _is_engine_issue(err_message) or _is_database_issue(err_message):
-                engine_name = engine or self.get_default_engine_name()
-                engine_size = self.config.get_default_engine_size()
-                self._poll_use_index(
-                    app_name=self.get_app_name(),
-                    sources=self.sources,
-                    model=database,
-                    engine_name=engine_name,
-                    engine_size=engine_size,
-                    headers=headers,
-                )
-                return self._exec_async_v2(
-                    database,
-                    engine,
-                    raw_code,
-                    inputs,
-                    readonly,
-                    nowait_durable,
-                    headers=headers,
-                    bypass_index=bypass_index,
-                    query_timeout_mins=query_timeout_mins,
-                )
-            else:
-                raise e
+        return self._execute_code(
+            database, engine, raw_code, inputs, readonly, nowait_durable,
+            headers, bypass_index, 'rel', query_timeout_mins
+        )
 
 
     def format_results(self, results, task:m.Task|None=None) -> Tuple[DataFrame, List[Any]]:
@@ -2045,7 +1980,7 @@ Otherwise, remove it from your '{profile}' configuration profile.
                         if query_timeout_mins is not None:
                             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(?, ?, ?, ?, ?, NULL, ?, {headers}, ?);", [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)
@@ -2082,8 +2017,8 @@ Otherwise, remove it from your '{profile}' configuration profile.
                         if rejected_rows:
                             debugging.warn(RowsDroppedFromTargetTableWarning(rejected_rows, rejected_rows_count, col_names_map))
             except Exception as e:
-                msg = str(e).lower()
-                if "no columns returned" in msg or "columns of results could not be determined" in msg:
+                messages = collect_error_messages(e)
+                if any("no columns returned" in msg or "columns of results could not be determined" in msg for msg in messages):
                     pass
                 else:
                     raise e
@@ -2114,6 +2049,10 @@ Otherwise, remove it from your '{profile}' configuration profile.
         self.sources.add(parser.identity)
         return ns
 
+    #--------------------------------------------------
+    # Source Management
+    #--------------------------------------------------
+
     def _check_source_updates(self, sources: Iterable[str]):
         if not sources:
             return {}
@@ -2376,19 +2315,6 @@ Otherwise, remove it from your '{profile}' configuration profile.
     #--------------------------------------------------
     # Transactions
     #--------------------------------------------------
-    def txn_list_to_dicts(self, transactions):
-        dicts = []
-        for txn in transactions:
-            dict = {}
-            txn_dict = txn.asDict()
-            for key in txn_dict:
-                mapValue = FIELD_MAP.get(key.lower())
-                if mapValue:
-                    dict[mapValue] = txn_dict[key]
-                else:
-                    dict[key.lower()] = txn_dict[key]
-            dicts.append(dict)
-        return dicts
 
     def get_transaction(self, transaction_id):
         results = self._exec(
@@ -2396,7 +2322,7 @@ Otherwise, remove it from your '{profile}' configuration profile.
         if not results:
             return None
 
-        results = self.txn_list_to_dicts(results)
+        results = txn_list_to_dicts(results)
 
         txn = {field: results[0][field] for field in GET_TXN_SQL_FIELDS}
 
@@ -2448,7 +2374,7 @@ Otherwise, remove it from your '{profile}' configuration profile.
         results = self._exec(query, [limit])
         if not results:
             return []
-        return self.txn_list_to_dicts(results)
+        return txn_list_to_dicts(results)
 
     def cancel_transaction(self, transaction_id):
         self._exec(f"CALL {APP_NAME}.api.cancel_own_transaction(?);", [transaction_id])
@@ -2482,66 +2408,15 @@ Otherwise, remove it from your '{profile}' configuration profile.
             return None
         return results[0][0]
 
-    def list_warehouses(self):
-        results = self._exec("SHOW WAREHOUSES")
-        if not results:
-            return []
-        return [{"name":name}
-                for (name, *rest) in results]
-
-    def list_compute_pools(self):
-        results = self._exec("SHOW COMPUTE POOLS")
-        if not results:
-            return []
-        return [{"name":name, "status":status, "min_nodes":min_nodes, "max_nodes":max_nodes, "instance_family":instance_family}
-                for (name, status, min_nodes, max_nodes, instance_family, *rest) in results]
-
-    def list_roles(self):
-        results = self._exec("SELECT CURRENT_AVAILABLE_ROLES()")
-        if not results:
-            return []
-        # the response is a single row with a single column containing
-        # a stringified JSON array of role names:
-        row = results[0]
-        if not row:
-            return []
-        return [{"name": name} for name in json.loads(row[0])]
+    # CLI methods (list_warehouses, list_compute_pools, list_roles, list_apps, 
+    # list_databases, list_sf_schemas, list_tables) are now in CLIResources class
+    # schema_info is kept in base Resources class since it's used by SnowflakeSchema._fetch_info()
 
-    def list_apps(self):
-        all_apps = self._exec(f"SHOW APPLICATIONS LIKE '{RAI_APP_NAME}'")
-        if not all_apps:
-            all_apps = self._exec("SHOW APPLICATIONS")
-            if not all_apps:
-                return []
-        return [{"name":name}
-                for (time, name, *rest) in all_apps]
-
-    def list_databases(self):
-        results = self._exec("SHOW DATABASES")
-        if not results:
-            return []
-        return [{"name":name}
-                for (time, name, *rest) in results]
-
-    def list_sf_schemas(self, database:str):
-        results = self._exec(f"SHOW SCHEMAS IN {database}")
-        if not results:
-            return []
-        return [{"name":name}
-                for (time, name, *rest) in results]
-
-    def list_tables(self, database:str, schema:str):
-        results = self._exec(f"SHOW OBJECTS IN {database}.{schema}")
-        items = []
-        if results:
-            for (time, name, db_name, schema_name, kind, *rest) in results:
-                items.append({"name":name, "kind":kind.lower()})
-        return items
-
-    def schema_info(self, database:str, schema:str, tables:Iterable[str]):
-        app_name = self.get_app_name()
-        # Only pass the db + schema as the identifier so that the resulting identity is correct
-        parser = IdentityParser(f"{database}.{schema}")
+    def schema_info(self, database: str, schema: str, tables: Iterable[str]):
+        """Get detailed schema information including primary keys, foreign keys, and columns."""
+        app_name = self.get_app_name()
+        # Only pass the db + schema as the identifier so that the resulting identity is correct
+        parser = IdentityParser(f"{database}.{schema}")
 
         with debugging.span("schema_info"):
             with debugging.span("primary_keys") as span:
@@ -2556,7 +2431,7 @@ Otherwise, remove it from your '{profile}' configuration profile.
 
             # IdentityParser will parse a single value (with no ".") and store it in this case in the db field
             with debugging.span("columns") as span:
-                tables = ", ".join([f"'{IdentityParser(t).db}'" for t in tables])
+                tables_str = ", ".join([f"'{IdentityParser(t).db}'" for t in tables])
                 query = textwrap.dedent(f"""
                     begin
                         SHOW COLUMNS IN SCHEMA {parser.identity};
@@ -2573,7 +2448,7 @@ Otherwise, remove it from your '{profile}' configuration profile.
                                     ELSE FALSE
                                 END as "supported_type"
                             FROM table(result_scan(-1)) as t
-                            WHERE "table_name" in ({tables})
+                            WHERE "table_name" in ({tables_str})
                         );
                         return table(r);
                     end;
@@ -2584,7 +2459,7 @@ Otherwise, remove it from your '{profile}' configuration profile.
             results = defaultdict(lambda: {"pks": [], "fks": {}, "columns": {}, "invalid_columns": {}})
             if pks:
                 for row in pks:
-                    results[row[3]]["pks"].append(row[4]) # type: ignore
+                    results[row[3]]["pks"].append(row[4])  # type: ignore
             if fks:
                 for row in fks:
                     results[row[7]]["fks"][row[8]] = row[3]
@@ -2726,7 +2601,7 @@ class SnowflakeSchema:
                 tables_with_invalid_columns[table_name] = table_info["invalid_columns"]
 
         if tables_with_invalid_columns:
-            from ..errors import UnsupportedColumnTypesWarning
+            from relationalai.errors import UnsupportedColumnTypesWarning
             UnsupportedColumnTypesWarning(tables_with_invalid_columns)
 
     def _add(self, name, is_imported=False):
@@ -2935,10 +2810,14 @@ class Provider(ProviderBase):
         if resources:
             self.resources = resources
         else:
-            resource_class = Resources
-            if config and config.get("use_direct_access", USE_DIRECT_ACCESS):
-                resource_class = DirectAccessResources
-            self.resources = resource_class(profile=profile, config=config, generation=generation)
+            from .resources_factory import create_resources_instance
+            self.resources = create_resources_instance(
+                config=config,
+                profile=profile,
+                generation=generation or Generation.V0,
+                dry_run=False,
+                language="rel",
+            )
 
     def list_streams(self, model:str):
         return self.resources.list_imports(model=model)
@@ -3107,19 +2986,31 @@ def Graph(
     nowait_durable: bool = True,
     format: str = "default",
 ):
+    from .resources_factory import create_resources_instance
+    from .use_index_resources import UseIndexResources
 
-    client_class = Client
-    resource_class = Resources
     use_graph_index = config.get("use_graph_index", USE_GRAPH_INDEX)
     use_monotype_operators = config.get("compiler.use_monotype_operators", False)
-    use_direct_access = config.get("use_direct_access", USE_DIRECT_ACCESS)
 
-    if use_graph_index:
+    # Create resources instance using factory
+    resources = create_resources_instance(
+        config=config,
+        profile=profile,
+        connection=connection,
+        generation=Generation.V0,
+        dry_run=False,  # Resources instance dry_run is separate from client dry_run
+        language="rel",
+    )
+
+    # Determine client class based on resources type and config
+    # SnowflakeClient is used for resources that support use_index functionality
+    if use_graph_index or isinstance(resources, UseIndexResources):
         client_class = SnowflakeClient
-    if use_direct_access:
-        resource_class = DirectAccessResources
+    else:
+        client_class = Client
+
     client = client_class(
-        resource_class(generation=Generation.V0, profile=profile, config=config, connection=connection),
+        resources,
         rel.Compiler(config),
         name,
         config,
@@ -3190,639 +3081,3 @@ def Graph(
     debugging.set_source(pyrel_base)
     client.install("pyrel_base", pyrel_base)
     return dsl.Graph(client, name, format=format)
-
-
-
-#--------------------------------------------------
-# Direct Access
-#--------------------------------------------------
-# Note: All direct access components should live in a separate file
-
-class DirectAccessResources(Resources):
-    """
-    Resources class for Direct Service Access avoiding Snowflake service functions.
-    """
-    def __init__(
-        self,
-        profile: Union[str, None] = None,
-        config: Union[Config, None] = None,
-        connection: Union[Session, None] = None,
-        dry_run: bool = False,
-        reset_session: bool = False,
-        generation: Optional[Generation] = None,
-        language: str = "rel",
-    ):
-        super().__init__(
-            generation=generation,
-            profile=profile,
-            config=config,
-            connection=connection,
-            reset_session=reset_session,
-            dry_run=dry_run,
-            language=language,
-        )
-        self._endpoint_info = ConfigStore(ENDPOINT_FILE)
-        self._service_endpoint = ""
-        self._direct_access_client = None
-        self.generation = generation
-        self.database = ""
-
-    @property
-    def service_endpoint(self) -> str:
-        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_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)
-        query = f"CALL {self.get_app_name()}.app.service_endpoint({not is_snowflake_notebook});"
-        result = self._exec(query)
-        assert result, f"Could not retrieve service endpoint for {self.get_app_name()}"
-        if is_snowflake_notebook:
-            self._service_endpoint = f"http://{result[0]['SERVICE_ENDPOINT']}"
-        else:
-            self._service_endpoint = f"https://{result[0]['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:
-            self._endpoint_info.save()
-        except Exception:
-            print("Failed to persist endpoints to file. This might slow down future executions.")
-
-        return self._service_endpoint
-
-    @property
-    def direct_access_client(self) -> DirectAccessClient:
-        if self._direct_access_client:
-            return self._direct_access_client
-        try:
-            service_endpoint = self.service_endpoint
-            self._direct_access_client = DirectAccessClient(
-                self.config, self.token_handler, service_endpoint, self.generation,
-            )
-        except Exception as e:
-            raise e
-        return self._direct_access_client
-
-    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,
-        skip_auto_create: bool = False,
-    ) -> requests.Response:
-        with debugging.span("direct_access_request"):
-            def _send_request():
-                return self.direct_access_client.request(
-                    endpoint=endpoint,
-                    payload=payload,
-                    headers=headers,
-                    path_params=path_params,
-                    query_params=query_params,
-                )
-            try:
-                response = _send_request()
-                if response.status_code != 200:
-                    # For 404 responses with skip_auto_create=True, return immediately to let caller handle it
-                    # (e.g., get_engine needs to check 404 and return None for auto_create_engine)
-                    # For skip_auto_create=False, continue to auto-creation logic below
-                    if response.status_code == 404 and skip_auto_create:
-                        return response
-
-                    try:
-                        message = response.json().get("message", "")
-                    except requests.exceptions.JSONDecodeError:
-                        # Can't parse JSON response. For skip_auto_create=True (e.g., get_engine),
-                        # this should have been caught by the 404 check above, so this is an error.
-                        # For skip_auto_create=False, we explicitly check status_code below,
-                        # so we don't need to parse the message.
-                        if skip_auto_create:
-                            raise ResponseStatusException(
-                                f"Failed to parse error response from endpoint {endpoint}.", response
-                            )
-                        message = ""  # Not used when we check status_code directly
-
-                    # fix engine on engine error and retry
-                    # Skip auto-retry if skip_auto_create is True to avoid recursion
-                    if (_is_engine_issue(message) and not skip_auto_create) or _is_database_issue(message):
-                        engine_name = payload.get("caller_engine_name", "") if payload else ""
-                        engine_name = engine_name or self.get_default_engine_name()
-                        engine_size = self.config.get_default_engine_size()
-                        self._poll_use_index(
-                            app_name=self.get_app_name(),
-                            sources=self.sources,
-                            model=self.database,
-                            engine_name=engine_name,
-                            engine_size=engine_size,
-                            headers=headers,
-                        )
-                        response = _send_request()
-            except requests.exceptions.ConnectionError as e:
-                if "NameResolutionError" in str(e):
-                    # when we can not resolve the service endpoint, we assume it is outdated
-                    # hence, we try to retrieve it again and query again.
-                    self.direct_access_client.service_endpoint = self._retrieve_service_endpoint(
-                        enforce_update=True,
-                    )
-                    return _send_request()
-                # raise in all other cases
-                raise e
-            return response
-
-    def _exec_async_v2(
-        self,
-        database: str,
-        engine: Union[str, None],
-        raw_code: str,
-        inputs: Dict | None = None,
-        readonly=True,
-        nowait_durable=False,
-        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:
-            with debugging.span("create_v2") as create_span:
-
-                use_graph_index = self.config.get("use_graph_index", USE_GRAPH_INDEX)
-
-                payload = {
-                    "dbname": database,
-                    "engine_name": engine,
-                    "query": raw_code,
-                    "v1_inputs": inputs,
-                    "nowait_durable": nowait_durable,
-                    "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(
-                    "create_txn", payload=payload, headers=headers, query_params=query_params,
-                )
-
-                if response.status_code != 200:
-                    raise ResponseStatusException("Failed to create transaction.", response)
-
-                artifact_info = {}
-                response_content = response.json()
-
-                txn_id = response_content["transaction"]['id']
-                state = response_content["transaction"]['state']
-
-                txn_span["txn_id"] = txn_id
-                create_span["txn_id"] = txn_id
-                debugging.event("transaction_created", txn_span, txn_id=txn_id)
-
-            # fast path: transaction already finished
-            if state in ["COMPLETED", "ABORTED"]:
-                if txn_id in self._pending_transactions:
-                    self._pending_transactions.remove(txn_id)
-
-                # Process rows to get the rest of the artifacts
-                for result in response_content.get("results", []):
-                    filename = result['filename']
-                    # making keys uppercase to match the old behavior
-                    artifact_info[filename] = {k.upper(): v for k, v in result.items()}
-
-            # Slow path: transaction not done yet; start polling
-            else:
-                self._pending_transactions.append(txn_id)
-                with debugging.span("wait", txn_id=txn_id):
-                    poll_with_specified_overhead(
-                        lambda: self._check_exec_async_status(txn_id, headers=headers), 0.1
-                    )
-                artifact_info = self._list_exec_async_artifacts(txn_id, headers=headers)
-
-            with debugging.span("fetch"):
-                return self._download_results(artifact_info, txn_id, state)
-
-    def _prepare_index(
-        self,
-        model: str,
-        engine_name: str,
-        engine_size: str = "",
-        language: str = "rel",
-        rai_relations: List[str] | None = None,
-        pyrel_program_id: str | None  = None,
-        skip_pull_relations: bool = False,
-        headers: Dict | None = None,
-    ):
-        """
-        Prepare the index for the given engine and model.
-        """
-        with debugging.span("prepare_index"):
-            if headers is None:
-                headers = {}
-
-            payload = {
-                "model_name": model,
-                "caller_engine_name": engine_name,
-                "language": language,
-                "pyrel_program_id": pyrel_program_id,
-                "skip_pull_relations": skip_pull_relations,
-                "rai_relations": rai_relations or [],
-                "user_agent": get_pyrel_version(self.generation),
-            }
-            # Only include engine_size if it has a non-empty string value
-            if engine_size and engine_size.strip():
-                payload["caller_engine_size"] = engine_size
-
-            response = self.request(
-                "prepare_index", payload=payload, headers=headers
-            )
-
-            if response.status_code != 200:
-                raise ResponseStatusException("Failed to prepare index.", response)
-
-            return response.json()
-
-    def _poll_use_index(
-        self,
-        app_name: str,
-        sources: Iterable[str],
-        model: str,
-        engine_name: str,
-        engine_size: str | None = None,
-        program_span_id: str | None = None,
-        headers: Dict | None = None,
-    ):
-        return DirectUseIndexPoller(
-            self,
-            app_name=app_name,
-            sources=sources,
-            model=model,
-            engine_name=engine_name,
-            engine_size=engine_size,
-            language=self.language,
-            program_span_id=program_span_id,
-            headers=headers,
-            generation=self.generation,
-        ).poll()
-
-    def maybe_poll_use_index(
-        self,
-        app_name: str,
-        sources: Iterable[str],
-        model: str,
-        engine_name: str,
-        engine_size: str | None = None,
-        program_span_id: str | None = None,
-        headers: Dict | None = None,
-    ):
-        """Only call poll() if there are sources to process and cache is not valid."""
-        sources_list = list(sources)
-        self.database = model
-        if sources_list:
-            poller = DirectUseIndexPoller(
-                self,
-                app_name=app_name,
-                sources=sources_list,
-                model=model,
-                engine_name=engine_name,
-                engine_size=engine_size,
-                language=self.language,
-                program_span_id=program_span_id,
-                headers=headers,
-                generation=self.generation,
-            )
-            # If cache is valid (data freshness has not expired), skip polling
-            if poller.cache.is_valid():
-                cached_sources = len(poller.cache.sources)
-                total_sources = len(sources_list)
-                cached_timestamp = poller.cache._metadata.get("cachedIndices", {}).get(poller.cache.key, {}).get("last_use_index_update_on", "")
-
-                message = f"Using cached data for {cached_sources}/{total_sources} data streams"
-                if cached_timestamp:
-                    print(f"\n{message} (cached at {cached_timestamp})\n")
-                else:
-                    print(f"\n{message}\n")
-            else:
-                return poller.poll()
-
-    def _check_exec_async_status(self, txn_id: str, headers: Dict[str, str] | None = None) -> bool:
-        """Check whether the given transaction has completed."""
-
-        with debugging.span("check_status"):
-            response = self.request(
-                "get_txn",
-                headers=headers,
-                path_params={"txn_id": txn_id},
-            )
-            assert response, f"No results from get_transaction('{txn_id}')"
-
-        response_content = response.json()
-        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"
-
-    def _list_exec_async_artifacts(self, txn_id: str, headers: Dict[str, str] | None = None) -> Dict[str, Dict]:
-        """Grab the list of artifacts produced in the transaction and the URLs to retrieve their contents."""
-        with debugging.span("list_results"):
-            response = self.request(
-                "get_txn_artifacts",
-                headers=headers,
-                path_params={"txn_id": txn_id},
-            )
-            assert response, f"No results from get_transaction_artifacts('{txn_id}')"
-            artifact_info = {}
-            for result in response.json()["results"]:
-                filename = result['filename']
-                # making keys uppercase to match the old behavior
-                artifact_info[filename] = {k.upper(): v for k, v in result.items()}
-            return artifact_info
-
-    def get_transaction_problems(self, txn_id: str) -> List[Dict[str, Any]]:
-        with debugging.span("get_transaction_problems"):
-            response = self.request(
-                "get_txn_problems",
-                path_params={"txn_id": txn_id},
-            )
-            response_content = response.json()
-            if not response_content:
-                return []
-            return response_content.get("problems", [])
-
-    def get_transaction_events(self, transaction_id: str, continuation_token: str = ''):
-        response = self.request(
-            "get_txn_events",
-            path_params={"txn_id": transaction_id, "stream_name": "profiler"},
-            query_params={"continuation_token": continuation_token},
-        )
-        response_content = response.json()
-        if not response_content:
-            return {
-                "events": [],
-                "continuation_token": None
-            }
-        return response_content
-
-    #--------------------------------------------------
-    # Databases
-    #--------------------------------------------------
-
-    def get_installed_packages(self, database: str) -> Union[Dict, None]:
-        use_graph_index = self.config.get("use_graph_index", USE_GRAPH_INDEX)
-        if use_graph_index:
-            response = self.request(
-                "get_model_package_versions",
-                payload={"model_name": database},
-            )
-        else:
-            response = self.request(
-                "get_package_versions",
-                path_params={"db_name": database},
-            )
-        if response.status_code == 404 and response.json().get("message", "") == "database not found":
-            return None
-        if response.status_code != 200:
-            raise ResponseStatusException(
-                f"Failed to retrieve package versions for {database}.", response
-            )
-
-        content = response.json()
-        if not content:
-            return None
-
-        return safe_json_loads(content["package_versions"])
-
-    def get_database(self, database: str):
-        with debugging.span("get_database", dbname=database):
-            if not database:
-                raise ValueError("Database name must be provided to get database.")
-            response = self.request(
-                "get_db",
-                path_params={},
-                query_params={"name": database},
-            )
-            if response.status_code != 200:
-                raise ResponseStatusException(f"Failed to get db. db:{database}", response)
-
-            response_content = response.json()
-
-            if (response_content.get("databases") and len(response_content["databases"]) == 1):
-                db = response_content["databases"][0]
-                return {
-                    "id": db["id"],
-                    "name": db["name"],
-                    "created_by": db.get("created_by"),
-                    "created_on": ms_to_timestamp(db.get("created_on")),
-                    "deleted_by": db.get("deleted_by"),
-                    "deleted_on": ms_to_timestamp(db.get("deleted_on")),
-                    "state": db["state"],
-                }
-            else:
-                return None
-
-    def create_graph(self, name: str):
-        with debugging.span("create_model", dbname=name):
-            return self._create_database(name,"")
-
-    def delete_graph(self, name:str, force=False, language: str = "rel"):
-        prop_hdrs = debugging.gen_current_propagation_headers()
-        if self.config.get("use_graph_index", USE_GRAPH_INDEX):
-            keep_database = not force and self.config.get("reuse_model", True)
-            with debugging.span("release_index", name=name, keep_database=keep_database, language=language):
-                response = self.request(
-                    "release_index",
-                    payload={
-                        "model_name": name,
-                        "keep_database": keep_database,
-                        "language": language,
-                        "user_agent": get_pyrel_version(self.generation),
-                    },
-                    headers=prop_hdrs,
-                )
-                if (
-                    response.status_code != 200
-                    and not (
-                        response.status_code == 404
-                        and "database not found" in response.json().get("message", "")
-                    )
-                ):
-                    raise ResponseStatusException(f"Failed to release index. Model: {name} ", response)
-        else:
-            with debugging.span("delete_model", name=name):
-                self._delete_database(name, headers=prop_hdrs)
-
-    def clone_graph(self, target_name:str, source_name:str, nowait_durable=True, force=False):
-        if force and self.get_graph(target_name):
-            self.delete_graph(target_name)
-        with debugging.span("clone_model", target_name=target_name, source_name=source_name):
-            return self._create_database(target_name,source_name)
-
-    def _delete_database(self, name:str, headers:Dict={}):
-        with debugging.span("_delete_database", dbname=name):
-            response = self.request(
-                "delete_db",
-                path_params={"db_name": name},
-                query_params={},
-                headers=headers,
-            )
-            if response.status_code != 200:
-                raise ResponseStatusException(f"Failed to delete db. db:{name} ", response)
-
-    def _create_database(self, name:str, source_name:str):
-        with debugging.span("_create_database", dbname=name):
-            payload = {
-                "name": name,
-                "source_name": source_name,
-            }
-            response = self.request(
-                "create_db", payload=payload, headers={}, query_params={},
-            )
-            if response.status_code != 200:
-                raise ResponseStatusException(f"Failed to create db. db:{name}", response)
-
-    #--------------------------------------------------
-    # Engines
-    #--------------------------------------------------
-
-    def list_engines(self, state: str | None = None):
-        response = self.request("list_engines")
-        if response.status_code != 200:
-            raise ResponseStatusException(
-                "Failed to retrieve engines.", response
-            )
-        response_content = response.json()
-        if not response_content:
-            return []
-        engines = [
-            {
-                "name": engine["name"],
-                "id": engine["id"],
-                "size": engine["size"],
-                "state": engine["status"], # callers are expecting 'state'
-                "created_by": engine["created_by"],
-                "created_on": engine["created_on"],
-                "updated_on": engine["updated_on"],
-            }
-            for engine in response_content.get("engines", [])
-            if state is None or engine.get("status") == state
-        ]
-        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"}, skip_auto_create=True)
-        if response.status_code == 404: # engine not found return 404
-            return None
-        elif response.status_code != 200:
-            raise ResponseStatusException(
-                f"Failed to retrieve engine {name}.", response
-            )
-        engine = response.json()
-        if not engine:
-            return None
-        engine_state: EngineState = {
-            "name": engine["name"],
-            "id": engine["id"],
-            "size": engine["size"],
-            "state": engine["status"], # callers are expecting 'state'
-            "created_by": engine["created_by"],
-            "created_on": engine["created_on"],
-            "updated_on": engine["updated_on"],
-            "version": engine["version"],
-            "auto_suspend": engine["auto_suspend_mins"],
-            "suspends_at": engine["suspends_at"],
-        }
-        return engine_state
-
-    def _create_engine(
-            self,
-            name: str,
-            size: str | None = None,
-            auto_suspend_mins: int | None = None,
-            is_async: bool = False,
-            headers: Dict[str, str] | None = None
-        ):
-        # only async engine creation supported via direct access
-        if not is_async:
-            return super()._create_engine(name, size, auto_suspend_mins, is_async, headers=headers)
-        payload:Dict[str, Any] = {
-            "name": name,
-        }
-        if auto_suspend_mins is not None:
-            payload["auto_suspend_mins"] = auto_suspend_mins
-        if size is not None:
-            payload["size"] = size
-        response = self.request(
-            "create_engine",
-            payload=payload,
-            path_params={"engine_type": "logic"},
-            headers=headers,
-            skip_auto_create=True,
-        )
-        if response.status_code != 200:
-            raise ResponseStatusException(
-                f"Failed to create engine {name} with size {size}.", response
-            )
-
-    def delete_engine(self, name:str, force:bool = False, headers={}):
-        response = self.request(
-            "delete_engine",
-            path_params={"engine_name": name, "engine_type": "logic"},
-            headers=headers,
-            skip_auto_create=True,
-        )
-        if response.status_code != 200:
-            raise ResponseStatusException(
-                f"Failed to delete engine {name}.", response
-            )
-
-    def suspend_engine(self, name:str):
-        response = self.request(
-            "suspend_engine",
-            path_params={"engine_name": name, "engine_type": "logic"},
-            skip_auto_create=True,
-        )
-        if response.status_code != 200:
-            raise ResponseStatusException(
-                f"Failed to suspend engine {name}.", response
-            )
-
-    def resume_engine_async(self, name:str, headers={}):
-        response = self.request(
-            "resume_engine",
-            path_params={"engine_name": name, "engine_type": "logic"},
-            headers=headers,
-            skip_auto_create=True,
-        )
-        if response.status_code != 200:
-            raise ResponseStatusException(
-                f"Failed to resume engine {name}.", response
-            )
-        return {}