relationalai 0.12.0__py3-none-any.whl → 0.12.2__py3-none-any.whl

relationalai/clients/use_index_poller.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import Iterable, Dict, Optional, List, cast, TYPE_CHECKING
 import json
 import logging
@@ -5,7 +7,11 @@ import uuid
 
 from relationalai import debugging
 from relationalai.clients.cache_store import GraphIndexCache
-from relationalai.clients.util import get_pyrel_version, poll_with_specified_overhead
+from relationalai.clients.util import (
+    get_pyrel_version,
+    normalize_datetime,
+    poll_with_specified_overhead,
+)
 from relationalai.errors import (
     ERPNotRunningError,
     EngineProvisioningFailed,
@@ -29,6 +35,7 @@ from relationalai.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
@@ -63,49 +70,49 @@ POLL_MAX_DELAY = 2.5  # Maximum delay between polls in seconds
 # 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 = """
-SELECT
-    FQ_OBJECT_NAME,
-    SHA2(
-        LISTAGG(
-            value:name::VARCHAR ||
+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 || ')'
+                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 ||
-            CASE WHEN value:nullable::BOOLEAN THEN 'YES' ELSE 'NO' END,
-            ','
-        ) WITHIN GROUP (ORDER BY value:name::VARCHAR),
-        256
-    ) AS STREAM_HASH
-FROM {app_name}.api.data_streams,
-LATERAL FLATTEN(input => COLUMNS) f
-WHERE RAI_DATABASE = '{rai_database}' AND FQ_OBJECT_NAME IN ({fqn_list})
-GROUP BY FQ_OBJECT_NAME;
+                    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;
 """
 
 
@@ -154,8 +161,9 @@ class UseIndexPoller:
         model: str,
         engine_name: str,
         engine_size: Optional[str],
-        program_span_id: Optional[str],
-        headers: Optional[Dict],
+        language: str = "rel",
+        program_span_id: Optional[str] = None,
+        headers: Optional[Dict] = None,
         generation: Optional[Generation] = None,
     ):
         self.res = resource
@@ -164,6 +172,7 @@ class UseIndexPoller:
         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
@@ -183,8 +192,8 @@ class UseIndexPoller:
         )
         current_user = self.res.get_sf_session().get_current_user()
         assert current_user is not None, "current_user must be set"
-        data_freshness = self.res.config.get_data_freshness_mins()
-        self.cache = GraphIndexCache(current_user, model, data_freshness, self.sources)
+        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())
@@ -296,9 +305,10 @@ class UseIndexPoller:
         Returns:
             List of truly stale sources that need to be deleted/recreated
 
-        A source is truly stale if:
-        - The stream doesn't exist (needs to be created), OR
-        - The column hashes don't match (needs to be 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)
 
@@ -306,14 +316,30 @@ class UseIndexPoller:
         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:
+            if stream_hash is None or source_hash != stream_hash or recreated_table:
                 logger.debug("  Action: DELETE (stale)")
                 truly_stale.append(source)
             else:
@@ -376,7 +402,7 @@ class UseIndexPoller:
             stale_sources = [
                 source
                 for source, info in self.source_info.items()
-                if info["state"] == "STALE"
+                if info.get("state") == "STALE"
             ]
 
         if not stale_sources:
@@ -462,6 +488,8 @@ class UseIndexPoller:
                 "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)
@@ -763,7 +791,7 @@ class UseIndexPoller:
                     # 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
@@ -829,11 +857,23 @@ class DirectUseIndexPoller(UseIndexPoller):
         model: str,
         engine_name: str,
         engine_size: Optional[str],
-        program_span_id: Optional[str],
-        headers: Optional[Dict],
+        language: str = "rel",
+        program_span_id: Optional[str] = None,
+        headers: Optional[Dict] = None,
         generation: Optional[Generation] = None,
     ):
-        super().__init__(resource, app_name, sources, model, engine_name, engine_size, program_span_id, headers, generation)
+        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.snowflake import DirectAccessResources
         self.res: DirectAccessResources = cast(DirectAccessResources, self.res)
 
@@ -854,6 +894,7 @@ class DirectUseIndexPoller(UseIndexPoller):
                         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,

relationalai/clients/util.py

@@ -80,6 +80,15 @@ def escape_for_f_string(code: str) -> str:
 def escape_for_sproc(code: str) -> str:
     return code.replace("$$", "\\$\\$")
 
+
+def normalize_datetime(value: object) -> datetime | None:
+    """Return a timezone-aware UTC datetime or None."""
+    if not isinstance(value, datetime):
+        return None
+    if value.tzinfo is None:
+        return value.replace(tzinfo=timezone.utc)
+    return value.astimezone(timezone.utc)
+
 # @NOTE: `overhead_rate` should fall between 0.05 and 0.5 depending on how time sensitive / expensive the operation in question is.
 def poll_with_specified_overhead(
     f,

relationalai/dsl.py

@@ -22,6 +22,7 @@ import sys
 from pandas import DataFrame
 
 from relationalai.environments import runtime_env, SnowbookEnvironment
+from relationalai.tools.constants import QUERY_ATTRIBUTES_HEADER
 
 from .clients.client import Client
 
@@ -34,9 +35,7 @@ from .errors import FilterAsValue, Errors, InvalidPropertySetException, Multiple
 #--------------------------------------------------
 
 RESERVED_PROPS = ["add", "set", "persist", "unpersist"]
-
 MAX_QUERY_ATTRIBUTE_LENGTH = 255
-QUERY_ATTRIBUTES_HEADER = "X-Query-Attributes"
 
 Value = Union[
     "Expression",

relationalai/environments/snowbook.py

@@ -20,7 +20,16 @@ class SnowbookEnvironment(NotebookRuntimeEnvironment, SessionEnvironment):
 
     def __init__(self):
         super().__init__()
-        self.runner = "container" if "snowflake.connector.auth" in sys.modules else "warehouse"
+        # Detect runner type based on module presence:
+        # - Warehouse runtime has '_snowflake' module
+        # - Container runtime has 'snowflake._legacy' module
+        if "_snowflake" in sys.modules:
+            self.runner = "warehouse"
+        elif "snowflake._legacy" in sys.modules:
+            self.runner = "container"
+        else:
+            # Fallback to original check
+            self.runner = "container" if "snowflake.connector.auth" in sys.modules else "warehouse"
 
     @classmethod
     def detect(cls):