relationalai 1.0.0a2__py3-none-any.whl → 1.0.0a4__py3-none-any.whl

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

@@ -5,21 +5,23 @@ import json
 import logging
 import uuid
 
-from v0.relationalai import debugging
-from v0.relationalai.clients.cache_store import GraphIndexCache
-from v0.relationalai.clients.util import (
+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 v0.relationalai.errors import (
+from ....errors import (
     ERPNotRunningError,
     EngineProvisioningFailed,
     SnowflakeChangeTrackingNotEnabledException,
     SnowflakeTableObjectsException,
     SnowflakeTableObject,
+    SnowflakeRaiAppNotStarted,
 )
-from v0.relationalai.tools.cli_controls import (
+from ....tools.cli_controls import (
     DebuggingSpan,
     create_progress,
     TASK_CATEGORY_INDEXING,
@@ -30,7 +32,7 @@ from v0.relationalai.tools.cli_controls import (
     TASK_CATEGORY_STATUS,
     TASK_CATEGORY_VALIDATION,
 )
-from v0.relationalai.tools.constants import WAIT_FOR_STREAM_SYNC, Generation
+from ....tools.constants import WAIT_FOR_STREAM_SYNC, Generation
 
 # Set up logger for this module
 logger = logging.getLogger(__name__)
@@ -44,8 +46,8 @@ except ImportError:
     Table = None
 
 if TYPE_CHECKING:
-    from v0.relationalai.clients.snowflake import Resources
-    from v0.relationalai.clients.snowflake import DirectAccessResources
+    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
@@ -187,6 +189,9 @@ class UseIndexPoller:
         # on every 5th iteration we reset the cdc status, so it will be checked again
         self.should_check_cdc = True
 
+        # Flag to only check data stream health once in the first call
+        self.check_data_stream_health = True
+
         self.wait_for_stream_sync = self.res.config.get(
             "wait_for_stream_sync", WAIT_FOR_STREAM_SYNC
         )
@@ -278,7 +283,7 @@ class UseIndexPoller:
         Raises:
             ValueError: If the query fails (permissions, table doesn't exist, etc.)
         """
-        from v0.relationalai.clients.snowflake import PYREL_ROOT_DB
+        from v0.relationalai.clients.resources.snowflake import PYREL_ROOT_DB
 
         # Build FQN list for SQL IN clause
         fqn_list = ", ".join([f"'{source}'" for source in sources])
@@ -427,7 +432,7 @@ class UseIndexPoller:
                     return
 
                 # Delete truly stale streams
-                from v0.relationalai.clients.snowflake import PYREL_ROOT_DB
+                from v0.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)
@@ -456,7 +461,8 @@ class UseIndexPoller:
                     )
 
                 # Don't raise if streams don't exist - this is expected
-                if "data streams do not exist" not in str(e).lower():
+                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:
@@ -500,6 +506,7 @@ class UseIndexPoller:
                 "init_engine_async": self.init_engine_async,
                 "language": self.language,
                 "data_freshness_mins": self.data_freshness,
+                "check_data_stream_health": self.check_data_stream_health
             })
 
             request_headers = debugging.add_current_propagation_headers(self.headers)
@@ -532,6 +539,7 @@ class UseIndexPoller:
                 errors = use_index_data.get("errors", [])
                 relations = use_index_data.get("relations", {})
                 cdc_enabled = use_index_data.get("cdcEnabled", False)
+                health_checked = use_index_data.get("healthChecked", False)
                 if self.check_ready_count % ERP_CHECK_FREQUENCY == 0 or not cdc_enabled:
                     self.should_check_cdc = True
                 else:
@@ -539,6 +547,9 @@ class UseIndexPoller:
 
                 if engines and self.init_engine_async:
                     self.init_engine_async = False
+                
+                if self.check_data_stream_health and health_checked:
+                    self.check_data_stream_health = False
 
                 break_loop = False
                 has_stream_errors = False
@@ -577,6 +588,9 @@ class UseIndexPoller:
                         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)
                                 )
@@ -702,6 +716,7 @@ class UseIndexPoller:
                                 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"))
                                 )
@@ -709,6 +724,7 @@ class UseIndexPoller:
                             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"))
                             )
@@ -737,6 +753,29 @@ class UseIndexPoller:
 
         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.
 
@@ -887,7 +926,7 @@ class DirectUseIndexPoller(UseIndexPoller):
             headers=headers,
             generation=generation,
         )
-        from v0.relationalai.clients.snowflake import DirectAccessResources
+        from v0.relationalai.clients.resources.snowflake import DirectAccessResources
         self.res: DirectAccessResources = cast(DirectAccessResources, self.res)
 
     def poll(self) -> None:

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

@@ -0,0 +1,188 @@
+"""
+Use Index Resources - Resources class with use_index functionality.
+This class keeps the use_index retry logic in _exec and provides use_index methods.
+"""
+from __future__ import annotations
+from typing import Iterable, Dict, Any
+
+from .use_index_poller import UseIndexPoller
+from ...config import Config
+from ....tools.constants import Generation
+from snowflake.snowpark import Session
+from .error_handlers import ErrorHandler, UseIndexRetryErrorHandler
+
+# Import Resources from snowflake - this creates a dependency but no circular import
+# since snowflake.py doesn't import from this file
+from .snowflake import Resources
+from .util import (
+    is_engine_issue as _is_engine_issue,
+    is_database_issue as _is_database_issue,
+    collect_error_messages,
+)
+
+
+class UseIndexResources(Resources):
+    """
+    Resources class with use_index functionality.
+    Provides use_index polling methods and keeps use_index retry logic in _exec.
+    """
+    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",
+    ):
+        super().__init__(
+            profile=profile,
+            config=config,
+            connection=connection,
+            dry_run=dry_run,
+            reset_session=reset_session,
+            generation=generation,
+        )
+        self.database = ""
+        self.language = language
+
+    def _is_db_or_engine_error(self, e: Exception) -> bool:
+        """Check if an exception indicates a database or engine error."""
+        messages = collect_error_messages(e)
+        for msg in messages:
+            if msg and (_is_database_issue(msg) or _is_engine_issue(msg)):
+                return True
+        return False
+
+    def _get_error_handlers(self) -> list[ErrorHandler]:
+        # Ensure use_index retry happens before standard database/engine error handlers.
+        return [UseIndexRetryErrorHandler(), *super()._get_error_handlers()]
+
+    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,
+    ):
+        """Poll use_index to prepare indices for the given sources."""
+        return UseIndexPoller(
+            self,
+            app_name,
+            sources,
+            model,
+            engine_name,
+            engine_size,
+            self.language,
+            program_span_id,
+            headers,
+            self.generation
+        ).poll()
+
+    def maybe_poll_use_index(
+        self,
+        app_name: str,
+        sources: Iterable[str],
+        model: str,
+        engine_name: str,
+        engine_size: str | None = None,
+        program_span_id: str | None = None,
+        headers: Dict | None = None,
+    ):
+        """Only call poll() if there are sources to process and cache is not valid."""
+        sources_list = list(sources)
+        self.database = model
+        if sources_list:
+            poller = UseIndexPoller(
+                self,
+                app_name,
+                sources_list,
+                model,
+                engine_name,
+                engine_size,
+                self.language,
+                program_span_id,
+                headers,
+                self.generation
+            )
+            # If cache is valid (data freshness has not expired), skip polling
+            if poller.cache.is_valid():
+                cached_sources = len(poller.cache.sources)
+                total_sources = len(sources_list)
+                cached_timestamp = poller.cache._metadata.get("cachedIndices", {}).get(poller.cache.key, {}).get("last_use_index_update_on", "")
+
+                message = f"Using cached data for {cached_sources}/{total_sources} data streams"
+                if cached_timestamp:
+                    print(f"\n{message} (cached at {cached_timestamp})\n")
+                else:
+                    print(f"\n{message}\n")
+            else:
+                return poller.poll()
+
+    def _exec_with_gi_retry(
+        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,
+    ):
+        """Execute with graph index retry logic.
+
+        Attempts execution with gi_setup_skipped=True first. If an engine or database
+        issue occurs, polls use_index and retries with gi_setup_skipped=False.
+        """
+        try:
+            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,
+            )
+        except Exception as e:
+            if not self._is_db_or_engine_error(e):
+                raise e
+
+            engine_name = engine or self.get_default_engine_name()
+            engine_size = self.config.get_default_engine_size()
+            self._poll_use_index(
+                app_name=self.get_app_name(),
+                sources=self.sources,
+                model=database,
+                engine_name=engine_name,
+                engine_size=engine_size,
+                headers=headers,
+            )
+
+            return self._exec_async_v2(
+                database, engine, raw_code, inputs, readonly, nowait_durable,
+                headers=headers, bypass_index=bypass_index, language=language,
+                query_timeout_mins=query_timeout_mins, gi_setup_skipped=False,
+            )
+
+    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:
+        """Override to use retry logic with use_index polling."""
+        return self._exec_with_gi_retry(
+            database, engine, raw_code, inputs, readonly, nowait_durable,
+            headers, bypass_index, language, query_timeout_mins
+        )