# pyright: reportUnusedExpression=false
from __future__ import annotations
import base64
import importlib.resources
import io
import re
import json
import time
import textwrap
import ast
import uuid
import warnings
import atexit
import hashlib
from dataclasses import dataclass

from ....auth.token_handler import TokenHandler
import snowflake.snowpark

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
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, Tuple, List, Literal, cast

from pandas import DataFrame

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, 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,
)


#--------------------------------------------------
# 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"

#--------------------------------------------------
# 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,
        profile: str | None = None,
        config: Config | None = None,
        connection: Session | None = None,
        dry_run: bool = False,
        reset_session: bool = False,
        generation: Generation | None = None,
        language: str = "rel",  # Accepted for backward compatibility, but not stored in base class
    ):
        super().__init__(profile, config=config)
        self._token_handler: TokenHandler | None = None
        self._session = connection
        self.generation = generation
        if self._session is None and not dry_run:
            try:
                # we may still be constructing the config, so this can fail now,
                # if so we'll create later
                self._session = self.get_sf_session(reset_session)
            except Exception:
                pass
        self._pending_transactions: list[str] = []
        self._ns_cache = {}
        # self.sources contains fully qualified Snowflake table/view names
        self.sources: set[str] = set()
        self._sproc_models = None
        # 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();"
        try:
            result = self._exec(query)
            # The result is a list of dictionaries, each with a "STATUS" key
            # The column name containing the result is "SERVICE_STATUS"
            services_status = json.loads(result[0]["SERVICE_STATUS"])
            # Find the dictionary with "name" of "main" and check if its "status" is "READY"
            for service in services_status:
                if service.get("name") == "main" and service.get("status") == "READY":
                    return True
            return False
        except Exception:
            return False

    def get_sf_session(self, reset_session: bool = False):
        if self._session:
            return self._session

        if isinstance(runtime_env, HexEnvironment):
            raise HexSessionException()
        if isinstance(runtime_env, SnowbookEnvironment):
            return get_active_session()
        else:
            # if there's already been a session created, try using that
            # if reset_session is true always try to get the new session
            if not reset_session:
                try:
                    return get_active_session()
                except Exception:
                    pass

            # otherwise, create a new session
            missing_keys = []
            connection_parameters = {}

            authenticator = self.config.get('authenticator', None)
            passcode = self.config.get("passcode", "")
            private_key_file = self.config.get("private_key_file", "")

            # If the authenticator is not set, we need to set it based on the provided parameters
            if authenticator is None:
                if private_key_file != "":
                    authenticator = "snowflake_jwt"
                elif passcode != "":
                    authenticator = "username_password_mfa"
                else:
                    authenticator = "snowflake"
                # set the default authenticator in the config so we can skip it when we check for missing keys
                self.config.set("authenticator", authenticator)

            if authenticator in SNOWFLAKE_AUTHS:
                required_keys = {
                    key for key, value in SNOWFLAKE_AUTHS[authenticator].items() if value.get("required", True)
                }
                for key in required_keys:
                    if self.config.get(key, None) is None:
                        default = SNOWFLAKE_AUTHS[authenticator][key].get("value", None)
                        if default is None or default == FIELD_PLACEHOLDER:
                            # No default value and no value in the config, add to missing keys
                            missing_keys.append(key)
                        else:
                            # Set the default value in the config from the auth defaults
                            self.config.set(key, default)
                if missing_keys:
                    profile = getattr(self.config, 'profile', None)
                    config_file_path = getattr(self.config, 'file_path', None)
                    raise SnowflakeMissingConfigValuesException(missing_keys, profile, config_file_path)
                for key in SNOWFLAKE_AUTHS[authenticator]:
                    connection_parameters[key] = self.config.get(key, None)
            else:
                raise ValueError(f'Authenticator "{authenticator}" not supported')

            return self._build_snowflake_session(connection_parameters)

    def _build_snowflake_session(self, connection_parameters: Dict[str, Any]) -> Session:
        try:
            tmp = {
                "client_session_keep_alive": True,
                "client_session_keep_alive_heartbeat_frequency": 60 * 5,
            }
            tmp.update(connection_parameters)
            connection_parameters = tmp
            # authenticator programmatic access token needs to be upper cased to work...
            connection_parameters["authenticator"] = connection_parameters["authenticator"].upper()
            if "authenticator" in connection_parameters and connection_parameters["authenticator"] == "OAUTH_AUTHORIZATION_CODE":
                # we are replicating OAUTH_AUTHORIZATION_CODE by first retrieving the token
                # and then authenticating with the token via the OAUTH authenticator
                connection_parameters["token"] = self.token_handler.get_session_login_token()
                connection_parameters["authenticator"] = "OAUTH"
            return Session.builder.configs(connection_parameters).create()
        except snowflake.connector.errors.Error as e:
            raise SnowflakeDatabaseException(e)
        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()),
            params
        )
        if raw:
            return sess_results
        return sess_results.collect()

    def _exec(
        self,
        code: str,
        params: List[Any] | Any | None = None,
        raw: bool = False,
        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:
            return self._exec_sql(ctx.code, ctx.params, raw=ctx.raw)
        except Exception as e:
            if not ctx.help:
                raise 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.

        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))

    #--------------------------------------------------
    # Feature Detection & Configuration
    #--------------------------------------------------

    def is_direct_access_enabled(self) -> bool:
        try:
            feature_enabled = self._exec(
                    f"call {APP_NAME}.APP.DIRECT_INGRESS_ENABLED();"
                )
            if not feature_enabled:
                return False

            # Even if the feature is enabled, customers still need to reactivate ERP to ensure the endpoint is available.
            endpoint = self._exec(
                    f"call {APP_NAME}.APP.SERVICE_ENDPOINT(true);"
                )
            if not endpoint or endpoint[0][0] is None:
                return False

            return feature_enabled[0][0]
        except Exception as e:
            raise Exception(f"Unable to determine if direct access is enabled. Details error: {e}") from e


    def is_account_flag_set(self, flag: str) -> bool:
        results = self._exec(
            f"SHOW PARAMETERS LIKE '%{flag}%' IN ACCOUNT;"
        )
        if not results:
            return False
        return results[0]["value"] == "true"

    #--------------------------------------------------
    # Databases
    #--------------------------------------------------

    def get_database(self, database: str):
        try:
            results = self._exec(
                f"call {APP_NAME}.api.get_database('{database}');"
            )
        except Exception as e:
            messages = collect_error_messages(e)
            if any("database does not exist" in msg for msg in messages):
                return None
            raise e

        if not results:
            return None
        db = results[0]
        if not db:
            return None
        return {
            "id": db["ID"],
            "name": db["NAME"],
            "created_by": db["CREATED_BY"],
            "created_on": db["CREATED_ON"],
            "deleted_by": db["DELETED_BY"],
            "deleted_on": db["DELETED_ON"],
            "state": db["STATE"],
        }

    def get_installed_packages(self, database: str) -> Dict | None:
        query = f"call {APP_NAME}.api.get_installed_package_versions('{database}');"
        try:
            results = self._exec(query)
        except Exception as 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 any("unknown user-defined function" in msg for msg in messages):
                return None
            raise e

        if not results:
            return None

        row = results[0]
        if not row:
            return None

        return safe_json_loads(row["PACKAGE_VERSIONS"])

    #--------------------------------------------------
    # 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 v0.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:
            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
        ]

    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}';"
        )
        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_default_engine_name(self) -> str:
        if self.config.get("engine_name", None) is not None:
            profile = self.config.profile
            raise InvalidAliasError(f"""
            'engine_name' is not a valid config option.
If you meant to use a specific engine, use 'engine' instead.
Otherwise, remove it from your '{profile}' configuration profile.
            """)
        engine = self.config.get("engine", None)
        if not engine and self.config.get("user", None):
            engine = _sanitize_user_name(str(self.config.get("user")))
        if not engine:
            engine = self.get_user_based_engine_name()
        self.config.set("engine", engine)
        return engine

    def is_valid_engine_state(self, name:str):
        return name in VALID_ENGINE_STATES

    def _create_engine(
            self,
            name: str,
            size: str | None = None,
            auto_suspend_mins: int | None= None,
            is_async: bool = False,
            headers: Dict | 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

    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_async(self, name:str, size:str|None=None, auto_suspend_mins:int|None=None):
        self._create_engine(name, size, auto_suspend_mins, 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 suspend_engine(self, name:str):
        self._exec(f"call {APP_NAME}.api.suspend_engine('{name}');")

    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_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 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});")

    #--------------------------------------------------
    # Graphs
    #--------------------------------------------------

    def list_graphs(self) -> List[AvailableModel]:
        with debugging.span("list_models"):
            query = textwrap.dedent(f"""
                    SELECT NAME, ID, CREATED_BY, CREATED_ON, STATE, DELETED_BY, DELETED_ON
                    FROM {APP_NAME}.api.databases
                    WHERE state <> 'DELETED'
                    ORDER BY NAME ASC;
                    """)
            results = self._exec(query)
            if not results:
                return []
            return [
                {
                    "name": row["NAME"],
                    "id": row["ID"],
                    "created_by": row["CREATED_BY"],
                    "created_on": row["CREATED_ON"],
                    "state": row["STATE"],
                    "deleted_by": row["DELETED_BY"],
                    "deleted_on": row["DELETED_ON"],
                }
                for row in results
            ]

    def get_graph(self, name: str):
        res = self.get_database(name)
        if res and res.get("state") != "DELETED":
            return res

    def create_graph(self, name: str):
        with debugging.span("create_model", name=name):
            self._exec(f"call {APP_NAME}.api.create_database('{name}', false, {debugging.gen_current_propagation_headers()});")

    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):
                #TODO add headers to release_index
                response = self._exec(f"call {APP_NAME}.api.release_index('{name}', OBJECT_CONSTRUCT('keep_database', {keep_database}, 'language', '{language}', 'user_agent', '{get_pyrel_version(self.generation)}'));")
                if response:
                    result = next(iter(response))
                    obj = json.loads(result["RELEASE_INDEX"])
                    error = obj.get('error', None)
                    if error and "Model database not found" not in error:
                        raise Exception(f"Error releasing index: {error}")
                else:
                    raise Exception("There was no response from the release index call.")
        else:
            with debugging.span("delete_model", name=name):
                self._exec(f"call {APP_NAME}.api.delete_database('{name}', false, {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):
            # not a mistake: the clone_database argument order is indeed target then source:
            headers = debugging.gen_current_propagation_headers()
            self._exec(f"call {APP_NAME}.api.clone_database('{target_name}', '{source_name}', {nowait_durable}, {headers});")

    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,
    ) -> 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,
        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,
    ) -> 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
    #--------------------------------------------------

    def list_models(self, database: str, engine: str):
        pass

    def create_models(self, database: str, engine: str | None, models:List[Tuple[str, str]]) -> List[Any]:
        rel_code = self.create_models_code(models)
        self.exec_raw(database, engine, rel_code, readonly=False)
        # TODO: handle SPCS errors once they're figured out
        return []

    def delete_model(self, database:str, engine:str | None, name:str):
        self.exec_raw(database, engine, f"def delete[:rel, :catalog, :model, \"{name}\"]: rel[:catalog, :model, \"{name}\"]", readonly=False)

    def create_models_code(self, models:List[Tuple[str, str]]) -> str:
        lines = []
        for (name, code) in models:
            name = name.replace("\"", "\\\"")
            assert "\"\"\"\"\"\"\"" not in code, "Code literals must use fewer than 7 quotes."

            lines.append(textwrap.dedent(f"""
            def delete[:rel, :catalog, :model, "{name}"]: rel[:catalog, :model, "{name}"]
            def insert[:rel, :catalog, :model, "{name}"]: raw\"\"\"\"\"\"\"
            """) + code + "\n\"\"\"\"\"\"\"")
        rel_code = "\n\n".join(lines)
        return rel_code

    #--------------------------------------------------
    # Exports
    #--------------------------------------------------

    def list_exports(self, database: str, engine: str):
        return []


    def get_export_code(self, params: ExportParams, all_installs):
        sql_inputs = ", ".join([f"{name} {type_to_sql(type)}" for (name, _, type) in params.inputs])
        input_names = [name for (name, *_) in params.inputs]
        has_return_hint = params.out_fields and isinstance(params.out_fields[0], tuple)
        if has_return_hint:
            sql_out = ", ".join([f"\"{name}\" {type_to_sql(type)}" for (name, type) in params.out_fields])
            sql_out_names = ", ".join([f"('{name}', '{type_to_sql(type)}')" for (ix, (name, type)) in enumerate(params.out_fields)])
            py_outs = ", ".join([f"StructField(\"{name}\", {type_to_snowpark(type)})" for (name, type) in params.out_fields])
        else:
            sql_out = ""
            sql_out_names = ", ".join([f"'{name}'" for name in params.out_fields])
            py_outs = ", ".join([f"StructField(\"{name}\", {type_to_snowpark(str)})" for name in params.out_fields])
        py_inputs = ", ".join([name for (name, *_) in params.inputs])
        safe_rel = escape_for_f_string(params.code).strip()
        clean_inputs = []
        for (name, var, type) in params.inputs:
            if type is str:
                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 = 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)
        file = "export_procedure.py.jinja"
        with importlib.resources.open_text(
            "relationalai.clients.resources.snowflake", file
        ) as f:
            template = f.read()
        def quote(s: str, f = False) -> str:
            return '"' + s + '"' if not f else 'f"' + s + '"'

        wait_for_stream_sync = self.config.get("wait_for_stream_sync", WAIT_FOR_STREAM_SYNC)
        # 1. Check the sources for staled sources
        # 2. Get the object references for the sources
        # TODO: this could be optimized to do it in the run time of the stored procedure
        #       instead of doing it here. It will make it more reliable when sources are
        #       modified after the stored procedure is created.
        checked_sources = self._check_source_updates(self.sources)
        source_obj_references = self._get_source_references(checked_sources)

        # Escape double quotes in the source object references
        escaped_source_obj_references = [source.replace('"', '\\"') for source in source_obj_references]
        escaped_proc_database = params.proc_database.replace('"', '\\"')

        normalized_func_name = IdentityParser(params.func_name).identity
        assert normalized_func_name is not None, "Function name must be set"
        skip_invalid_data = params.skip_invalid_data
        python_code = process_jinja_template(
            template,
            func_name=quote(normalized_func_name),
            database=quote(params.root_database),
            proc_database=quote(escaped_proc_database),
            engine=quote(params.engine),
            rel_code=quote(safe_rel, f=True),
            APP_NAME=quote(APP_NAME),
            input_names=input_names,
            outputs=sql_out,
            sql_out_names=sql_out_names,
            clean_inputs=clean_inputs,
            py_inputs=py_inputs,
            py_outs=py_outs,
            skip_invalid_data=skip_invalid_data,
            source_references=", ".join(escaped_source_obj_references),
            install_code=all_installs.replace("\\", "\\\\").replace("\n", "\\n"),
            has_return_hint=has_return_hint,
            wait_for_stream_sync=wait_for_stream_sync,
        ).strip()
        return_clause = f"TABLE({sql_out})" if sql_out else "STRING"
        destination_input = "" if sql_out else "save_as_table STRING DEFAULT NULL,"
        module_name = sanitize_module_name(normalized_func_name)
        stage = f"@{self.get_app_name()}.app_state.stored_proc_code_stage"
        file_loc = f"{stage}/{module_name}.py"
        python_code = python_code.replace(APP_NAME, self.get_app_name())

        hash = hashlib.sha256()
        hash.update(python_code.encode('utf-8'))
        code_hash = hash.hexdigest()
        print(code_hash)

        sql_code = textwrap.dedent(f"""
            CREATE OR REPLACE PROCEDURE {normalized_func_name}({sql_inputs}{sql_inputs and ',' or ''} {destination_input} engine STRING DEFAULT NULL)
            RETURNS {return_clause}
            LANGUAGE PYTHON
            RUNTIME_VERSION = '3.10'
            IMPORTS = ('{file_loc}')
            PACKAGES = ('snowflake-snowpark-python')
            HANDLER = 'checked_handle'
            EXECUTE AS CALLER
            AS
            $$
            import {module_name}
            import inspect, hashlib, os, sys
            def checked_handle(*args, **kwargs):
                import_dir = sys._xoptions["snowflake_import_directory"]
                wheel_path = os.path.join(import_dir, '{module_name}.py')
                h = hashlib.sha256()
                with open(wheel_path, 'rb') as f:
                    for chunk in iter(lambda: f.read(1<<20), b''):
                        h.update(chunk)
                code_hash = h.hexdigest()
                if code_hash != '{code_hash}':
                    raise RuntimeError("Code hash mismatch. The code has been modified since it was uploaded.")
                # Call the handle function with the provided arguments
                return {module_name}.handle(*args, **kwargs)

            $$;
        """)
        # print(f"\n--- python---\n{python_code}\n--- end python---\n")
        # This check helps catch invalid code early and for dry runs:
        try:
            ast.parse(python_code)
        except SyntaxError:
            raise ValueError(f"Internal error: invalid Python code generated:\n{python_code}")
        return (sql_code, python_code, file_loc)

    def get_sproc_models(self, params: ExportParams):
        if self._sproc_models is not None:
            return self._sproc_models

        with debugging.span("get_sproc_models"):
            code = """
            def output(name, model):
                rel(:catalog, :model, name, model)
                and not starts_with(name, "rel/")
                and not starts_with(name, "pkg/rel")
                and not starts_with(name, "pkg/std")
                and starts_with(name, "pkg/")
            """
            res = self.exec_raw(params.model_database, params.engine, code, readonly=True, nowait_durable=True)
            df, errors = result_helpers.format_results(res, None, ["name", "model"])
            models = []
            for row in df.itertuples():
                models.append((row.name, row.model))
            self._sproc_models = models
            return models

    def create_export(self, params: ExportParams):
        with debugging.span("create_export") as span:
            if params.dry_run:
                (sql_code, python_code, file_loc) = self.get_export_code(params, params.install_code)
                span["sql"] = sql_code
                return

            start = time.perf_counter()
            use_graph_index = self.config.get("use_graph_index", USE_GRAPH_INDEX)
            # for the non graph index case we need to create the cloned proc database
            if not use_graph_index:
                raise RAIException(
                    "To ensure permissions are properly accounted for, stored procedures require using the graph index. "
                    "Set use_graph_index=True in your config to proceed."
                )

            models = self.get_sproc_models(params)
            lib_installs = self.create_models_code(models)
            all_installs = lib_installs + "\n\n" + params.install_code

            (sql_code, python_code, file_loc) = self.get_export_code(params, all_installs)

            span["sql"] = sql_code
            assert self._session

            with debugging.span("upload_sproc_code"):
                code_bytes = python_code.encode('utf-8')
                code_stream = io.BytesIO(code_bytes)
                self._session.file.put_stream(code_stream, file_loc, auto_compress=False, overwrite=True)

            with debugging.span("sql_install"):
                self._exec(sql_code)

            debugging.time("export", time.perf_counter() - start, DataFrame(), code=sql_code.replace(APP_NAME, self.get_app_name()))


    def create_export_table(self, database: str, engine: str, table: str, relation: str, columns: Dict[str, str], code: str, refresh: str|None=None):
        print("Snowflake doesn't support creating export tables yet. Try creating the table manually first.")
        pass

    def delete_export(self, database: str, engine: str, name: str):
        pass

    #--------------------------------------------------
    # Imports
    #--------------------------------------------------


    def change_stream_status(self, stream_id: str, model:str, suspend: bool):
        if stream_id and model:
            if suspend:
                self._exec(f"CALL {APP_NAME}.api.suspend_data_stream('{stream_id}', '{model}');")
            else:
                self._exec(f"CALL {APP_NAME}.api.resume_data_stream('{stream_id}', '{model}');")

    def change_imports_status(self, suspend: bool):
        if suspend:
            self._exec(f"CALL {APP_NAME}.app.suspend_cdc();")
        else:
            self._exec(f"CALL {APP_NAME}.app.resume_cdc();")

    def get_imports_status(self) -> ImportsStatus|None:
        # NOTE: We expect there to only ever be one result?
        results = self._exec(f"CALL {APP_NAME}.app.cdc_status();")
        if results:
            result = next(iter(results))
            engine = result['CDC_ENGINE_NAME']
            engine_status = result['CDC_ENGINE_STATUS']
            engine_size = result['CDC_ENGINE_SIZE']
            task_status = result['CDC_TASK_STATUS']
            info = result['CDC_TASK_INFO']
            enabled = result['CDC_ENABLED']
            return {"engine": engine, "engine_size": engine_size, "engine_status": engine_status, "status": task_status, "enabled": enabled, "info": info }
        return None

    def set_imports_engine_size(self, size:str):
        try:
            self._exec(f"CALL {APP_NAME}.app.alter_cdc_engine_size('{size}');")
        except Exception as e:
            raise e

    def list_imports(
        self,
        id:str|None = None,
        name:str|None = None,
        model:str|None = None,
        status:str|None = None,
        creator:str|None = None,
    ) -> list[Import]:
        where = []
        if id and isinstance(id, str):
            where.append(f"LOWER(ID) = '{id.lower()}'")
        if name and isinstance(name, str):
            where.append(f"LOWER(FQ_OBJECT_NAME) = '{name.lower()}'")
        if model and isinstance(model, str):
            where.append(f"LOWER(RAI_DATABASE) = '{model.lower()}'")
        if creator and isinstance(creator, str):
            where.append(f"LOWER(CREATED_BY) = '{creator.lower()}'")
        if status and isinstance(status, str):
            where.append(f"LOWER(batch_status) = '{status.lower()}'")
        where_clause = " AND ".join(where)

        # This is roughly inspired by the native app code because we don't have a way to
        # get the status of multiple streams at once and doing them individually is way
        # too slow. We use window functions to get the status of the stream and the batch
        # details.
        statement = f"""
            SELECT
                ID,
                RAI_DATABASE,
                FQ_OBJECT_NAME,
                CREATED_AT,
                CREATED_BY,
                CASE
                    WHEN nextBatch.quarantined > 0 THEN 'quarantined'
                    ELSE nextBatch.status
                END as batch_status,
                nextBatch.processing_errors,
                nextBatch.batches
            FROM {APP_NAME}.api.data_streams as ds
            LEFT JOIN (
                SELECT DISTINCT
                    data_stream_id,
                    -- Get status from the progress record using window functions
                    FIRST_VALUE(status) OVER (
                        PARTITION BY data_stream_id
                        ORDER BY
                            CASE WHEN unloaded IS NOT NULL THEN 1 ELSE 0 END DESC,
                            unloaded ASC
                    ) as status,
                    -- Get batch_details from the same record
                    FIRST_VALUE(batch_details) OVER (
                        PARTITION BY data_stream_id
                        ORDER BY
                            CASE WHEN unloaded IS NOT NULL THEN 1 ELSE 0 END DESC,
                            unloaded ASC
                    ) as batch_details,
                    -- Aggregate the other fields
                    FIRST_VALUE(processing_details:processingErrors) OVER (
                        PARTITION BY data_stream_id
                        ORDER BY
                            CASE WHEN unloaded IS NOT NULL THEN 1 ELSE 0 END DESC,
                            unloaded ASC
                    ) as processing_errors,
                    MIN(unloaded) OVER (PARTITION BY data_stream_id) as unloaded,
                    COUNT(*) OVER (PARTITION BY data_stream_id) as batches,
                    COUNT_IF(status = 'quarantined') OVER (PARTITION BY data_stream_id) as quarantined
                FROM {APP_NAME}.api.data_stream_batches
            ) nextBatch
            ON ds.id = nextBatch.data_stream_id
            {f"where {where_clause}" if where_clause else ""}
            ORDER BY FQ_OBJECT_NAME ASC;
        """
        results = self._exec(statement)
        items = []
        if results:
            for stream in results:
                (id, db, name, created_at, created_by, status, processing_errors, batches) = stream
                if status and isinstance(status, str):
                    status = status.upper()
                if processing_errors:
                    if status in ["QUARANTINED", "PENDING"]:
                        start = processing_errors.rfind("Error")
                        if start != -1:
                            processing_errors = processing_errors[start:-1]
                    else:
                        processing_errors = None
                items.append(cast(Import, {
                    "id": id,
                    "model": db,
                    "name": name,
                    "created": created_at,
                    "creator": created_by,
                    "status": status.upper() if status else None,
                    "errors": processing_errors if processing_errors != "[]" else None,
                    "batches": f"{batches}" if batches else "",
                }))
        return items

    def poll_imports(self, sources:List[str], model:str):
        source_set = self._create_source_set(sources)
        def check_imports():
            imports = [
                import_
                for import_ in self.list_imports(model=model)
                if import_["name"] in source_set
            ]
            # loop through printing status for each in the format (index): (name) - (status)
            statuses = [import_["status"] for import_ in imports]
            if all(status == "LOADED" for status in statuses):
                return True
            if any(status == "QUARANTINED" for status in statuses):
                failed_imports = [import_["name"] for import_ in imports if import_["status"] == "QUARANTINED"]
                raise RAIException("Imports failed:" + ", ".join(failed_imports)) from None
            # this check is necessary in case some of the tables are empty;
            # such tables may be synced even though their status is None:
            def synced(import_):
                if import_["status"] == "LOADED":
                    return True
                if import_["status"] is None:
                    import_full_status = self.get_import_stream(import_["name"], model)
                    if import_full_status and import_full_status[0]["data_sync_status"] == "SYNCED":
                        return True
                return False
            if all(synced(import_) for import_ in imports):
                return True
        poll_with_specified_overhead(check_imports, overhead_rate=0.1, max_delay=10)

    def _create_source_set(self, sources: List[str]) -> set:
        return {
            source.upper() if not IdentityParser(source).has_double_quoted_identifier else IdentityParser(source).identity
            for source in sources
        }

    def get_import_stream(self, name:str|None, model:str|None):
        results = self._exec(f"CALL {APP_NAME}.api.get_data_stream('{name}', '{model}');")
        if not results:
            return None
        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."
        object = source.fqn

        # Parse only to the schema level
        schemaParser = IdentityParser(f"{source.database}.{source.schema}")

        if object.lower() in [x["name"].lower() for x in self.list_imports(model=model)]:
            return

        query = f"SHOW OBJECTS LIKE '{source.table}' IN {schemaParser.identity}"

        info = self._exec(query)
        if not info:
            raise ValueError(f"Object {source.table} not found in schema {schemaParser.identity}")
        else:
            data = info[0]
            if not data:
                raise ValueError(f"Object {source.table} not found in {schemaParser.identity}")
            # (time, name, db_name, schema_name, kind, *rest)
            kind = data["kind"]

        relation_name = to_fqn_relation_name(object)

        command = f"""call {APP_NAME}.api.create_data_stream(
            {APP_NAME}.api.object_reference('{kind}', '{object}'),
            '{model}',
            '{relation_name}');"""

        def create_stream(tracking_just_changed=False):
            try:
                self._exec(command)
            except Exception as 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;")
                            create_stream(tracking_just_changed=True)
                        except Exception:
                            pass
                    else:
                        print("\n")
                        exception = SnowflakeChangeTrackingNotEnabledException((object, kind))
                        raise exception from None
                elif any("database does not exist" in msg for msg in messages):
                    print("\n")
                    raise ModelNotFoundException(model) from None
                raise e

        create_stream()

    def create_import_snapshot(self, source:ImportSource, model:str, options: dict|None = None):
        raise Exception("Snowflake integration doesn't support snapshot imports yet")

    def delete_import(self, import_name:str, model:str, force = False):
        engine = self.get_default_engine_name()
        rel_name = to_fqn_relation_name(import_name)
        try:
            self._exec(f"""call {APP_NAME}.api.delete_data_stream(
                '{import_name}',
                '{model}'
            );""")
        except RAIException as err:
            if "streams do not exist" not in str(err) or not force:
                raise

        # if force is true, we delete the leftover relation to free up the name (in case the user re-creates the stream)
        if force:
            self.exec_raw(model, engine, f"""
                declare ::{rel_name}
                def delete[:\"{rel_name}\"]: {{ {rel_name} }}
            """, readonly=False, bypass_index=True)

    #--------------------------------------------------
    # Exec Async
    #--------------------------------------------------

    def _check_exec_async_status(self, txn_id: str, headers: Dict | None = None):
        """Check whether the given transaction has completed."""
        if headers is None:
            headers = {}

        with debugging.span("check_status"):
            response = self._exec(f"CALL {APP_NAME}.api.get_transaction('{txn_id}',{headers});")
            assert response, f"No results from get_transaction('{txn_id}')"

        response_row = next(iter(response)).asDict()
        status: str = response_row['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 response_row.get("ABORT_REASON", "") == TXN_ABORT_REASON_TIMEOUT:
            config_file_path = getattr(self.config, 'file_path', None)
            # todo: use the timeout returned alongside the transaction as soon as it's exposed
            timeout_mins = int(self.config.get("query_timeout_mins", DEFAULT_QUERY_TIMEOUT_MINS) or DEFAULT_QUERY_TIMEOUT_MINS)
            raise QueryTimeoutExceededException(
                timeout_mins=timeout_mins,
                query_id=txn_id,
                config_file_path=config_file_path,
            )

        # @TODO: Find some way to tunnel the ABORT_REASON out. Azure doesn't have this, but it's handy
        return status == "COMPLETED" or status == "ABORTED"


    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."""
        if headers is None:
            headers = {}
        with debugging.span("list_results"):
            response = self._exec(
                f"CALL {APP_NAME}.api.get_own_transaction_artifacts('{txn_id}',{headers});"
            )
            assert response, f"No results from get_own_transaction_artifacts('{txn_id}')"
            return {row["FILENAME"]: row for row in response}

    def _fetch_exec_async_artifacts(
        self, artifact_info: Dict[str, Dict[str, Any]]
    ) -> Dict[str, Any]:
        """Grab the contents of the given artifacts from SF in parallel using threads."""

        with requests.Session() as session:
            def _fetch_data(name_info):
                filename, metadata = name_info

                try:
                    # Extract the presigned URL and encryption material from metadata
                    url_key = self.get_url_key(metadata)
                    presigned_url = metadata[url_key]
                    encryption_material = metadata["ENCRYPTION_MATERIAL"]

                    response = get_with_retries(session, presigned_url, config=self.config)
                    response.raise_for_status()  # Throw if something goes wrong

                    decrypted = self._maybe_decrypt(response.content, encryption_material)
                    return (filename, decrypted)

                except requests.RequestException as e:
                    raise scrub_exception(wrap_with_request_id(e))

            # Create a list of tuples for the map function
            name_info_pairs = list(artifact_info.items())

            with ThreadPoolExecutor(max_workers=5) as executor:
                results = executor.map(_fetch_data, name_info_pairs)

                return {name: data for (name, data) in results}

    def _maybe_decrypt(self, content: bytes, encryption_material: str) -> bytes:
        # Decrypt if encryption material is present
        if encryption_material:
            # if there's no padding, the initial file was empty
            if len(content) == 0:
                return b""

            return decrypt_artifact(content, encryption_material)

        # otherwise, return content directly
        return content

    def _parse_exec_async_results(self, arrow_files: List[Tuple[str, bytes]]):
        """Mimics the logic in _parse_arrow_results of railib/api.py#L303 without requiring a wrapping multipart form."""
        results = []

        for file_name, file_content in arrow_files:
            with pa.ipc.open_stream(file_content) as reader:
                schema = reader.schema
                batches = [batch for batch in reader]
                table = pa.Table.from_batches(batches=batches, schema=schema)
                results.append({"relationId": file_name, "table": table})

        return results

    def _download_results(
        self, artifact_info: Dict[str, Dict], txn_id: str, state: str
    ) -> TransactionAsyncResponse:
        with debugging.span("download_results"):
            # Fetch artifacts
            artifacts = self._fetch_exec_async_artifacts(artifact_info)

            # Directly use meta_json as it is fetched
            meta_json_bytes = artifacts["metadata.json"]

            # Decode the bytes and parse the JSON
            meta_json_str = meta_json_bytes.decode('utf-8')
            meta_json = json.loads(meta_json_str)  # Parse the JSON string

            # Use the metadata to map arrow files to the relations they contain
            try:
                arrow_files_to_relations = {
                    artifact["filename"]: artifact["relationId"]
                    for artifact in meta_json
                }
            except KeyError:
                # TODO: Remove this fallback mechanism later once several engine versions are updated
                arrow_files_to_relations = {
                    f"{ix}.arrow": artifact["relationId"]
                    for ix, artifact in enumerate(meta_json)
                }

            # Hydrate the arrow files into tables
            results = self._parse_exec_async_results(
                [
                    (arrow_files_to_relations[name], content)
                    for name, content in artifacts.items()
                    if name.endswith(".arrow")
                ]
            )

            # Create and return the response
            rsp = TransactionAsyncResponse()
            rsp.transaction = {
                "id": txn_id,
                "state": state,
                "response_format_version": None,
            }
            rsp.metadata = meta_json
            rsp.problems = artifacts.get(
                "problems.json"
            )  # Safely access possible missing keys
            rsp.results = results
            return rsp

    def get_transaction_problems(self, txn_id: str) -> List[Dict[str, Any]]:
        with debugging.span("get_own_transaction_problems"):
            response = self._exec(
                f"select * from table({APP_NAME}.api.get_own_transaction_problems('{txn_id}'));"
            )
            if not response:
                return []
            return response

    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 is_azure_url(metadata['PRESIGNED_URL']):
            return 'PRESIGNED_URL'

        configured = self.config.get("download_url_type", None)
        if configured == "internal":
            return 'PRESIGNED_URL_AP'
        elif configured == "external":
            return "PRESIGNED_URL"

        if is_container_runtime():
            return 'PRESIGNED_URL_AP'

        return 'PRESIGNED_URL'

    def _exec_rai_app(
        self,
        database: str,
        engine: str | None,
        raw_code: str,
        inputs: Dict,
        readonly=True,
        nowait_durable=False,
        request_headers: Dict | None = None,
        bypass_index=False,
        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)
        # Depending on the shape of the input, the behavior of exec_async_v2 changes.
        # When using the new format (with an object), the function retrieves the
        # 'rai' database by hashing the model and username. In contrast, the
        # current version directly uses the passed database value.
        # Therefore, we must use the original exec_async_v2 when not using the
        # graph index to ensure the correct database is utilized.
        use_graph_index = self.config.get("use_graph_index", USE_GRAPH_INDEX)
        if use_graph_index and not bypass_index:
            payload = {
                'database': database,
                'engine': engine,
                'inputs': inputs,
                'readonly': readonly,
                'nowait_durable': nowait_durable,
                'language': language,
                'headers': request_headers
            }
            if query_timeout_mins is not None:
                payload["timeout_mins"] = query_timeout_mins
            sql_string = f"CALL {APP_NAME}.api.exec_async_v2(?, {payload});"
        else:
            if query_timeout_mins is not None:
                sql_string = f"CALL {APP_NAME}.api.exec_async_v2('{database}','{engine}', ?, {inputs}, {readonly}, {nowait_durable}, '{language}', {query_timeout_mins}, {request_headers});"
            else:
                sql_string = f"CALL {APP_NAME}.api.exec_async_v2('{database}','{engine}', ?, {inputs}, {readonly}, {nowait_durable}, '{language}', {request_headers});"
        # 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")
        return response

    def _exec_async_v2(
        self,
        database: str,
        engine: str | None,
        raw_code: str,
        inputs: Dict | None = None,
        readonly=True,
        nowait_durable=False,
        headers: Dict | None = None,
        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)
        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,
                )

                artifact_info = {}
                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']
                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)

            # 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

            # 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)

            with debugging.span("fetch"):
                return self._download_results(artifact_info, txn_id, state)

    def get_user_based_engine_name(self):
        if not self._session:
            self._session = self.get_sf_session()
        user_table = self._session.sql("select current_user()").collect()
        user = user_table[0][0]
        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)
        return engine and engine["state"] == "READY"

    def auto_create_engine(self, name: str | None = None, 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()
            if active:
                return active

            # 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))

                # 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:
                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 active

        with Spinner(
            "Checking engine status",
            leading_newline=True,
        ) 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)

                try:
                    # Get current engine state
                    engine = self.get_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,
                    )

                    # 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,
        engine: str | None,
        raw_code: bytes,
        readonly=True,
        *,
        inputs: Dict | None = None,
        nowait_durable=False,
        headers: Dict | None = None,
        bypass_index=False,
        query_timeout_mins: int | None = None,
    ):
        """Execute LQP code."""
        raw_code_b64 = base64.b64encode(raw_code).decode("utf-8")
        return self._execute_code(
            database, engine, raw_code_b64, inputs, readonly, nowait_durable,
            headers, bypass_index, 'lqp', query_timeout_mins
        )

    def exec_raw(
        self,
        database: str,
        engine: str | None,
        raw_code: str,
        readonly=True,
        *,
        inputs: Dict | None = None,
        nowait_durable=False,
        headers: Dict | None = None,
        bypass_index=False,
        query_timeout_mins: int | None = None,
    ):
        """Execute raw code."""
        raw_code = raw_code.replace("'", "\\'")
        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]]:
        return result_helpers.format_results(results, task)

    #--------------------------------------------------
    # Exec format
    #--------------------------------------------------

    def exec_format(
        self,
        database: str,
        engine: str,
        raw_code: str,
        cols: List[str],
        format: str,
        inputs: Dict | None = None,
        readonly=True,
        nowait_durable=False,
        skip_invalid_data=False,
        headers: Dict | None = None,
        query_timeout_mins: int | None = None,
    ):
        if inputs is None:
            inputs = {}
        if headers is None:
            headers = {}
        if 'user-agent' not in headers:
            headers['user-agent'] = get_pyrel_version(self.generation)
        if query_timeout_mins is None and (timeout_value := self.config.get("query_timeout_mins", DEFAULT_QUERY_TIMEOUT_MINS)) is not None:
            query_timeout_mins = int(timeout_value)
        # TODO: add headers
        start = time.perf_counter()
        output_table = "out" + str(uuid.uuid4()).replace("-", "_")
        temp_table = f"temp_{output_table}"
        use_graph_index = self.config.get("use_graph_index", USE_GRAPH_INDEX)
        txn_id = None
        rejected_rows = None
        col_names_map = None
        artifacts = None
        assert self._session
        temp = self._session.createDataFrame([], StructType([StructField(name, StringType()) for name in cols]))
        with debugging.span("transaction") as txn_span:
            try:
                # In the graph index case we need to use the new exec_into_table proc as it obfuscates the db name
                with debugging.span("exec_format"):
                    if use_graph_index:
                        # we do not provide a default value for query_timeout_mins so that we can control the default on app level
                        if query_timeout_mins is not None:
                            res = self._exec(f"call {APP_NAME}.api.exec_into_table(?, ?, ?, ?, ?, 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])
                        txn_id = json.loads(res[0]["EXEC_INTO_TABLE"])["rai_transaction_id"]
                        rejected_rows = json.loads(res[0]["EXEC_INTO_TABLE"]).get("rejected_rows", [])
                        rejected_rows_count = json.loads(res[0]["EXEC_INTO_TABLE"]).get("rejected_rows_count", 0)
                    else:
                        if query_timeout_mins is not None:
                            res = self._exec(f"call {APP_NAME}.api.exec_into(?, ?, ?, ?, ?, {inputs}, ?, {headers}, ?, ?);", [database, engine, raw_code, output_table, readonly, nowait_durable, skip_invalid_data, query_timeout_mins])
                        else:
                            res = self._exec(f"call {APP_NAME}.api.exec_into(?, ?, ?, ?, ?, {inputs}, ?, {headers}, ?);", [database, engine, raw_code, output_table, readonly, nowait_durable, skip_invalid_data])
                        txn_id = json.loads(res[0]["EXEC_INTO"])["rai_transaction_id"]
                        rejected_rows = json.loads(res[0]["EXEC_INTO"]).get("rejected_rows", [])
                        rejected_rows_count = json.loads(res[0]["EXEC_INTO"]).get("rejected_rows_count", 0)
                    debugging.event("transaction_created", txn_span, txn_id=txn_id)
                    debugging.time("exec_format", time.perf_counter() - start, DataFrame())

                with debugging.span("temp_table_swap", txn_id=txn_id):
                    out_sample = self._exec(f"select * from {APP_NAME}.results.{output_table} limit 1;")
                    if out_sample:
                        keys = set([k.lower() for k in out_sample[0].as_dict().keys()])
                        col_names_map = {}
                        for ix, name in enumerate(cols):
                            col_key = f"col{ix:03}"
                            if col_key in keys:
                                col_names_map[col_key] = IdentityParser(name).identity
                            else:
                                col_names_map[col_key] = name

                        names = ", ".join([
                            f"{col_key} as {alias}" if col_key in keys else f"NULL as {alias}"
                            for col_key, alias in col_names_map.items()
                        ])
                        self._exec(f"CREATE TEMPORARY TABLE {APP_NAME}.results.{temp_table} AS SELECT {names} FROM {APP_NAME}.results.{output_table};")
                        self._exec(f"call {APP_NAME}.api.drop_result_table(?)", [output_table])
                        temp = cast(snowflake.snowpark.DataFrame, self._exec(f"select * from {APP_NAME}.results.{temp_table}", raw=True))
                        if rejected_rows:
                            debugging.warn(RowsDroppedFromTargetTableWarning(rejected_rows, rejected_rows_count, col_names_map))
            except Exception as e:
                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
            if txn_id:
                artifact_info = self._list_exec_async_artifacts(txn_id)
                with debugging.span("fetch"):
                    artifacts = self._download_results(artifact_info, txn_id, "ABORTED")
            return (temp, artifacts)

    #--------------------------------------------------
    # Custom model types
    #--------------------------------------------------

    def _get_ns(self, model:dsl.Graph):
        if model not in self._ns_cache:
            self._ns_cache[model] = _Snowflake(model)
        return self._ns_cache[model]

    def to_model_type(self, model:dsl.Graph, name: str, source:str):
        parser = IdentityParser(source)
        if not parser.is_complete:
            raise SnowflakeInvalidSource(Errors.call_source(), source)
        ns = self._get_ns(model)
        # skip the last item in the list (the full identifier)
        for part in parser.to_list()[:-1]:
            ns = ns._safe_get(part)
        assert parser.identity, f"Error parsing source in to_model_type: {source}"
        self.sources.add(parser.identity)
        return ns

    #--------------------------------------------------
    # Source Management
    #--------------------------------------------------

    def _check_source_updates(self, sources: Iterable[str]):
        if not sources:
            return {}
        app_name = self.get_app_name()

        source_types = dict[str, SourceInfo]()
        partitioned_sources: dict[str, dict[str, list[dict[str, str]]]] = defaultdict(
            lambda: defaultdict(list)
        )
        fqn_to_parts: dict[str, tuple[str, str, str]] = {}

        for source in sources:
            parser = IdentityParser(source, True)
            parsed = parser.to_list()
            assert len(parsed) == 4, f"Invalid source: {source}"
            db, schema, entity, identity = parsed
            assert db and schema and entity and identity, f"Invalid source: {source}"
            source_types[identity] = cast(
                SourceInfo,
                {
                    "type": None,
                    "state": "",
                    "columns_hash": None,
                    "table_created_at": None,
                    "stream_created_at": None,
                    "last_ddl": None,
                },
            )
            partitioned_sources[db][schema].append({"entity": entity, "identity": identity})
            fqn_to_parts[identity] = (db, schema, entity)

        if not partitioned_sources:
            return source_types

        state_queries: list[str] = []
        for db, schemas in partitioned_sources.items():
            select_rows: list[str] = []
            for schema, tables in schemas.items():
                for table_info in tables:
                    select_rows.append(
                        "SELECT "
                        f"{IdentityParser.to_sql_value(db)} AS catalog_name, "
                        f"{IdentityParser.to_sql_value(schema)} AS schema_name, "
                        f"{IdentityParser.to_sql_value(table_info['entity'])} AS table_name"
                    )

            if not select_rows:
                continue

            target_entities_clause = "\n                        UNION ALL\n                        ".join(select_rows)
            # Main query:
            #   1. Enumerate the target tables via target_entities.
            #   2. Pull their metadata (last_altered, type) from INFORMATION_SCHEMA.TABLES.
            #   3. Look up the most recent stream activity for those FQNs only.
            #   4. Capture creation timestamps and use last_ddl vs created_at to classify each target,
            #      so we mark tables as stale when they were recreated even if column hashes still match.
            state_queries.append(
                f"""WITH target_entities AS (
                        {target_entities_clause}
                    ),
                    table_info AS (
                        SELECT
                            {app_name}.api.normalize_fq_ids(
                                ARRAY_CONSTRUCT(
                                    CASE
                                        WHEN t.table_catalog = UPPER(t.table_catalog) THEN t.table_catalog
                                        ELSE '"' || t.table_catalog || '"'
                                    END || '.' ||
                                    CASE
                                        WHEN t.table_schema = UPPER(t.table_schema) THEN t.table_schema
                                        ELSE '"' || t.table_schema || '"'
                                    END || '.' ||
                                    CASE
                                        WHEN t.table_name = UPPER(t.table_name) THEN t.table_name
                                        ELSE '"' || t.table_name || '"'
                                    END
                                )
                            )[0]:identifier::string AS fqn,
                            CONVERT_TIMEZONE('UTC', t.last_altered) AS last_ddl,
                            CONVERT_TIMEZONE('UTC', t.created) AS table_created_at,
                            t.table_type AS kind
                        FROM {db}.INFORMATION_SCHEMA.tables t
                        JOIN target_entities te
                            ON t.table_catalog = te.catalog_name
                            AND t.table_schema = te.schema_name
                            AND t.table_name = te.table_name
                    ),
                    stream_activity AS (
                        SELECT
                            sa.fqn,
                            MAX(sa.created_at) AS created_at
                        FROM (
                            SELECT
                                {app_name}.api.normalize_fq_ids(ARRAY_CONSTRUCT(fq_object_name))[0]:identifier::string AS fqn,
                                created_at
                            FROM {app_name}.api.data_streams
                            WHERE rai_database = '{PYREL_ROOT_DB}'
                        ) sa
                        JOIN table_info ti
                            ON sa.fqn = ti.fqn
                        GROUP BY sa.fqn
                    )
                    SELECT
                        ti.fqn,
                        ti.kind,
                        ti.last_ddl,
                        ti.table_created_at,
                        sa.created_at AS stream_created_at,
                        IFF(
                            DATEDIFF(second, sa.created_at::timestamp, ti.last_ddl::timestamp) > 0,
                            'STALE',
                            'CURRENT'
                        ) AS state
                    FROM table_info ti
                    LEFT JOIN stream_activity sa
                        ON sa.fqn = ti.fqn
                """
            )

        stale_fqns: list[str] = []
        for state_query in state_queries:
            for row in self._exec(state_query):
                row_dict = row.as_dict() if hasattr(row, "as_dict") else dict(row)
                row_fqn = row_dict["FQN"]
                parser = IdentityParser(row_fqn, True)
                fqn = parser.identity
                assert fqn, f"Error parsing returned FQN: {row_fqn}"

                source_types[fqn]["type"] = (
                    "TABLE" if row_dict["KIND"] == "BASE TABLE" else row_dict["KIND"]
                )
                source_types[fqn]["state"] = row_dict["STATE"]
                source_types[fqn]["last_ddl"] = normalize_datetime(row_dict.get("LAST_DDL"))
                source_types[fqn]["table_created_at"] = normalize_datetime(row_dict.get("TABLE_CREATED_AT"))
                source_types[fqn]["stream_created_at"] = normalize_datetime(row_dict.get("STREAM_CREATED_AT"))
                if row_dict["STATE"] == "STALE":
                    stale_fqns.append(fqn)

        if not stale_fqns:
            return source_types

        # We batch stale tables by database/schema so each Snowflake query can hash
        # multiple objects at once instead of issuing one statement per table.
        stale_partitioned: dict[str, dict[str, list[dict[str, str]]]] = defaultdict(
            lambda: defaultdict(list)
        )
        for fqn in stale_fqns:
            db, schema, table = fqn_to_parts[fqn]
            stale_partitioned[db][schema].append({"table": table, "identity": fqn})

        # Build one hash query per database, grouping schemas/tables inside so we submit
        # at most a handful of set-based statements to Snowflake.
        for db, schemas in stale_partitioned.items():
            column_select_rows: list[str] = []
            for schema, tables in schemas.items():
                for table_info in tables:
                    # Build the literal rows for this db/schema so we can join back
                    # against INFORMATION_SCHEMA.COLUMNS in a single statement.
                    column_select_rows.append(
                        "SELECT "
                        f"{IdentityParser.to_sql_value(db)} AS catalog_name, "
                        f"{IdentityParser.to_sql_value(schema)} AS schema_name, "
                        f"{IdentityParser.to_sql_value(table_info['table'])} AS table_name"
                    )

            if not column_select_rows:
                continue

            target_entities_clause = "\n                        UNION ALL\n                        ".join(column_select_rows)
            # Main query: compute deterministic column hashes for every stale table
            # in this database/schema batch so we can compare schemas without a round trip per table.
            column_query = f"""WITH target_entities AS (
                        {target_entities_clause}
                    ),
                    column_info AS (
                        SELECT
                            {app_name}.api.normalize_fq_ids(
                                ARRAY_CONSTRUCT(
                                    CASE
                                        WHEN c.table_catalog = UPPER(c.table_catalog) THEN c.table_catalog
                                        ELSE '"' || c.table_catalog || '"'
                                    END || '.' ||
                                    CASE
                                        WHEN c.table_schema = UPPER(c.table_schema) THEN c.table_schema
                                        ELSE '"' || c.table_schema || '"'
                                    END || '.' ||
                                    CASE
                                        WHEN c.table_name = UPPER(c.table_name) THEN c.table_name
                                        ELSE '"' || c.table_name || '"'
                                    END
                                )
                            )[0]:identifier::string AS fqn,
                            c.column_name,
                            CASE
                                WHEN c.numeric_precision IS NOT NULL AND c.numeric_scale IS NOT NULL
                                    THEN c.data_type || '(' || c.numeric_precision || ',' || c.numeric_scale || ')'
                                WHEN c.datetime_precision IS NOT NULL
                                    THEN c.data_type || '(0,' || c.datetime_precision || ')'
                                WHEN c.character_maximum_length IS NOT NULL
                                    THEN c.data_type || '(' || c.character_maximum_length || ')'
                                ELSE c.data_type
                            END AS type_signature,
                            IFF(c.is_nullable = 'YES', 'YES', 'NO') AS nullable_flag
                        FROM {db}.INFORMATION_SCHEMA.COLUMNS c
                        JOIN target_entities te
                            ON c.table_catalog = te.catalog_name
                            AND c.table_schema = te.schema_name
                            AND c.table_name = te.table_name
                    )
                    SELECT
                        fqn,
                        HEX_ENCODE(
                            HASH_AGG(
                                HASH(
                                    column_name,
                                    type_signature,
                                    nullable_flag
                                )
                            )
                        ) AS columns_hash
                    FROM column_info
                    GROUP BY fqn
                """

            for row in self._exec(column_query):
                row_fqn = row["FQN"]
                parser = IdentityParser(row_fqn, True)
                fqn = parser.identity
                assert fqn, f"Error parsing returned FQN: {row_fqn}"
                source_types[fqn]["columns_hash"] = row["COLUMNS_HASH"]

        return source_types

    def _get_source_references(self, source_info: dict[str, SourceInfo]):
        app_name = self.get_app_name()
        missing_sources = []
        invalid_sources = {}
        source_references = []
        for source, info in source_info.items():
            source_type = info.get("type")
            if source_type is None:
                missing_sources.append(source)
            elif source_type not in ("TABLE", "VIEW"):
                invalid_sources[source] = source_type
            else:
                source_references.append(f"{app_name}.api.object_reference('{source_type}', '{source}')")

        if missing_sources:
            current_role = self.get_sf_session().get_current_role()
            if current_role is None:
                current_role = self.config.get("role", None)
            debugging.warn(UnknownSourceWarning(missing_sources, current_role))

        if invalid_sources:
            debugging.warn(InvalidSourceTypeWarning(invalid_sources))

        self.source_references = source_references
        return source_references

    #--------------------------------------------------
    # Transactions
    #--------------------------------------------------

    def get_transaction(self, transaction_id):
        results = self._exec(
            f"CALL {APP_NAME}.api.get_transaction(?);", [transaction_id])
        if not results:
            return None

        results = txn_list_to_dicts(results)

        txn = {field: results[0][field] for field in GET_TXN_SQL_FIELDS}

        state = txn.get("state")
        created_on = txn.get("created_on")
        finished_at = txn.get("finished_at")
        if created_on:
            # Transaction is still running
            if state not in TERMINAL_TXN_STATES:
                tz_info = created_on.tzinfo
                txn['duration'] = datetime.now(tz_info) - created_on
            # Transaction is terminal
            elif finished_at:
                txn['duration'] = finished_at - created_on
            # Transaction is still running and we have no state or finished_at
            else:
                txn['duration'] = timedelta(0)
        return txn

    def list_transactions(self, **kwargs):
        id = kwargs.get("id", None)
        state = kwargs.get("state", None)
        engine = kwargs.get("engine", None)
        limit = kwargs.get("limit", 100)
        all_users = kwargs.get("all_users", False)
        created_by = kwargs.get("created_by", None)
        only_active = kwargs.get("only_active", False)
        where_clause_arr = []

        if id:
            where_clause_arr.append(f"id = '{id}'")
        if state:
            where_clause_arr.append(f"state = '{state.upper()}'")
        if engine:
            where_clause_arr.append(f"LOWER(engine_name) = '{engine.lower()}'")
        else:
            if only_active:
                where_clause_arr.append("state in ('CREATED', 'RUNNING', 'PENDING')")
        if not all_users and created_by is not None:
            where_clause_arr.append(f"LOWER(created_by) = '{created_by.lower()}'")

        if len(where_clause_arr):
            where_clause = f'WHERE {" AND ".join(where_clause_arr)}'
        else:
            where_clause = ""

        sql_fields = ", ".join(LIST_TXN_SQL_FIELDS)
        query = f"SELECT {sql_fields} from {APP_NAME}.api.transactions {where_clause} ORDER BY created_on DESC LIMIT ?"
        results = self._exec(query, [limit])
        if not results:
            return []
        return txn_list_to_dicts(results)

    def cancel_transaction(self, transaction_id):
        self._exec(f"CALL {APP_NAME}.api.cancel_own_transaction(?);", [transaction_id])
        if transaction_id in self._pending_transactions:
            self._pending_transactions.remove(transaction_id)

    def cancel_pending_transactions(self):
        for txn_id in self._pending_transactions:
            self.cancel_transaction(txn_id)

    def get_transaction_events(self, transaction_id: str, continuation_token:str=''):
        results = self._exec(
            f"SELECT {APP_NAME}.api.get_own_transaction_events(?, ?);",
            [transaction_id, continuation_token],
        )
        if not results:
            return {
                "events": [],
                "continuation_token": None
            }
        row = results[0][0]
        return json.loads(row)

    #--------------------------------------------------
    # Snowflake specific
    #--------------------------------------------------

    def get_version(self):
        results = self._exec(f"SELECT {APP_NAME}.app.get_release()")
        if not results:
            return None
        return results[0][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 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:
                pk_query = f"SHOW PRIMARY KEYS IN SCHEMA {parser.identity};"
                pks = self._exec(pk_query)
                span["sql"] = pk_query

            with debugging.span("foreign_keys") as span:
                fk_query = f"SHOW IMPORTED KEYS IN SCHEMA {parser.identity};"
                fks = self._exec(fk_query)
                span["sql"] = fk_query

            # 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_str = ", ".join([f"'{IdentityParser(t).db}'" for t in tables])
                query = textwrap.dedent(f"""
                    begin
                        SHOW COLUMNS IN SCHEMA {parser.identity};
                        let r resultset := (
                            SELECT
                                CASE
                                    WHEN "table_name" = UPPER("table_name") THEN "table_name"
                                ELSE '"' || "table_name" || '"'
                                END as "table_name",
                                "column_name",
                                "data_type",
                                CASE
                                    WHEN ARRAY_CONTAINS(PARSE_JSON("data_type"):"type", {app_name}.app.get_supported_column_types()) THEN TRUE
                                    ELSE FALSE
                                END as "supported_type"
                            FROM table(result_scan(-1)) as t
                            WHERE "table_name" in ({tables_str})
                        );
                        return table(r);
                    end;
                """)
                span["sql"] = query
                columns = self._exec(query)

            results = defaultdict(lambda: {"pks": [], "fks": {}, "columns": {}, "invalid_columns": {}})
            if pks:
                for row in pks:
                    results[row[3]]["pks"].append(row[4])  # type: ignore
            if fks:
                for row in fks:
                    results[row[7]]["fks"][row[8]] = row[3]
            if columns:
                # It seems that a SF parameter (QUOTED_IDENTIFIERS_IGNORE_CASE) can control
                # whether snowflake will ignore case on `row.data_type`,
                # so we have to use column indexes instead :(
                for row in columns:
                    table_name = row[0]
                    column_name = row[1]
                    data_type = row[2]
                    supported_type = row[3]
                    # Filter out unsupported types
                    if supported_type:
                        results[table_name]["columns"][column_name] = data_type
                    else:
                        results[table_name]["invalid_columns"][column_name] = data_type
        return results

    def get_cloud_provider(self) -> str:
        """
        Detect whether this is Snowflake on Azure, or AWS using Snowflake's CURRENT_REGION().
        Returns 'azure' or 'aws'.
        """
        if self._session:
            try:
                # Query Snowflake's current region using the built-in function
                result = self._session.sql("SELECT CURRENT_REGION()").collect()
                if result:
                    region_info = result[0][0]
                    # Check if the region string contains the cloud provider name
                    if isinstance(region_info, str):
                        region_str = region_info.lower()
                        # Check for cloud providers in the region string
                        if 'azure' in region_str:
                            return 'azure'
                        else:
                            return 'aws'
            except Exception:
                pass

        # Fallback to AWS as default if detection fails
        return 'aws'

#--------------------------------------------------
# Snowflake Wrapper
#--------------------------------------------------

class PrimaryKey:
    pass

class _Snowflake:
    def __init__(self, model, auto_import=False):
        self._model = model
        self._auto_import = auto_import
        if not isinstance(model._client.resources, Resources):
            raise ValueError("Snowflake model must be used with a snowflake config")
        self._dbs = {}
        imports = model._client.resources.list_imports(model=model.name)
        self._import_structure(imports)

    def _import_structure(self, imports: list[Import]):
        tree = self._dbs
        # pre-create existing imports
        schemas = set()
        for item in imports:
            parser = IdentityParser(item["name"])
            database_name, schema_name, table_name = parser.to_list()[:-1]
            database = getattr(self, database_name)
            schema = getattr(database, schema_name)
            schemas.add(schema)
            schema._add(table_name, is_imported=True)
        return tree

    def _safe_get(self, name:str) -> 'SnowflakeDB':
        name = name
        if name in self._dbs:
            return self._dbs[name]
        self._dbs[name] = SnowflakeDB(self, name)
        return self._dbs[name]

    def __getattr__(self, name: str) -> 'SnowflakeDB':
        return self._safe_get(name)


class Snowflake(_Snowflake):
    def __init__(self, model: dsl.Graph, auto_import=False):
        if model._config.get_bool("use_graph_index", USE_GRAPH_INDEX):
            raise SnowflakeProxySourceError()
        else:
            debugging.warn(SnowflakeProxyAPIDeprecationWarning())

        super().__init__(model, auto_import)

class SnowflakeDB:
    def __init__(self, parent, name):
        self._name = name
        self._parent = parent
        self._model = parent._model
        self._schemas = {}

    def _safe_get(self, name: str) -> 'SnowflakeSchema':
        name = name
        if name in self._schemas:
            return self._schemas[name]
        self._schemas[name] = SnowflakeSchema(self, name)
        return self._schemas[name]

    def __getattr__(self, name: str) -> 'SnowflakeSchema':
        return self._safe_get(name)

class SnowflakeSchema:
    def __init__(self, parent, name):
        self._name = name
        self._parent = parent
        self._model = parent._model
        self._tables = {}
        self._imported = set()
        self._table_info = defaultdict(lambda: {"pks": [], "fks": {}, "columns": {}, "invalid_columns": {}})
        self._dirty = True

    def _fetch_info(self):
        if not self._dirty:
            return
        self._table_info = self._model._client.resources.schema_info(self._parent._name, self._name, list(self._tables.keys()))

        check_column_types = self._model._config.get("check_column_types", True)

        if check_column_types:
            self._check_and_confirm_invalid_columns()

        self._dirty = False

    def _check_and_confirm_invalid_columns(self):
        """Check for invalid columns across the schema's tables."""
        tables_with_invalid_columns = {}
        for table_name, table_info in self._table_info.items():
            if table_info.get("invalid_columns"):
                tables_with_invalid_columns[table_name] = table_info["invalid_columns"]

        if tables_with_invalid_columns:
            from v0.relationalai.errors import UnsupportedColumnTypesWarning
            UnsupportedColumnTypesWarning(tables_with_invalid_columns)

    def _add(self, name, is_imported=False):
        if name in self._tables:
            return self._tables[name]
        self._dirty = True
        if is_imported:
            self._imported.add(name)
        else:
            self._tables[name] = SnowflakeTable(self, name)
        return self._tables.get(name)

    def _safe_get(self, name: str) -> 'SnowflakeTable | None':
        table = self._add(name)
        return table

    def __getattr__(self, name: str) -> 'SnowflakeTable | None':
        return self._safe_get(name)


class SnowflakeTable(dsl.Type):
    def __init__(self, parent, name):
        super().__init__(parent._model, f"sf_{name}")
        # hack to make this work for pathfinder
        self._type.parents.append(m.Builtins.PQFilterAnnotation)
        self._name = name
        self._model = parent._model
        self._parent = parent
        self._aliases = {}
        self._finalzed = False
        self._source = runtime_env.get_source()
        relation_name = to_fqn_relation_name(self.fqname())
        self._model.install_raw(f"declare {relation_name}")

    def __call__(self, *args, **kwargs):
        self._lazy_init()
        return super().__call__(*args, **kwargs)

    def add(self, *args, **kwargs):
        self._lazy_init()
        return super().add(*args, **kwargs)

    def extend(self, *args, **kwargs):
        self._lazy_init()
        return super().extend(*args, **kwargs)

    def known_properties(self):
        self._lazy_init()
        return super().known_properties()

    def _lazy_init(self):
        if self._finalzed:
            return

        parent = self._parent
        name = self._name
        use_graph_index = self._model._config.get("use_graph_index", USE_GRAPH_INDEX)

        if not use_graph_index and name not in parent._imported:
            if self._parent._parent._parent._auto_import:
                with Spinner(f"Creating stream for {self.fqname()}", f"Stream for {self.fqname()} created successfully"):
                    db_name = parent._parent._name
                    schema_name = parent._name
                    self._model._client.resources.create_import_stream(ImportSourceTable(db_name, schema_name, name), self._model.name)
                print("")
                parent._imported.add(name)
            else:
                imports = self._model._client.resources.list_imports(model=self._model.name)
                for item in imports:
                    cur_name = item["name"].lower().split(".")[-1]
                    parent._imported.add(cur_name)
            if name not in parent._imported:
                exception = SnowflakeImportMissingException(runtime_env.get_source(), self.fqname(), self._model.name)
                raise exception from None

        parent._fetch_info()
        self._finalize()

    def _finalize(self):
        if self._finalzed:
            return

        self._finalzed = True
        self._schema = self._parent._table_info[self._name]

        # Set the relation name to the sanitized version of the fully qualified name
        relation_name = to_fqn_relation_name(self.fqname())

        model:dsl.Graph = self._model
        edb = getattr(std.rel, relation_name)
        edb._rel.parents.append(m.Builtins.EDB)
        id_rel = getattr(std.rel, f"{relation_name}_pyrel_id")

        with model.rule(globalize=True, source=self._source):
            id, val = dsl.create_vars(2)
            edb(dsl.Symbol("METADATA$ROW_ID"), id, val)
            std.rel.SHA1(id)
            id_rel.add(id)

        with model.rule(dynamic=True, globalize=True, source=self._source):
            prop, id, val = dsl.create_vars(3)
            id_rel(id)
            std.rel.SHA1(id)
            self.add(snowflake_id=id)

        for prop, prop_type in self._schema["columns"].items():
            _prop = prop
            if _prop.startswith("_"):
                _prop = "col" + prop

            prop_ident = sanitize_identifier(_prop.lower())

            with model.rule(dynamic=True, globalize=True, source=self._source):
                id, val = dsl.create_vars(2)
                edb(dsl.Symbol(prop), id, val)
                std.rel.SHA1(id)
                _prop = getattr(self, prop_ident)
                if not _prop:
                    raise ValueError(f"Property {_prop} couldn't be accessed on {self.fqname()}")
                if _prop.is_multi_valued:
                    inst = self(snowflake_id=id)
                    getattr(inst, prop_ident).add(val)
                else:
                    self(snowflake_id=id).set(**{prop_ident: val})

        # Because we're bypassing a bunch of the normal Type.add machinery here,
        # we need to manually account for the case where people are using value types.
        def wrapped(x):
            if not model._config.get("compiler.use_value_types", False):
                return x
            other_id = dsl.create_var()
            model._action(dsl.build.construct(self._type, [x, other_id]))
            return other_id

        # new UInt128 schema mapping rules
        with model.rule(dynamic=True, globalize=True, source=self._source):
            id = dsl.create_var()
            # This will generate an arity mismatch warning when used with the old SHA-1 Data Streams.
            # Ideally we have the `@no_diagnostics(:ARITY_MISMATCH)` attribute on the relation using
            # the METADATA$KEY column but that ended up being a more involved change then expected
            # for avoiding a non-blocking warning
            edb(dsl.Symbol("METADATA$KEY"), id)
            std.rel.UInt128(id)
            self.add(wrapped(id), snowflake_id=id)

        for prop, prop_type in self._schema["columns"].items():
            _prop = prop
            if _prop.startswith("_"):
                _prop = "col" + prop

            prop_ident = sanitize_identifier(_prop.lower())
            with model.rule(dynamic=True, globalize=True, source=self._source):
                id, val = dsl.create_vars(2)
                edb(dsl.Symbol(prop), id, val)
                std.rel.UInt128(id)
                _prop = getattr(self, prop_ident)
                if not _prop:
                    raise ValueError(f"Property {_prop} couldn't be accessed on {self.fqname()}")
                if _prop.is_multi_valued:
                    inst = self(id)
                    getattr(inst, prop_ident).add(val)
                else:
                    model._check_property(_prop._prop)
                    raw_relation = getattr(std.rel, prop_ident)
                    dsl.tag(raw_relation, dsl.Builtins.FunctionAnnotation)
                    raw_relation.add(wrapped(id), val)

    def namespace(self):
        return f"{self._parent._parent._name}.{self._parent._name}"

    def fqname(self):
        return f"{self.namespace()}.{self._name}"

    def describe(self, **kwargs):
        model = self._model
        for k, v in kwargs.items():
            if v is PrimaryKey:
                self._schema["pks"] = [k]
            elif isinstance(v, tuple):
                (table, name) = v
                if isinstance(table, SnowflakeTable):
                    fk_table = table
                    pk = fk_table._schema["pks"]
                    with model.rule():
                        inst = fk_table()
                        me = self()
                        getattr(inst, pk[0]) == getattr(me, k)
                        if getattr(self, name).is_multi_valued:
                            getattr(me, name).add(inst)
                        else:
                            me.set(**{name: inst})
                else:
                    raise ValueError(f"Invalid foreign key {v}")
            else:
                raise ValueError(f"Invalid column {k}={v}")
        return self

class Provider(ProviderBase):
    def __init__(
        self,
        profile: str | None = None,
        config: Config | None = None,
        resources: Resources | None = None,
        generation: Generation | None = None,
    ):
        if resources:
            self.resources = resources
        else:
            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)

    def create_streams(self, sources:List[str], model:str, force=False):
        if not self.resources.get_graph(model):
            self.resources.create_graph(model)
        def parse_source(raw:str):
            parser = IdentityParser(raw)
            assert parser.is_complete, "Snowflake table imports must be in `database.schema.table` format"
            return ImportSourceTable(*parser.to_list())
        for source in sources:
            source_table = parse_source(source)
            try:
                with Spinner(f"Creating stream for {source_table.name}", f"Stream for {source_table.name} created successfully"):
                    if force:
                        self.resources.delete_import(source_table.name, model, True)
                    self.resources.create_import_stream(source_table, model)
            except Exception as e:
                if "stream already exists" in f"{e}":
                    raise Exception(f"\n\nStream'{source_table.name.upper()}' already exists.")
                elif "engine not found" in f"{e}":
                    raise Exception("\n\nNo engines found in a READY state. Please use `engines:create` to create an engine that will be used to initialize the target relation.")
                else:
                    raise e
        with Spinner("Waiting for imports to complete", "Imports complete"):
            self.resources.poll_imports(sources, model)

    def delete_stream(self, stream_id: str, model: str):
        return self.resources.delete_import(stream_id, model)

    def sql(self, query:str, params:List[Any]=[], format:Literal["list", "pandas", "polars", "lazy"]="list"):
        # note: default format cannot be pandas because .to_pandas() only works on SELECT queries
        result = self.resources._exec(query, params, raw=True, help=False)
        if format == "lazy":
            return cast(snowflake.snowpark.DataFrame, result)
        elif format == "list":
            return cast(list, result.collect())
        elif format == "pandas":
            import pandas as pd
            try:
                # use to_pandas for SELECT queries
                return cast(pd.DataFrame, result.to_pandas())
            except Exception:
                # handle non-SELECT queries like SHOW
                return pd.DataFrame(result.collect())
        elif format == "polars":
            import polars as pl # type: ignore
            return pl.DataFrame(
                [row.as_dict() for row in result.collect()],
                orient="row",
                strict=False,
                infer_schema_length=None
            )
        else:
            raise ValueError(f"Invalid format {format}. Should be one of 'list', 'pandas', 'polars', 'lazy'")

    def activate(self):
        with Spinner("Activating RelationalAI app...", "RelationalAI app activated"):
            self.sql("CALL RELATIONALAI.APP.ACTIVATE();")

    def deactivate(self):
        with Spinner("Deactivating RelationalAI app...", "RelationalAI app deactivated"):
            self.sql("CALL RELATIONALAI.APP.DEACTIVATE();")

    def drop_service(self):
        warnings.warn(
            "The drop_service method has been deprecated in favor of deactivate",
            DeprecationWarning,
            stacklevel=2,
        )
        self.deactivate()

    def resume_service(self):
        warnings.warn(
            "The resume_service method has been deprecated in favor of activate",
            DeprecationWarning,
            stacklevel=2,
        )
        self.activate()


#--------------------------------------------------
# SnowflakeClient
#--------------------------------------------------
class SnowflakeClient(Client):
    def create_database(self, isolated=True, nowait_durable=True, headers: Dict | None = None):
        from v0.relationalai.tools.cli_helpers import validate_engine_name

        assert isinstance(self.resources, Resources)

        if self.last_database_version == len(self.resources.sources):
            return

        model = self._source_database
        app_name = self.resources.get_app_name()
        engine_name = self.resources.get_default_engine_name()
        engine_size = self.resources.config.get_default_engine_size()

        # Validate engine name
        is_name_valid, _ = validate_engine_name(engine_name)
        if not is_name_valid:
            raise EngineNameValidationException(engine_name)

        # Validate engine size
        valid_sizes = self.resources.get_engine_sizes()
        if not isinstance(engine_size, str) or engine_size not in valid_sizes:
            raise InvalidEngineSizeError(str(engine_size), valid_sizes)

        program_span_id = debugging.get_program_span_id()

        query_attrs_dict = json.loads(headers.get("X-Query-Attributes", "{}")) if headers else {}
        with debugging.span("poll_use_index", sources=self.resources.sources, model=model, engine=engine_name, **query_attrs_dict):
            self.maybe_poll_use_index(
                app_name=app_name,
                sources=self.resources.sources,
                model=model,
                engine_name=engine_name,
                engine_size=engine_size,
                program_span_id=program_span_id,
                headers=headers
            )

        self.last_database_version = len(self.resources.sources)
        self._manage_packages()

        if isolated and not self.keep_model:
            atexit.register(self.delete_database)

    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_use_index if there are sources to process."""
        assert isinstance(self.resources, Resources)
        return self.resources.maybe_poll_use_index(
            app_name=app_name,
            sources=sources,
            model=model,
            engine_name=engine_name,
            engine_size=engine_size,
            program_span_id=program_span_id,
            headers=headers
        )


#--------------------------------------------------
# Graph
#--------------------------------------------------

def Graph(
    name,
    *,
    profile: str | None = None,
    config: Config,
    dry_run: bool = False,
    isolated: bool = True,
    connection: Session | None = None,
    keep_model: bool = False,
    nowait_durable: bool = True,
    format: str = "default",
):
    from .resources_factory import create_resources_instance
    from .use_index_resources import UseIndexResources

    use_graph_index = config.get("use_graph_index", USE_GRAPH_INDEX)
    use_monotype_operators = config.get("compiler.use_monotype_operators", False)

    # 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
    else:
        client_class = Client

    client = client_class(
        resources,
        rel.Compiler(config),
        name,
        config,
        dry_run=dry_run,
        isolated=isolated,
        keep_model=keep_model,
        nowait_durable=nowait_durable
    )
    base_rel = """
        @inline
        def make_identity(x..., z):
            rel_primitive_hash_tuple_uint128(x..., z)

        @inline
        def pyrel_default({F}, c, k..., v):
            F(k..., v) or (not F(k..., _) and v = c)

        @inline
        def pyrel_unwrap(x in UInt128, y): y = x

        @inline
        def pyrel_dates_period_days(x in Date, y in Date, z in Int):
            exists((u) | dates_period_days(x, y , u) and u = ::std::common::^Day[z])

        @inline
        def pyrel_datetimes_period_milliseconds(x in DateTime, y in DateTime, z in Int):
            exists((u) | datetimes_period_milliseconds(x, y , u) and u = ^Millisecond[z])

        @inline
        def pyrel_bool_filter(a, b, {F}, z): { z = if_then_else[F(a, b), boolean_true, boolean_false] }

        @inline
        def pyrel_strftime(v, fmt, tz in String, s in String):
            (Date(v) and s = format_date[v, fmt])
            or (DateTime(v) and s = format_datetime[v, fmt, tz])

        @inline
        def pyrel_regex_match_all(pattern, string in String, pos in Int, offset in Int, match in String):
            regex_match_all(pattern, string, offset, match) and offset >= pos

        @inline
        def pyrel_regex_match(pattern, string in String, pos in Int, offset in Int, match in String):
            pyrel_regex_match_all(pattern, string, pos, offset, match) and offset = pos

        @inline
        def pyrel_regex_search(pattern, string in String, pos in Int, offset in Int, match in String):
            enumerate(pyrel_regex_match_all[pattern, string, pos], 1, offset, match)

        @inline
        def pyrel_regex_sub(pattern, repl in String, string in String, result in String):
            string_replace_multiple(string, {(last[regex_match_all[pattern, string]], repl)}, result)

        @inline
        def pyrel_capture_group(regex in Pattern, string in String, pos in Int, index, match in String):
            (Integer(index) and capture_group_by_index(regex, string, pos, index, match)) or
            (String(index) and capture_group_by_name(regex, string, pos, index, match))

        declare __resource
        declare __compiled_patterns
    """
    if use_monotype_operators:
        base_rel += """

        // use monotyped operators
        from ::std::monotype import +, -, *, /, <, <=, >, >=
        """
    pyrel_base = dsl.build.raw_task(base_rel)
    debugging.set_source(pyrel_base)
    client.install("pyrel_base", pyrel_base)
    return dsl.Graph(client, name, format=format)
