relationalai 1.0.0a3__py3-none-any.whl → 1.0.0a5__py3-none-any.whl

v0/relationalai/clients/resources/snowflake/snowflake.py

@@ -15,6 +15,7 @@ import hashlib
 from dataclasses import dataclass
 
 from ....auth.token_handler import TokenHandler
+from v0.relationalai.clients.exec_txn_poller import ExecTxnPoller
 import snowflake.snowpark
 
 from ....rel_utils import sanitize_identifier, to_fqn_relation_name
@@ -38,6 +39,7 @@ from ...types import AvailableModel, EngineState, Import, ImportSource, ImportSo
 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 .engine_service import EngineServiceSQL, EngineType
 from .util import (
     collect_error_messages,
     process_jinja_template,
@@ -54,7 +56,7 @@ from .util import (
 )
 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 ....errors import EngineProvisioningFailed, EngineNameValidationException, Errors, GuardRailsException, 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, timedelta
 from snowflake.snowpark.types import StringType, StructField, StructType
@@ -63,6 +65,7 @@ from .error_handlers import (
     ErrorHandler,
     DuoSecurityErrorHandler,
     AppMissingErrorHandler,
+    AppFunctionMissingErrorHandler,
     DatabaseErrorsHandler,
     EngineErrorsHandler,
     ServiceNotStartedErrorHandler,
@@ -89,22 +92,26 @@ from .engine_state_handlers import (
 # Constants
 #--------------------------------------------------
 
-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"]
 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"]
 # Note: ENGINE_ERRORS, ENGINE_NOT_READY_MSGS, DATABASE_ERRORS moved to util.py
 PYREL_ROOT_DB = 'pyrel_root_db'
 
 TERMINAL_TXN_STATES = ["COMPLETED", "ABORTED"]
 
 TXN_ABORT_REASON_TIMEOUT = "transaction timeout"
+GUARDRAILS_ABORT_REASON = "guard rail violation"
+
+PRINT_TXN_PROGRESS_FLAG = "print_txn_progress"
+
+#--------------------------------------------------
+# Helpers
+#--------------------------------------------------
+
+def should_print_txn_progress(config) -> bool:
+    return bool(config.get(PRINT_TXN_PROGRESS_FLAG, False))
 
 #--------------------------------------------------
 # Resources
@@ -131,6 +138,19 @@ class ExecContext:
             skip_engine_db_error_retry=self.skip_engine_db_error_retry
         )
 
+
+@dataclass
+class TxnCreationResult:
+    """Result of creating a transaction via _create_v2_txn.
+
+    This standardizes the response format between different implementations
+    (SQL stored procedure vs HTTP direct access).
+    """
+    txn_id: str
+    state: str
+    artifact_info: Dict[str, Dict]  # Populated if fast-path (state is COMPLETED/ABORTED)
+
+
 class Resources(ResourcesBase):
     def __init__(
         self,
@@ -160,11 +180,17 @@ class Resources(ResourcesBase):
         self._sproc_models = None
         # Store language for backward compatibility (used by child classes for use_index polling)
         self.language = language
+        # Engine subsystem (composition: keeps engine CRUD isolated from the core Resources class)
+        self._engines = EngineServiceSQL(self)
         # Register error and state handlers
         self._register_handlers()
         # Register atexit callback to cancel pending transactions
         atexit.register(self.cancel_pending_transactions)
 
+    @property
+    def engines(self) -> EngineServiceSQL:
+        return self._engines
+
     #--------------------------------------------------
     # Initialization & Properties
     #--------------------------------------------------
@@ -193,11 +219,12 @@ class Resources(ResourcesBase):
                 return handlers
         """
         return [
-            DuoSecurityErrorHandler(),
             AppMissingErrorHandler(),
+            AppFunctionMissingErrorHandler(),
+            ServiceNotStartedErrorHandler(),
+            DuoSecurityErrorHandler(),
             DatabaseErrorsHandler(),
             EngineErrorsHandler(),
-            ServiceNotStartedErrorHandler(),
             TransactionAbortedErrorHandler(),
         ]
 
@@ -589,7 +616,7 @@ class Resources(ResourcesBase):
 
         # Validate engine size
         if engine_size:
-            is_size_valid, sizes = self.validate_engine_size(engine_size)
+            is_size_valid, sizes = self._engines.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:
@@ -679,67 +706,27 @@ class Resources(ResourcesBase):
         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:
-            cloud_provider = self.get_cloud_provider()
-        if cloud_provider == 'azure':
-            sizes = ENGINE_SIZES_AZURE
-        else:
-            sizes = ENGINE_SIZES_AWS
-        if self.config.show_all_engine_sizes():
-            return INTERNAL_ENGINE_SIZES + sizes
-        else:
-            return sizes
-
-    def list_engines(self, state: str | None = None):
-        where_clause = f"WHERE STATUS = '{state.upper()}'" if state else ""
-        statement = f"SELECT NAME, ID, SIZE, STATUS, CREATED_BY, CREATED_ON, UPDATED_ON FROM {APP_NAME}.api.engines {where_clause} ORDER BY NAME ASC;"
-        results = self._exec(statement)
-        if not results:
-            return []
-        return [
-            {
-                "name": row["NAME"],
-                "id": row["ID"],
-                "size": row["SIZE"],
-                "state": row["STATUS"], # callers are expecting 'state'
-                "created_by": row["CREATED_BY"],
-                "created_on": row["CREATED_ON"],
-                "updated_on": row["UPDATED_ON"],
-            }
-            for row in results
-        ]
+        return self._engines.get_engine_sizes(cloud_provider=cloud_provider)
 
-    def get_engine(self, name: str):
-        results = self._exec(
-            f"SELECT NAME, ID, SIZE, STATUS, CREATED_BY, CREATED_ON, UPDATED_ON, VERSION, AUTO_SUSPEND_MINS, SUSPENDS_AT FROM {APP_NAME}.api.engines WHERE NAME='{name}';"
+    def list_engines(
+        self,
+        state: str | None = None,
+        name: str | None = None,
+        type: str | None = None,
+        size: str | None = None,
+        created_by: str | None = None,
+    ):
+        return self._engines.list_engines(
+            state=state,
+            name=name,
+            type=type,
+            size=size,
+            created_by=created_by,
         )
-        if not results:
-            return None
-        engine = results[0]
-        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 get_engine(self, name: str, type: str):
+        return self._engines.get_engine(name, type)
 
     def get_default_engine_name(self) -> str:
         if self.config.get("engine_name", None) is not None:
@@ -760,60 +747,82 @@ Otherwise, remove it from your '{profile}' configuration profile.
     def is_valid_engine_state(self, name:str):
         return name in VALID_ENGINE_STATES
 
+    # Can be overridden by subclasses (e.g. DirectAccessResources)
     def _create_engine(
             self,
             name: str,
+            type: str = EngineType.LOGIC,
             size: str | None = None,
             auto_suspend_mins: int | None= None,
             is_async: bool = False,
             headers: Dict | None = None,
+            settings: Dict[str, Any] | None = None,
         ):
-        api = "create_engine_async" if is_async else "create_engine"
-        if size is None:
-            size = self.config.get_default_engine_size()
-        # if auto_suspend_mins is None, get the default value from the config
-        if auto_suspend_mins is None:
-            auto_suspend_mins = self.config.get_default_auto_suspend_mins()
-        try:
-            headers = debugging.gen_current_propagation_headers()
-            with debugging.span(api, name=name, size=size, auto_suspend_mins=auto_suspend_mins):
-                # check in case the config default is missing
-                if auto_suspend_mins is None:
-                    self._exec(f"call {APP_NAME}.api.{api}('{name}', '{size}', null, {headers});")
-                else:
-                    self._exec(f"call {APP_NAME}.api.{api}('{name}', '{size}', PARSE_JSON('{{\"auto_suspend_mins\": {auto_suspend_mins}}}'), {headers});")
-        except Exception as e:
-            raise EngineProvisioningFailed(name, e) from e
+        return self._engines._create_engine(
+            name=name,
+            type=type,
+            size=size,
+            auto_suspend_mins=auto_suspend_mins,
+            is_async=is_async,
+            headers=headers,
+            settings=settings,
+        )
 
-    def create_engine(self, name:str, size:str|None=None, auto_suspend_mins:int|None=None, headers: Dict | None = None):
-        self._create_engine(name, size, auto_suspend_mins, headers=headers)
+    def create_engine(
+        self,
+        name: str,
+        type: str | None = None,
+        size: str | None = None,
+        auto_suspend_mins: int | None = None,
+        headers: Dict | None = None,
+        settings: Dict[str, Any] | None = None,
+    ):
+        if type is None:
+            type = EngineType.LOGIC
+        # Route through _create_engine so subclasses (e.g. DirectAccessResources)
+        # can override engine creation behavior.
+        return self._create_engine(
+            name=name,
+            type=type,
+            size=size,
+            auto_suspend_mins=auto_suspend_mins,
+            is_async=False,
+            headers=headers,
+            settings=settings,
+        )
 
-    def create_engine_async(self, name:str, size:str|None=None, auto_suspend_mins:int|None=None):
-        self._create_engine(name, size, auto_suspend_mins, True)
+    def create_engine_async(
+        self,
+        name: str,
+        type: str = EngineType.LOGIC,
+        size: str | None = None,
+        auto_suspend_mins: int | None = None,
+    ):
+        # Route through _create_engine so subclasses (e.g. DirectAccessResources)
+        # can override async engine creation behavior.
+        return self._create_engine(
+            name=name,
+            type=type,
+            size=size,
+            auto_suspend_mins=auto_suspend_mins,
+            is_async=True,
+        )
 
-    def delete_engine(self, name:str, force:bool = False, headers: Dict | None = None):
-        request_headers = debugging.add_current_propagation_headers(headers)
-        self._exec(f"call {APP_NAME}.api.delete_engine('{name}', {force},{request_headers});")
+    def delete_engine(self, name: str, type: str):
+        return self._engines.delete_engine(name, type)
 
-    def suspend_engine(self, name:str):
-        self._exec(f"call {APP_NAME}.api.suspend_engine('{name}');")
+    def suspend_engine(self, name: str, type: str | None = None):
+        return self._engines.suspend_engine(name, type)
 
-    def resume_engine(self, name:str, headers: Dict | None = None) -> Dict:
-        request_headers = debugging.add_current_propagation_headers(headers)
-        self._exec(f"call {APP_NAME}.api.resume_engine('{name}',{request_headers});")
-        # returning empty dict to match the expected return type
-        return {}
+    def resume_engine(self, name: str, type: str | None = None, headers: Dict | None = None) -> Dict:
+        return self._engines.resume_engine(name, type=type, headers=headers)
 
-    def resume_engine_async(self, name:str, headers: Dict | None = None) -> Dict:
-        if headers is None:
-            headers = {}
-        self._exec(f"call {APP_NAME}.api.resume_engine_async('{name}',{headers});")
-        # returning empty dict to match the expected return type
-        return {}
+    def resume_engine_async(self, name: str, type: str | None = None, headers: Dict | None = None) -> Dict:
+        return self._engines.resume_engine_async(name, type=type, headers=headers)
 
     def alter_engine_pool(self, size:str|None=None, mins:int|None=None, maxs:int|None=None):
         """Alter engine pool node limits for Snowflake."""
-        self._exec(f"call {APP_NAME}.api.alter_engine_pool_node_limits('{size}', {mins}, {maxs});")
+        return self._engines.alter_engine_pool(size=size, mins=mins, maxs=maxs)
 
     #--------------------------------------------------
     # Graphs
@@ -1411,15 +1420,18 @@ Otherwise, remove it from your '{profile}' configuration profile.
             if txn_id in self._pending_transactions:
                 self._pending_transactions.remove(txn_id)
 
-        if status == "ABORTED" and response_row.get("ABORT_REASON", "") == TXN_ABORT_REASON_TIMEOUT:
-            config_file_path = getattr(self.config, 'file_path', None)
-            # todo: use the timeout returned alongside the transaction as soon as it's exposed
-            timeout_mins = int(self.config.get("query_timeout_mins", DEFAULT_QUERY_TIMEOUT_MINS) or DEFAULT_QUERY_TIMEOUT_MINS)
-            raise QueryTimeoutExceededException(
-                timeout_mins=timeout_mins,
-                query_id=txn_id,
-                config_file_path=config_file_path,
-            )
+        if status == "ABORTED":
+            if response_row.get("ABORT_REASON", "") == TXN_ABORT_REASON_TIMEOUT:
+                config_file_path = getattr(self.config, 'file_path', None)
+                # todo: use the timeout returned alongside the transaction as soon as it's exposed
+                timeout_mins = int(self.config.get("query_timeout_mins", DEFAULT_QUERY_TIMEOUT_MINS) or DEFAULT_QUERY_TIMEOUT_MINS)
+                raise QueryTimeoutExceededException(
+                    timeout_mins=timeout_mins,
+                    query_id=txn_id,
+                    config_file_path=config_file_path,
+                )
+            elif response_row.get("ABORT_REASON", "") == GUARDRAILS_ABORT_REASON:
+                raise GuardRailsException()
 
         # @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"
@@ -1654,6 +1666,72 @@ Otherwise, remove it from your '{profile}' configuration profile.
             raise Exception("Failed to create transaction")
         return response
 
+    def _create_v2_txn(
+        self,
+        database: str,
+        engine: str | None,
+        raw_code: str,
+        inputs: Dict,
+        headers: Dict[str, str],
+        readonly: bool,
+        nowait_durable: bool,
+        bypass_index: bool,
+        language: str,
+        query_timeout_mins: int | None,
+    ) -> TxnCreationResult:
+        """
+        Create a transaction and return the result.
+
+        This method handles calling the RAI app stored procedure to create a transaction
+        and parses the response into a standardized TxnCreationResult format.
+
+        This method can be overridden by subclasses (e.g., DirectAccessResources)
+        to use different transport mechanisms (HTTP instead of SQL).
+
+        Args:
+            database: Database/model name
+            engine: Engine name (optional)
+            raw_code: Code to execute (REL, LQP, or SQL)
+            inputs: Input parameters for the query
+            headers: HTTP headers (must be prepared by caller)
+            readonly: Whether the transaction is read-only
+            nowait_durable: Whether to wait for durable writes
+            bypass_index: Whether to bypass graph index setup
+            language: Query language ("rel" or "lqp")
+            query_timeout_mins: Optional query timeout in minutes
+
+        Returns:
+            TxnCreationResult containing txn_id, state, and artifact_info
+        """
+        response = self._exec_rai_app(
+            database=database,
+            engine=engine,
+            raw_code=raw_code,
+            inputs=inputs,
+            readonly=readonly,
+            nowait_durable=nowait_durable,
+            request_headers=headers,
+            bypass_index=bypass_index,
+            language=language,
+            query_timeout_mins=query_timeout_mins,
+        )
+
+        rows = list(iter(response))
+
+        # process the first row since txn_id and state are the same for all rows
+        first_row = rows[0]
+        txn_id = first_row['ID']
+        state = first_row['STATE']
+
+        # Build artifact_info if transaction completed immediately (fast path)
+        artifact_info: Dict[str, Dict] = {}
+        if state in ["COMPLETED", "ABORTED"]:
+            for row in rows:
+                filename = row['FILENAME']
+                artifact_info[filename] = row
+
+        return TxnCreationResult(txn_id=txn_id, state=state, artifact_info=artifact_info)
+
     def _exec_async_v2(
         self,
         database: str,
@@ -1672,15 +1750,20 @@ Otherwise, remove it from your '{profile}' configuration profile.
         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
+        1. Creates a transaction by calling _create_v2_txn
         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).
+        This method is called by _execute_code (base implementation), and calls the
+        following methods that can be overridden by child classes (e.g.,
+        DirectAccessResources uses HTTP instead):
+            - _create_v2_txn
+            - _check_exec_async_status
+            - _list_exec_async_artifacts
+            - _download_results
 
         Args:
             database: Database/model name
@@ -1704,57 +1787,62 @@ Otherwise, remove it from your '{profile}' configuration profile.
         query_attrs_dict = json.loads(request_headers.get("X-Query-Attributes", "{}"))
 
         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,
-                    raw_code=raw_code,
-                    inputs=inputs,
-                    readonly=readonly,
-                    nowait_durable=nowait_durable,
-                    request_headers=request_headers,
-                    bypass_index=bypass_index,
-                    language=language,
-                    query_timeout_mins=query_timeout_mins,
-                )
+            txn_start_time = time.time()
+            print_txn_progress = should_print_txn_progress(self.config)
+
+            with ExecTxnPoller(
+                print_txn_progress=print_txn_progress,
+                resource=self, txn_id=None, headers=request_headers,
+                txn_start_time=txn_start_time
+            ) as poller:
+                with debugging.span("create_v2", **query_attrs_dict) as create_span:
+                    # Prepare headers for transaction creation
+                    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 ""
+
+                    # Create the transaction
+                    result = self._create_v2_txn(
+                        database=database,
+                        engine=engine,
+                        raw_code=raw_code,
+                        inputs=inputs,
+                        headers=request_headers,
+                        readonly=readonly,
+                        nowait_durable=nowait_durable,
+                        bypass_index=bypass_index,
+                        language=language,
+                        query_timeout_mins=query_timeout_mins,
+                    )
 
-                artifact_info = {}
-                rows = list(iter(response))
+                    txn_id = result.txn_id
+                    state = result.state
 
-                # process the first row since txn_id and state are the same for all rows
-                first_row = rows[0]
-                txn_id = first_row['ID']
-                state = first_row['STATE']
-                filename = first_row['FILENAME']
+                    txn_span["txn_id"] = txn_id
+                    create_span["txn_id"] = txn_id
+                    debugging.event("transaction_created", txn_span, txn_id=txn_id)
 
-                txn_span["txn_id"] = txn_id
-                create_span["txn_id"] = txn_id
-                debugging.event("transaction_created", txn_span, txn_id=txn_id)
+                # Set the transaction ID now that we have it, to update the progress text
+                poller.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)
+                # 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 row in rows:
-                    filename = row['FILENAME']
-                    artifact_info[filename] = row
+                    artifact_info = result.artifact_info
 
-            # 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=request_headers), 0.1
-                    )
-                artifact_info = self._list_exec_async_artifacts(txn_id, headers=request_headers)
+                # Slow path: transaction not done yet; start polling
+                else:
+                    self._pending_transactions.append(txn_id)
+                    # Use the interactive poller for transaction status
+                    with debugging.span("wait", txn_id=txn_id):
+                        poller.poll()
+
+                    artifact_info = self._list_exec_async_artifacts(txn_id, headers=request_headers)
 
-            with debugging.span("fetch"):
-                return self._download_results(artifact_info, txn_id, state)
+                with debugging.span("fetch"):
+                    return self._download_results(artifact_info, txn_id, state)
 
     def get_user_based_engine_name(self):
         if not self._session:
@@ -1764,11 +1852,17 @@ Otherwise, remove it from your '{profile}' configuration profile.
         assert isinstance(user, str), f"current_user() must return a string, not {type(user)}"
         return _sanitize_user_name(user)
 
-    def is_engine_ready(self, engine_name: str):
-        engine = self.get_engine(engine_name)
+    def is_engine_ready(self, engine_name: str, type: str = EngineType.LOGIC):
+        engine = self.get_engine(engine_name, type)
         return engine and engine["state"] == "READY"
 
-    def auto_create_engine(self, name: str | None = None, size: str | None = None, headers: Dict | None = None):
+    def auto_create_engine(
+        self,
+        name: str | None = None,
+        type: str = EngineType.LOGIC,
+        size: str | None = None,
+        headers: Dict | None = None,
+    ):
         """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()
@@ -1776,18 +1870,19 @@ Otherwise, remove it from your '{profile}' configuration profile.
                 return active
 
             # Resolve and validate parameters
-            engine_name, engine_size = self._prepare_engine_params(name, size)
+            name, size = self._prepare_engine_params(name, size)
 
             try:
                 # Get current engine state
-                engine = self.get_engine(engine_name)
+                engine = self.get_engine(name, type)
                 if engine:
                     span.update(cast(dict, engine))
 
                 # Create context for state handling
                 context = EngineContext(
-                    engine_name=engine_name,
-                    engine_size=engine_size,
+                    name=name,
+                    size=size,
+                    type=type,
                     headers=headers,
                     requested_size=size,
                     span=span,
@@ -1797,12 +1892,14 @@ Otherwise, remove it from your '{profile}' configuration profile.
                 self._process_engine_state(engine, context, self._sync_engine_state_handlers)
 
             except Exception as e:
-                self._handle_engine_creation_errors(e, engine_name)
+                self._handle_engine_creation_errors(e, name)
 
-            return engine_name
+            return name
 
-    def auto_create_engine_async(self, name: str | None = None):
+    def auto_create_engine_async(self, name: str | None = None, type: str | None = None):
         """Asynchronously create/ensure an engine, returns immediately."""
+        if type is None:
+            type = EngineType.LOGIC
         active = self._get_active_engine()
         if active and (active == name or name is None):
             return active
@@ -1813,16 +1910,17 @@ Otherwise, remove it from your '{profile}' configuration profile.
         ) as spinner:
             with debugging.span("auto_create_engine_async", active=self._active_engine):
                 # Resolve and validate parameters (use_default_size=True for async)
-                engine_name, engine_size = self._prepare_engine_params(name, None, use_default_size=True)
+                name, size = self._prepare_engine_params(name, None, use_default_size=True)
 
                 try:
                     # Get current engine state
-                    engine = self.get_engine(engine_name)
+                    engine = self.get_engine(name, type)
 
                     # Create context for state handling
                     context = EngineContext(
-                        engine_name=engine_name,
-                        engine_size=engine_size,
+                        name=name,
+                        size=size,
+                        type=type,
                         headers=None,
                         requested_size=None,
                         spinner=spinner,
@@ -1833,11 +1931,11 @@ Otherwise, remove it from your '{profile}' configuration profile.
 
                 except Exception as e:
                     spinner.update_messages({
-                        "finished_message": f"Failed to create engine {engine_name}",
+                        "finished_message": f"Failed to create engine {name}",
                     })
-                    self._handle_engine_creation_errors(e, engine_name, preserve_rai_exception=True)
+                    self._handle_engine_creation_errors(e, name, preserve_rai_exception=True)
 
-                return engine_name
+                return name
 
     #--------------------------------------------------
     # Exec
@@ -2408,7 +2506,7 @@ Otherwise, remove it from your '{profile}' configuration profile.
             return None
         return results[0][0]
 
-    # CLI methods (list_warehouses, list_compute_pools, list_roles, list_apps, 
+    # 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()