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

relationalai/clients/resources/snowflake/use_index_poller.py

@@ -1,1011 +0,0 @@
-from __future__ import annotations
-
-from typing import Iterable, Dict, Optional, List, cast, TYPE_CHECKING
-import json
-import logging
-import uuid
-
-from .... import debugging
-from .cache_store import GraphIndexCache
-from .util import collect_error_messages
-from ...util import (
-    get_pyrel_version,
-    normalize_datetime,
-    poll_with_specified_overhead,
-)
-from ....errors import (
-    ERPNotRunningError,
-    EngineProvisioningFailed,
-    SnowflakeChangeTrackingNotEnabledException,
-    SnowflakeTableObjectsException,
-    SnowflakeTableObject,
-    SnowflakeRaiAppNotStarted,
-)
-from ....tools.cli_controls import (
-    DebuggingSpan,
-    create_progress,
-    TASK_CATEGORY_INDEXING,
-    TASK_CATEGORY_PROVISIONING,
-    TASK_CATEGORY_CHANGE_TRACKING,
-    TASK_CATEGORY_CACHE,
-    TASK_CATEGORY_RELATIONS,
-    TASK_CATEGORY_STATUS,
-    TASK_CATEGORY_VALIDATION,
-)
-from ....tools.constants import WAIT_FOR_STREAM_SYNC, Generation
-
-# Set up logger for this module
-logger = logging.getLogger(__name__)
-
-
-try:
-    from rich.console import Console
-    from rich.table import Table
-except ImportError:
-    Console = None
-    Table = None
-
-if TYPE_CHECKING:
-    from .snowflake import Resources
-    from .direct_access_resources import DirectAccessResources
-
-# Maximum number of items to show individual subtasks for
-# If more items than this, show a single summary subtask instead
-MAX_INDIVIDUAL_SUBTASKS = 5
-
-# Special engine name for CDC managed engine
-CDC_MANAGED_ENGINE = "CDC_MANAGED_ENGINE"
-
-# Maximum number of data source subtasks to show simultaneously
-# When one completes, the next one from the queue will be added
-MAX_DATA_SOURCE_SUBTASKS = 10
-
-# How often to check ERP status (every N iterations)
-# To limit performance overhead, we only check ERP status periodically
-ERP_CHECK_FREQUENCY = 15
-
-# Polling behavior constants
-POLL_OVERHEAD_RATE = 0.1  # Overhead rate for exponential backoff
-POLL_MAX_DELAY = 2.5  # Maximum delay between polls in seconds
-
-# SQL query template for getting stream column hashes
-# This query calculates a hash of column metadata (name, type, precision, scale, nullable)
-# to detect if source table schema has changed since stream was created
-STREAM_COLUMN_HASH_QUERY = """
-WITH stream_columns AS (
-    SELECT
-        fq_object_name,
-        HASH(
-            value:name::VARCHAR,
-            CASE
-                WHEN value:precision IS NOT NULL AND value:scale IS NOT NULL THEN CASE value:type::VARCHAR
-                    WHEN 'FIXED' THEN 'NUMBER'
-                    WHEN 'REAL' THEN 'FLOAT'
-                    WHEN 'TEXT' THEN 'TEXT'
-                    ELSE value:type::VARCHAR
-                END || '(' || value:precision || ',' || value:scale || ')'
-                WHEN value:precision IS NOT NULL AND value:scale IS NULL THEN CASE value:type::VARCHAR
-                    WHEN 'FIXED' THEN 'NUMBER'
-                    WHEN 'REAL' THEN 'FLOAT'
-                    WHEN 'TEXT' THEN 'TEXT'
-                    ELSE value:type::VARCHAR
-                END || '(0,' || value:precision || ')'
-                WHEN value:length IS NOT NULL THEN CASE value:type::VARCHAR
-                    WHEN 'FIXED' THEN 'NUMBER'
-                    WHEN 'REAL' THEN 'FLOAT'
-                    WHEN 'TEXT' THEN 'TEXT'
-                    ELSE value:type::VARCHAR
-                END || '(' || value:length || ')'
-                ELSE CASE value:type::VARCHAR
-                    WHEN 'FIXED' THEN 'NUMBER'
-                    WHEN 'REAL' THEN 'FLOAT'
-                    WHEN 'TEXT' THEN 'TEXT'
-                    ELSE value:type::VARCHAR
-                END
-            END,
-            IFF(value:nullable::BOOLEAN, 'YES', 'NO')
-        ) AS column_signature
-    FROM {app_name}.api.data_streams,
-        LATERAL FLATTEN(input => columns)
-    WHERE rai_database = '{rai_database}'
-        AND fq_object_name IN ({fqn_list})
-)
-SELECT
-    fq_object_name AS FQ_OBJECT_NAME,
-    HEX_ENCODE(HASH_AGG(column_signature)) AS STREAM_HASH
-FROM stream_columns
-GROUP BY fq_object_name;
-"""
-
-
-class UseIndexPoller:
-    """
-    Encapsulates the polling logic for `use_index` streams.
-    """
-
-    def _add_stream_subtask(self, progress, fq_name: str, status: str, batches_count: int) -> bool:
-        """Add a stream subtask if we haven't reached the limit.
-
-        Returns:
-            True if subtask was added, False if limit reached
-        """
-        if fq_name not in self.stream_task_ids and len(self.stream_task_ids) < MAX_DATA_SOURCE_SUBTASKS:
-            # Get the position in the stream order (should already be there)
-            if fq_name in self.stream_order:
-                stream_position = self.stream_order.index(fq_name) + 1
-            else:
-                # Fallback if not in order (shouldn't happen)
-                stream_position = 1
-
-            # Build initial message based on status and batch count
-            if status == "synced":
-                initial_message = f"{fq_name} already synced"
-            elif batches_count > 0:
-                # Show stream position (x/y) before batch count
-                initial_message = f"Syncing {fq_name} ({stream_position}/{self.total_streams}), batches: {batches_count}"
-            else:
-                initial_message = f"Syncing {fq_name} ({stream_position}/{self.total_streams})"
-
-            self.stream_task_ids[fq_name] = progress.add_sub_task(initial_message, task_id=fq_name, category=TASK_CATEGORY_INDEXING)
-
-            # Complete immediately if already synced (without recording completion time)
-            if status == "synced":
-                progress.complete_sub_task(fq_name, record_time=False)
-
-            return True
-        return False
-
-    def __init__(
-        self,
-        resource: "Resources",
-        app_name: str,
-        sources: Iterable[str],
-        model: str,
-        engine_name: str,
-        engine_size: Optional[str],
-        language: str = "rel",
-        program_span_id: Optional[str] = None,
-        headers: Optional[Dict] = None,
-        generation: Optional[Generation] = None,
-    ):
-        self.res = resource
-        self.app_name = app_name
-        self.sources = list(sources)
-        self.model = model
-        self.engine_name = engine_name
-        self.engine_size = engine_size or self.res.config.get_default_engine_size()
-        self.language = language
-        self.program_span_id = program_span_id
-        self.headers = headers or {}
-        self.counter = 1
-        self.check_ready_count = 0
-        self.tables_with_not_enabled_change_tracking: List = []
-        self.table_objects_with_other_errors: List = []
-        self.engine_errors: List = []
-        # Flag to only ensure the engine is created asynchronously the initial call
-        self.init_engine_async = True
-        # Initially, we assume that cdc is not checked,
-        # then on subsequent calls, if we get if cdc is enabled, if it is not, we will check it
-        # on every 5th iteration we reset the cdc status, so it will be checked again
-        self.should_check_cdc = True
-
-        self.wait_for_stream_sync = self.res.config.get(
-            "wait_for_stream_sync", WAIT_FOR_STREAM_SYNC
-        )
-        current_user = self.res.get_sf_session().get_current_user()
-        assert current_user is not None, "current_user must be set"
-        self.data_freshness = self.res.config.get_data_freshness_mins()
-        self.cache = GraphIndexCache(current_user, model, self.data_freshness, self.sources)
-        self.sources = self.cache.choose_sources()
-        # execution_id is allowed to group use_index call, which belongs to the same loop iteration
-        self.execution_id = str(uuid.uuid4())
-
-        self.pyrel_version = get_pyrel_version(generation)
-
-        self.source_info = self.res._check_source_updates(self.sources)
-
-        # Track subtask IDs for streams, engines, and relations across multiple poll iterations
-        self.stream_task_ids = {}
-        self.engine_task_ids = {}
-        self.relations_task_id = None
-        self._erp_check_task_id = None
-
-        # Track total number of streams and current stream position for (x/y) progress display
-        self.total_streams = 0
-        self.stream_position = 0
-        self.stream_order = []  # Track the order of streams as they appear in data
-
-        # Timing will be tracked by TaskProgress
-
-    def poll(self) -> None:
-        """
-        Standard stream-based polling for use_index.
-        """
-        # Read show_duration_summary config flag (defaults to True for backward compatibility)
-        show_duration_summary = bool(self.res.config.get("show_duration_summary", True))
-
-        with create_progress(
-            description="Initializing data index",
-            success_message="",  # We'll handle this in the context manager
-            leading_newline=True,
-            trailing_newline=True,
-            show_duration_summary=show_duration_summary,
-        ) as progress:
-            # Set process start time
-            progress.set_process_start_time()
-            progress.update_main_status("Validating data sources")
-            self._maybe_delete_stale(progress)
-
-            # Add cache usage subtask
-            self._add_cache_subtask(progress)
-
-            progress.update_main_status("Initializing data index")
-            self._poll_loop(progress)
-            self._post_check(progress)
-
-            # Set process end time (summary will be automatically printed by __exit__)
-            progress.set_process_end_time()
-
-    def _add_cache_subtask(self, progress) -> None:
-        """Add a subtask showing cache usage information only when cache is used."""
-        if self.cache.using_cache:
-            # Cache was used - show how many sources were cached
-            total_sources = len(self.cache.sources)
-            cached_sources = total_sources - len(self.sources)
-
-            # Get the timestamp when sources were cached
-            entry = self.cache._metadata.get("cachedIndices", {}).get(self.cache.key, {})
-            cached_timestamp = entry.get("last_use_index_update_on", "")
-
-            message = f"Using cached data for {cached_sources}/{total_sources} data streams"
-            # Format the message with timestamp
-            if cached_timestamp:
-                message += f" (cached at {cached_timestamp})"
-
-            progress.add_sub_task(message, task_id="cache_usage", category=TASK_CATEGORY_CACHE)
-            # Complete the subtask immediately since it's just informational
-            progress.complete_sub_task("cache_usage")
-
-    def _get_stream_column_hashes(self, sources: List[str], progress) -> Dict[str, str]:
-        """
-        Query data_streams to get current column hashes for the given sources.
-
-        Args:
-            sources: List of source FQNs to query
-            progress: TaskProgress instance for updating status on error
-
-        Returns:
-            Dict mapping FQN -> column hash
-
-        Raises:
-            ValueError: If the query fails (permissions, table doesn't exist, etc.)
-        """
-        from relationalai.clients.resources.snowflake import PYREL_ROOT_DB
-
-        # Build FQN list for SQL IN clause
-        fqn_list = ", ".join([f"'{source}'" for source in sources])
-
-        # Format query template with actual values
-        hash_query = STREAM_COLUMN_HASH_QUERY.format(
-            app_name=self.app_name,
-            rai_database=PYREL_ROOT_DB,
-            fqn_list=fqn_list
-        )
-
-        try:
-            hash_results = self.res._exec(hash_query)
-            return {row["FQ_OBJECT_NAME"]: row["STREAM_HASH"] for row in hash_results}
-
-        except Exception as e:
-            logger.error(f"Failed to query stream column hashes: {e}")
-            logger.error(f"  Query: {hash_query[:200]}...")
-            logger.error(f"  Sources: {sources}")
-            progress.update_main_status("❌ Failed to validate data stream metadata")
-            raise ValueError(
-                f"Failed to validate stream column hashes. This may indicate a permissions "
-                f"issue or missing data_streams table. Error: {e}"
-            ) from e
-
-    def _filter_truly_stale_sources(self, stale_sources: List[str], progress) -> List[str]:
-        """
-        Filter stale sources to only include those with mismatched column hashes.
-
-        Args:
-            stale_sources: List of source FQNs marked as stale
-            progress: TaskProgress instance for updating status on error
-
-        Returns:
-            List of truly stale sources that need to be deleted/recreated
-
-        A source is truly stale if any of the following apply:
-        - The stream doesn't exist (needs to be created)
-        - The source table was recreated after the stream (table creation timestamp is newer)
-        - The column hashes don't match (schema drift needs cleanup)
-        """
-        stream_hashes = self._get_stream_column_hashes(stale_sources, progress)
-
-        truly_stale = []
-        for source in stale_sources:
-            source_hash = self.source_info[source].get("columns_hash")
-            stream_hash = stream_hashes.get(source)
-            table_created_at_raw = self.source_info[source].get("table_created_at")
-            stream_created_at_raw = self.source_info[source].get("stream_created_at")
-
-            table_created_at = normalize_datetime(table_created_at_raw)
-            stream_created_at = normalize_datetime(stream_created_at_raw)
-
-            recreated_table = False
-            if table_created_at is not None and stream_created_at is not None:
-                # If the source table was recreated (new creation timestamp) but kept
-                # the same column definitions, we still need to recycle the stream so
-                # that Snowflake picks up the new table instance.
-                recreated_table = table_created_at > stream_created_at
-
-            # Log hash comparison for debugging
-            logger.debug(f"Source: {source}")
-            logger.debug(f"  Source table hash: {source_hash}")
-            logger.debug(f"  Stream hash: {stream_hash}")
-            logger.debug(f"  Match: {source_hash == stream_hash}")
-            if recreated_table:
-                logger.debug("  Table appears to have been recreated (table_created_at > stream_created_at)")
-                logger.debug(f"    table_created_at: {table_created_at}")
-                logger.debug(f"    stream_created_at: {stream_created_at}")
-
-            if stream_hash is None or source_hash != stream_hash or recreated_table:
-                logger.debug("  Action: DELETE (stale)")
-                truly_stale.append(source)
-            else:
-                logger.debug("  Action: KEEP (valid)")
-
-        logger.debug(f"Stale sources summary: {len(truly_stale)}/{len(stale_sources)} truly stale")
-
-        return truly_stale
-
-    def _add_deletion_subtasks(self, progress, sources: List[str]) -> None:
-        """Add progress subtasks for source deletion.
-
-        Args:
-            progress: TaskProgress instance
-            sources: List of source FQNs to be deleted
-        """
-        if len(sources) <= MAX_INDIVIDUAL_SUBTASKS:
-            for i, source in enumerate(sources):
-                progress.add_sub_task(
-                    f"Removing stale stream {source} ({i+1}/{len(sources)})",
-                    task_id=f"stale_source_{i}",
-                    category=TASK_CATEGORY_VALIDATION
-                )
-        else:
-            progress.add_sub_task(
-                f"Removing {len(sources)} stale data sources",
-                task_id="stale_sources_summary",
-                category=TASK_CATEGORY_VALIDATION
-            )
-
-    def _complete_deletion_subtasks(self, progress, sources: List[str], deleted_count: int) -> None:
-        """Complete progress subtasks for source deletion.
-
-        Args:
-            progress: TaskProgress instance
-            sources: List of source FQNs that were processed
-            deleted_count: Number of sources successfully deleted
-        """
-        if len(sources) <= MAX_INDIVIDUAL_SUBTASKS:
-            for i in range(len(sources)):
-                if f"stale_source_{i}" in progress._tasks:
-                    progress.complete_sub_task(f"stale_source_{i}")
-        else:
-            if "stale_sources_summary" in progress._tasks:
-                if deleted_count > 0:
-                    s = "s" if deleted_count > 1 else ""
-                    progress.update_sub_task(
-                        "stale_sources_summary",
-                        f"Removed {deleted_count} stale data source{s}"
-                    )
-                progress.complete_sub_task("stale_sources_summary")
-
-    def _maybe_delete_stale(self, progress) -> None:
-        """Check for and delete stale data streams that need recreation.
-
-        Args:
-            progress: TaskProgress instance for tracking deletion progress
-        """
-        with debugging.span("check_sources"):
-            stale_sources = [
-                source
-                for source, info in self.source_info.items()
-                if info.get("state") == "STALE"
-            ]
-
-        if not stale_sources:
-            return
-
-        with DebuggingSpan("validate_sources"):
-            try:
-                # Validate which sources truly need deletion by comparing column hashes
-                truly_stale = self._filter_truly_stale_sources(stale_sources, progress)
-
-                if not truly_stale:
-                    return
-
-                # Delete truly stale streams
-                from relationalai.clients.resources.snowflake import PYREL_ROOT_DB
-                query = f"CALL {self.app_name}.api.delete_data_streams({truly_stale}, '{PYREL_ROOT_DB}');"
-
-                self._add_deletion_subtasks(progress, truly_stale)
-
-                delete_response = self.res._exec(query)
-                delete_json_str = delete_response[0]["DELETE_DATA_STREAMS"].lower()
-                delete_data = json.loads(delete_json_str)
-                deleted_count = delete_data.get("deleted", 0)
-
-                self._complete_deletion_subtasks(progress, truly_stale, deleted_count)
-
-                # Check for errors
-                diff = len(truly_stale) - deleted_count
-                if diff > 0:
-                    errors = delete_data.get("errors", None)
-                    if errors:
-                        raise Exception(f"Error(s) deleting streams with modified sources: {errors}")
-
-            except Exception as e:
-                # Complete any remaining subtasks
-                self._complete_deletion_subtasks(progress, stale_sources, 0)
-                if "stale_sources_summary" in progress._tasks:
-                    progress.update_sub_task(
-                        "stale_sources_summary",
-                        f"❌ Failed to remove stale sources: {str(e)}"
-                    )
-
-                # Don't raise if streams don't exist - this is expected
-                messages = collect_error_messages(e)
-                if not any("data streams do not exist" in msg for msg in messages):
-                    raise e from None
-
-    def _poll_loop(self, progress) -> None:
-        """Main polling loop for use_index streams.
-
-        Args:
-            progress: TaskProgress instance for tracking polling progress
-        """
-        source_references = self.res._get_source_references(self.source_info)
-        sources_object_references_str = ", ".join(source_references)
-
-        def check_ready(progress) -> bool:
-            self.check_ready_count += 1
-
-            # To limit the performance overhead, we only check if ERP is running every N iterations
-            if self.check_ready_count % ERP_CHECK_FREQUENCY == 0:
-                with debugging.span("check_erp_status"):
-                    # Add subtask for ERP status check
-                    if self._erp_check_task_id is None:
-                        self._erp_check_task_id = progress.add_sub_task("Checking system status", task_id="erp_check", category=TASK_CATEGORY_STATUS)
-
-                    if not self.res.is_erp_running(self.app_name):
-                        progress.update_sub_task("erp_check", "❌ System status check failed")
-                        progress.complete_sub_task("erp_check")
-                        raise ERPNotRunningError
-                    else:
-                        progress.update_sub_task("erp_check", "System status check complete")
-                        progress.complete_sub_task("erp_check")
-
-            use_index_id = f"{self.model}_{self.execution_id}"
-
-            params = json.dumps({
-                "model": self.model,
-                "engine": self.engine_name,
-                "default_engine_size": self.engine_size, # engine_size
-                "user_agent": self.pyrel_version,
-                "use_index_id": use_index_id,
-                "pyrel_program_id": self.program_span_id,
-                "wait_for_stream_sync": self.wait_for_stream_sync,
-                "should_check_cdc": self.should_check_cdc,
-                "init_engine_async": self.init_engine_async,
-                "language": self.language,
-                "data_freshness_mins": self.data_freshness,
-            })
-
-            request_headers = debugging.add_current_propagation_headers(self.headers)
-
-            sql_string = f"CALL {self.app_name}.api.use_index([{sources_object_references_str}], PARSE_JSON(?), {request_headers});"
-
-            with debugging.span("wait", counter=self.counter, use_index_id=use_index_id) as span:
-                results = self.res._exec(sql_string, [params])
-
-                # Extract the JSON string from the `USE_INDEX` field
-                use_index_json_str = results[0]["USE_INDEX"]
-
-                # Parse the JSON string into a Python dictionary
-                try:
-                    use_index_data = json.loads(use_index_json_str)
-                except json.JSONDecodeError as e:
-                    logger.error(f"Invalid JSON from use_index API: {e}")
-                    logger.error(f"Raw response (first 500 chars): {use_index_json_str[:500]}")
-                    progress.update_main_status("❌ Received invalid response from server")
-                    raise ValueError(f"Invalid JSON response from use_index: {e}") from e
-
-                span.update(use_index_data)
-
-                # Log the full use_index_data for debugging if needed
-                logger.debug(f"use_index_data: {json.dumps(use_index_data, indent=4)}")
-
-                all_data = use_index_data.get("data", [])
-                ready = use_index_data.get("ready", False)
-                engines = use_index_data.get("engines", [])
-                errors = use_index_data.get("errors", [])
-                relations = use_index_data.get("relations", {})
-                cdc_enabled = use_index_data.get("cdcEnabled", False)
-                if self.check_ready_count % ERP_CHECK_FREQUENCY == 0 or not cdc_enabled:
-                    self.should_check_cdc = True
-                else:
-                    self.should_check_cdc = False
-
-                if engines and self.init_engine_async:
-                    self.init_engine_async = False
-
-                break_loop = False
-                has_stream_errors = False
-                has_general_errors = False
-
-                # Update main progress message
-                if ready:
-                    progress.update_main_status("Done")
-
-                # Handle streams data
-                if not ready and all_data:
-                    progress.update_main_status("Processing background tasks. This may take a while...")
-
-                    # Build complete stream order first (only on first iteration with data)
-                    if self.total_streams == 0:
-                        for data in all_data:
-                            if data is None:
-                                continue
-                            fq_name = data.get("fq_object_name", "Unknown")
-                            if fq_name not in self.stream_order:
-                                self.stream_order.append(fq_name)
-
-                        # Set total streams count based on complete order (only once)
-                        self.total_streams = len(self.stream_order)
-
-                    # Add new streams as subtasks if we haven't reached the limit
-                    for data in all_data:
-                        fq_name = data.get("fq_object_name", "Unknown")
-                        status = data.get("data_sync_status", "").lower() if data else ""
-                        batches_count = data.get("pending_batches_count", 0)
-
-                        # Only add if we haven't seen this stream and we're under the limit
-                        self._add_stream_subtask(progress, fq_name, status, batches_count)
-
-                        # Handle errors for existing streams
-                        if fq_name in self.stream_task_ids and data.get("errors", []):
-                            for error in data.get("errors", []):
-                                error_msg = f"{error.get('error')}, source: {error.get('source')}"
-                                # Some failures indicate the RAI app is not started/active; surface
-                                # them as a rich, actionable error instead of aggregating.
-                                self._raise_if_app_not_started(error_msg)
-                                self.table_objects_with_other_errors.append(
-                                    SnowflakeTableObject(error_msg, fq_name)
-                                )
-                            # Mark stream as failed
-                            progress.update_sub_task(fq_name, f"❌ Failed: {fq_name}")
-                            has_stream_errors = True
-
-                        # Update stream status (only for streams that aren't already completed)
-                        if fq_name in self.stream_task_ids and fq_name in progress._tasks and not progress._tasks[fq_name].completed:
-                            # Get the stream position from the stream order
-                            if fq_name in self.stream_order:
-                                stream_position = self.stream_order.index(fq_name) + 1
-                            else:
-                                # Fallback to 1 if not in order (shouldn't happen)
-                                stream_position = 1
-
-                            # Build status message
-                            if batches_count > 0 and status == 'syncing':
-                                status_message = f"Syncing {fq_name} ({stream_position}/{self.total_streams}), batches: {batches_count}"
-                            else:
-                                status_message = f"Pending {fq_name} ({stream_position}/{self.total_streams})..."
-
-                            progress.update_sub_task(fq_name, status_message)
-
-                            # Complete the stream if it's synced
-                            if status == "synced":
-                                progress.complete_sub_task(fq_name)
-
-                    # Add more streams from the queue if we have space and more streams exist
-                    if len(self.stream_task_ids) < MAX_DATA_SOURCE_SUBTASKS:
-                        for data in all_data:
-                            fq_name = data.get("fq_object_name", "Unknown")
-                            status = data.get("data_sync_status", "").lower()
-                            batches_count = data.get("pending_batches_count", 0)
-
-                            self._add_stream_subtask(progress, fq_name, status, batches_count)
-
-                    self.counter += 1
-
-                # Handle engines data
-                if not ready and engines:
-                    # Add new engines as subtasks if they don't exist
-                    for engine in engines:
-                        if not engine or not isinstance(engine, dict):
-                            continue
-                        size = self.engine_size
-                        name = engine.get("name", "Unknown")
-                        status = (engine.get("status") or "").lower()
-                        sub_task_id = self.engine_task_ids.get(name, None)
-                        sub_task_status_message = ""
-
-                        # Complete the sub task if it exists and the engine status is ready
-                        if sub_task_id and name in progress._tasks and not progress._tasks[name].completed and (status == "ready"):
-                            sub_task_status_message = f"Engine {name} ({size}) ready"
-                            progress.update_sub_task(name, sub_task_status_message)
-                            progress.complete_sub_task(name)
-
-                        # Add the sub task if it doesn't exist and the engine status is pending
-                        if not sub_task_id and status == "pending":
-                            writer = engine.get("writer", False)
-                            engine_type = "writer engine" if writer else "engine"
-                            sub_task_status_message = f"Provisioning {engine_type} {name} ({size})"
-                            self.engine_task_ids[name] = progress.add_sub_task(sub_task_status_message, task_id=name, category=TASK_CATEGORY_PROVISIONING)
-
-                    # Special handling for CDC_MANAGED_ENGINE - mark ready when any stream starts processing
-                    cdc_task = progress._tasks.get(CDC_MANAGED_ENGINE) if CDC_MANAGED_ENGINE in progress._tasks else None
-                    if CDC_MANAGED_ENGINE in self.engine_task_ids and cdc_task and not cdc_task.completed:
-
-                        has_processing_streams = any(
-                            stream.get("next_batch_status", "") == "processing"
-                            for stream in all_data
-                        )
-                        if has_processing_streams and cdc_task and not cdc_task.completed:
-                            progress.update_sub_task(CDC_MANAGED_ENGINE, f"Engine {CDC_MANAGED_ENGINE} ({self.engine_size}) ready")
-                            progress.complete_sub_task(CDC_MANAGED_ENGINE)
-
-                    self.counter += 1
-
-                # Handle relations data
-                if relations and isinstance(relations, dict):
-                    txn = relations.get("txn", {}) or {}
-                    txn_id = txn.get("id", None)
-
-                    # Only show relations subtask if there is a valid txn object
-                    if txn_id:
-                        status = relations.get("status", "").upper()
-
-                        # Create relations subtask if it doesn't exist
-                        if self.relations_task_id is None:
-                            self.relations_task_id = progress.add_sub_task("Populating relations", task_id="relations", category=TASK_CATEGORY_RELATIONS)
-
-                        # Set the start time from the JSON if available (always update)
-                        start_time_ms = relations.get("start_time")
-                        if start_time_ms:
-                            start_time_seconds = start_time_ms / 1000.0
-                            progress._tasks["relations"].added_time = start_time_seconds
-
-                        # Update relations status
-                        if status == "COMPLETED":
-                            progress.update_sub_task("relations", f"Relations populated (txn: {txn_id})")
-
-                            # Set the completion time from the JSON if available
-                            end_time_ms = relations.get("end_time")
-                            if end_time_ms:
-                                end_time_seconds = end_time_ms / 1000.0
-                                progress._tasks["relations"].completed_time = end_time_seconds
-
-                            progress.complete_sub_task("relations", record_time=False)  # Don't record local time
-                        else:
-                            progress.update_sub_task("relations", f"Relations populating (txn: {txn_id})")
-
-                        self.counter += 1
-
-                # Handle errors
-                if not ready and errors:
-                    for error in errors:
-                        if error is None:
-                            continue
-                        if error.get("type") == "data":
-                            message = error.get("message", "").lower()
-                            if ("change_tracking" in message or "change tracking" in message):
-                                err_source = error.get("source")
-                                err_source_type = self.source_info.get(err_source, {}).get("type")
-                                self.tables_with_not_enabled_change_tracking.append((err_source, err_source_type))
-                            else:
-                                self._raise_if_app_not_started(error.get("message", ""))
-                                self.table_objects_with_other_errors.append(
-                                    SnowflakeTableObject(error.get("message"), error.get("source"))
-                                )
-                        elif error.get("type") == "engine":
-                            self.engine_errors.append(error)
-                        else:
-                            # Other types of errors, e.g. "validation"
-                            self._raise_if_app_not_started(error.get("message", ""))
-                            self.table_objects_with_other_errors.append(
-                                SnowflakeTableObject(error.get("message"), error.get("source"))
-                            )
-                    has_general_errors = True
-
-                # If ready, complete all remaining subtasks
-                if ready:
-                    self.cache.record_update(self.source_info)
-                    # Complete any remaining stream subtasks
-                    for fq_name in self.stream_task_ids:
-                        if fq_name in progress._tasks and not progress._tasks[fq_name].completed:
-                            progress.complete_sub_task(fq_name)
-                    # Complete any remaining engine subtasks
-                    for name in self.engine_task_ids:
-                        if name in progress._tasks and not progress._tasks[name].completed:
-                            progress.complete_sub_task(name)
-                    # Complete relations subtask if it exists and isn't completed
-                    if self.relations_task_id and "relations" in progress._tasks and not progress._tasks["relations"].completed:
-                        progress.complete_sub_task("relations")
-                    break_loop = True
-                elif has_stream_errors or has_general_errors:
-                    # Break the loop if there are errors, but only after reporting all progress
-                    break_loop = True
-
-                return break_loop
-
-        poll_with_specified_overhead(lambda: check_ready(progress), overhead_rate=POLL_OVERHEAD_RATE, max_delay=POLL_MAX_DELAY)
-
-    def _raise_if_app_not_started(self, message: str) -> None:
-        """Detect Snowflake-side 'app not active / service not started' messages and raise a rich exception.
-
-        The use_index stored procedure reports many failures inside the returned JSON payload
-        (use_index_data['errors']) rather than raising them as Snowflake exceptions, so the
-        standard `_exec()` error handlers won't run. We detect the known activation-needed
-        signals here and raise `SnowflakeRaiAppNotStarted` for nicer formatting.
-        """
-        if not message:
-            return
-        msg = str(message).lower()
-        if (
-            "service has not been started" in msg
-            or "call app.activate()" in msg
-            or "app_not_active_exception" in msg
-            or "application is not active" in msg
-            or "use the app.activate()" in msg
-        ):
-            app_name = self.res.config.get("rai_app_name", "") if hasattr(self.res, "config") else ""
-            if not isinstance(app_name, str) or not app_name:
-                app_name = self.app_name
-            raise SnowflakeRaiAppNotStarted(app_name)
-
-    def _post_check(self, progress) -> None:
-        """Run post-processing checks including change tracking enablement.
-
-        Args:
-            progress: TaskProgress instance for tracking progress
-
-        Raises:
-            SnowflakeChangeTrackingNotEnabledException: If change tracking cannot be enabled
-            SnowflakeTableObjectsException: If there are table-related errors
-            EngineProvisioningFailed: If engine provisioning fails
-        """
-        num_tables_altered = 0
-        failed_tables = []  # Track tables that failed to enable change tracking
-
-        enabled_tables = []
-        if (
-            self.tables_with_not_enabled_change_tracking
-            and self.res.config.get("ensure_change_tracking", False)
-        ):
-            tables_to_process = self.tables_with_not_enabled_change_tracking
-            # Track timing for change tracking
-
-            # Add subtasks based on count
-            if len(tables_to_process) <= MAX_INDIVIDUAL_SUBTASKS:
-                # Add individual subtasks for each table
-                for i, table in enumerate(tables_to_process):
-                    fqn, kind = table
-                    progress.add_sub_task(f"Enabling change tracking on {fqn} ({i+1}/{len(tables_to_process)})", task_id=f"change_tracking_{i}", category=TASK_CATEGORY_CHANGE_TRACKING)
-            else:
-                # Add single summary subtask for many tables
-                progress.add_sub_task(f"Enabling change tracking on {len(tables_to_process)} tables", task_id="change_tracking_summary", category=TASK_CATEGORY_CHANGE_TRACKING)
-
-            # Process tables
-            for i, table in enumerate(tables_to_process):
-                fqn, kind = table  # Unpack outside try block to ensure fqn is defined
-
-                try:
-                    # Validate table_type to prevent SQL injection
-                    # Should only be TABLE or VIEW
-                    if kind not in ("TABLE", "VIEW"):
-                        logger.error(f"Invalid table kind '{kind}' for {fqn}, skipping")
-                        failed_tables.append((fqn, f"Invalid table kind: {kind}"))
-                        # Mark as failed in progress
-                        if len(tables_to_process) <= MAX_INDIVIDUAL_SUBTASKS:
-                            if f"change_tracking_{i}" in progress._tasks:
-                                progress.update_sub_task(f"change_tracking_{i}", f"❌ Invalid type: {fqn}")
-                                progress.complete_sub_task(f"change_tracking_{i}")
-                        continue
-
-                    # Execute ALTER statement
-                    # Note: fqn should already be properly quoted from source_info
-                    self.res._exec(f"ALTER {kind} {fqn} SET CHANGE_TRACKING = TRUE;")
-                    enabled_tables.append(table)
-                    num_tables_altered += 1
-
-                    # Update progress based on subtask type
-                    if len(tables_to_process) <= MAX_INDIVIDUAL_SUBTASKS:
-                        # Complete individual table subtask
-                        progress.complete_sub_task(f"change_tracking_{i}")
-                    else:
-                        # Update summary subtask with progress
-                        progress.update_sub_task("change_tracking_summary",
-                            f"Enabling change tracking on {len(tables_to_process)} tables... ({i+1}/{len(tables_to_process)})")
-                except Exception as e:
-                    # Log the error for debugging
-                    logger.warning(f"Failed to enable change tracking on {fqn}: {e}")
-                    failed_tables.append((fqn, str(e)))
-
-                    # Handle errors based on subtask type
-                    if len(tables_to_process) <= MAX_INDIVIDUAL_SUBTASKS:
-                        # Mark the individual subtask as failed and complete it
-                        if f"change_tracking_{i}" in progress._tasks:
-                            progress.update_sub_task(f"change_tracking_{i}", f"❌ Failed: {fqn}")
-                            progress.complete_sub_task(f"change_tracking_{i}")
-                    # Continue processing other tables despite this failure
-
-            # Complete summary subtask if used
-            if len(tables_to_process) > MAX_INDIVIDUAL_SUBTASKS and "change_tracking_summary" in progress._tasks:
-                if num_tables_altered > 0:
-                    s = "s" if num_tables_altered > 1 else ""
-                    success_msg = f"Enabled change tracking on {num_tables_altered} table{s}"
-                    if failed_tables:
-                        success_msg += f" ({len(failed_tables)} failed)"
-                    progress.update_sub_task("change_tracking_summary", success_msg)
-                elif failed_tables:
-                    progress.update_sub_task("change_tracking_summary", f"❌ Failed on {len(failed_tables)} table(s)")
-                progress.complete_sub_task("change_tracking_summary")
-
-            # Log summary of failed tables
-            if failed_tables:
-                logger.warning(f"Failed to enable change tracking on {len(failed_tables)} table(s)")
-                for fqn, error in failed_tables:
-                    logger.warning(f"  {fqn}: {error}")
-
-            # Remove the tables that were successfully enabled from the list of not enabled tables
-            # so that we don't raise an exception for them later
-            self.tables_with_not_enabled_change_tracking = [
-                t for t in self.tables_with_not_enabled_change_tracking if t not in enabled_tables
-            ]
-
-        if self.tables_with_not_enabled_change_tracking:
-            progress.update_main_status("Errors found. See below for details.")
-            raise SnowflakeChangeTrackingNotEnabledException(
-                self.tables_with_not_enabled_change_tracking
-            )
-
-        if self.table_objects_with_other_errors:
-            progress.update_main_status("Errors found. See below for details.")
-            raise SnowflakeTableObjectsException(self.table_objects_with_other_errors)
-        if self.engine_errors:
-            progress.update_main_status("Errors found. See below for details.")
-            # if there is an engine error, probably auto create engine failed
-            # Create a synthetic exception from the first engine error
-            first_error = self.engine_errors[0]
-            error_message = first_error.get("message", "Unknown engine error")
-            synthetic_exception = Exception(f"Engine error: {error_message}")
-            raise EngineProvisioningFailed(self.engine_name, synthetic_exception)
-
-        if num_tables_altered > 0:
-            self._poll_loop(progress)
-
-class DirectUseIndexPoller(UseIndexPoller):
-    """
-    Extends UseIndexPoller to handle direct-access prepare_index when no sources.
-    """
-    def __init__(
-        self,
-        resource: "DirectAccessResources",
-        app_name: str,
-        sources: Iterable[str],
-        model: str,
-        engine_name: str,
-        engine_size: Optional[str],
-        language: str = "rel",
-        program_span_id: Optional[str] = None,
-        headers: Optional[Dict] = None,
-        generation: Optional[Generation] = None,
-    ):
-        super().__init__(
-            resource=resource,
-            app_name=app_name,
-            sources=sources,
-            model=model,
-            engine_name=engine_name,
-            engine_size=engine_size,
-            language=language,
-            program_span_id=program_span_id,
-            headers=headers,
-            generation=generation,
-        )
-        from relationalai.clients.resources.snowflake import DirectAccessResources
-        self.res: DirectAccessResources = cast(DirectAccessResources, self.res)
-
-    def poll(self) -> None:
-        if not self.sources:
-            from relationalai.errors import RAIException
-            collected_errors: List[Dict] = []
-            attempt = 1
-
-            def check_direct(progress) -> bool:
-                nonlocal attempt
-                with debugging.span("wait", counter=self.counter) as span:
-                    span.update({"attempt": attempt, "engine_name": self.engine_name, "model": self.model})
-                    # we are skipping pulling relations here, as direct access only handle non-sources cases
-                    # and we don't need to pull relations for that, therefore, we pass empty list for rai_relations
-                    # and set skip_pull_relations to True
-                    resp = self.res._prepare_index(
-                        model=self.model,
-                        engine_name=self.engine_name,
-                        engine_size=self.engine_size,
-                        language=self.language,
-                        rai_relations=[],
-                        pyrel_program_id=self.program_span_id,
-                        skip_pull_relations=True,
-                        headers=self.headers,
-                    )
-                    span.update(resp)
-                    caller_engine = resp.get("caller_engine", {})
-                    # Handle case where caller_engine might be None
-                    ce_status = caller_engine.get("status", "").lower() if caller_engine else ""
-                    errors = resp.get("errors", [])
-
-                    ready = resp.get("ready", False)
-
-                    # Update main progress message
-                    if ready:
-                        progress.update_main_status("Done")
-                    else:
-                        progress.update_main_status("Preparing your data...")
-
-                    if ready:
-                        return True
-                    else:
-                        if ce_status == "pending":
-                            # Add or update engine subtask
-                            engine_name = caller_engine.get('name', self.engine_name)
-                            if not hasattr(progress, '_engine_task_id'):
-                                progress._engine_task_id = progress.add_sub_task(f"Waiting for engine '{engine_name}' to be ready...", task_id=engine_name)
-                            else:
-                                progress.update_sub_task(engine_name, f"Waiting for engine '{engine_name}' to be ready...")
-                        else:
-                            # Handle errors as subtasks
-                            if errors:
-                                progress.update_main_status("Encountered errors during preparation...")
-                                for i, err in enumerate(errors):
-                                    error_id = f"error_{i}"
-                                    if not hasattr(progress, f'_error_task_id_{i}'):
-                                        error_msg = err.get('message', 'Unknown error')
-                                        setattr(progress, f'_error_task_id_{i}', progress.add_sub_task(f"❌ {error_msg}", task_id=error_id))
-                                    else:
-                                        error_msg = err.get('message', 'Unknown error')
-                                        progress.update_sub_task(error_id, f"❌ {error_msg}")
-                                    collected_errors.append(err)
-
-                    attempt += 1
-                    return False
-
-            # Read show_duration_summary config flag (defaults to True for backward compatibility)
-            show_duration_summary = bool(self.res.config.get("show_duration_summary", True))
-
-            with create_progress(
-                description="Preparing your data...",
-                success_message="Done",
-                leading_newline=True,
-                trailing_newline=True,
-                show_duration_summary=show_duration_summary,
-            ) as progress:
-                # Add cache usage subtask
-                self._add_cache_subtask(progress)
-
-                with debugging.span("poll_direct"):
-                    poll_with_specified_overhead(lambda: check_direct(progress), overhead_rate=POLL_OVERHEAD_RATE, max_delay=POLL_MAX_DELAY)
-
-                # Run the same post-check logic as UseIndexPoller
-                self._post_check(progress)
-
-            if collected_errors:
-                msg = "; ".join(e.get("message", "") for e in collected_errors)
-                raise RAIException(msg)
-        else:
-            super().poll()