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

relationalai/clients/resources/snowflake/error_handlers.py

@@ -0,0 +1,240 @@
+"""
+Error handlers for Snowflake Resources using Strategy Pattern.
+
+Each error handler encapsulates the detection logic and exception creation
+for a specific type of error. Handlers are processed in order until one matches.
+"""
+from __future__ import annotations
+import re
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any
+
+from ....errors import (
+    DuoSecurityFailed,
+    EngineNotFoundException,
+    EnginePending,
+    EngineNameValidationException,
+    EngineProvisioningFailed,
+    EngineResumeFailed,
+    RAIAbortedTransactionError,
+    RAIException,
+    SnowflakeAppMissingException,
+    SnowflakeDatabaseException,
+    SnowflakeRaiAppNotStarted,
+)
+
+if TYPE_CHECKING:
+    from .snowflake import ExecContext
+    from .snowflake import Resources
+
+from .util import is_database_issue, is_engine_issue, collect_error_messages
+
+class ErrorHandler(ABC):
+    """Base class for error handlers using Strategy Pattern."""
+
+    @abstractmethod
+    def matches(self, error: Exception, message: str, ctx: 'ExecContext', resources: 'Resources') -> bool:
+        """Check if this handler can process the error."""
+        pass
+
+    @abstractmethod
+    def handle(self, error: Exception, ctx: 'ExecContext', resources: 'Resources') -> Any | None:
+        """Handle the error and either raise an exception or return an alternate result."""
+        pass
+
+
+class DuoSecurityErrorHandler(ErrorHandler):
+    """Handle Duo security authentication errors."""
+
+    def matches(self, error: Exception, message: str, ctx: 'ExecContext', resources: 'Resources') -> bool:
+        messages = collect_error_messages(error)
+        return any("duo security" in msg for msg in messages)
+
+    def handle(self, error: Exception, ctx: 'ExecContext', resources: 'Resources') -> Any | None:
+        raise DuoSecurityFailed(error)
+
+
+class AppMissingErrorHandler(ErrorHandler):
+    """Handle missing RAI app database errors."""
+
+    def matches(self, error: Exception, message: str, ctx: 'ExecContext', resources: 'Resources') -> bool:
+        rai_app = resources.config.get("rai_app_name", "")
+        if not isinstance(rai_app, str):
+            return False
+        pattern = f"database '{rai_app}' does not exist or not authorized."
+        messages = collect_error_messages(error)
+        return any(re.search(pattern.lower(), msg) for msg in messages)
+
+    def handle(self, error: Exception, ctx: 'ExecContext', resources: 'Resources') -> Any | None:
+        rai_app = resources.config.get("rai_app_name", "")
+        assert isinstance(rai_app, str), f"rai_app_name must be a string, not {type(rai_app)}"
+        raise SnowflakeAppMissingException(rai_app, resources.config.get("role"))
+
+
+class AppFunctionMissingErrorHandler(ErrorHandler):
+    """Handle missing RAI app when app-scoped UDFs are unknown.
+
+    When the RelationalAI Snowflake native application isn't installed (or is installed
+    under a different name / not authorized), Snowflake can surface errors like:
+
+        SQL compilation error: Unknown user-defined function <APP>.<SCHEMA>.<FUNC>.
+
+    This should be presented to users as an "app missing" configuration issue, not
+    as a raw SQL error.
+    """
+
+    def matches(self, error: Exception, message: str, ctx: 'ExecContext', resources: 'Resources') -> bool:
+        rai_app = resources.config.get("rai_app_name", "")
+        if not isinstance(rai_app, str) or not rai_app:
+            return False
+
+        rai_app_lower = rai_app.lower()
+        # Normalize whitespace/newlines because Snowpark error strings may wrap.
+        messages = [" ".join(msg.split()).lower() for msg in collect_error_messages(error)]
+
+        # Examples:
+        # - "unknown user-defined function sqlib_ia_na_app.experimental.resume_engine_async."
+        # - "unknown user-defined function sqlib_ia_na_app.api.alter_engine_pool_node_limits."
+        needle = f"unknown user-defined function {rai_app_lower}."
+        return any(needle in msg for msg in messages)
+
+    def handle(self, error: Exception, ctx: 'ExecContext', resources: 'Resources') -> Any | None:
+        rai_app = resources.config.get("rai_app_name", "")
+        assert isinstance(rai_app, str), f"rai_app_name must be a string, not {type(rai_app)}"
+        raise SnowflakeAppMissingException(rai_app, resources.config.get("role"))
+
+
+class DatabaseErrorsHandler(ErrorHandler):
+    """Handle database-related errors from Snowflake/RAI."""
+
+    def matches(self, error: Exception, message: str, ctx: 'ExecContext', resources: 'Resources') -> bool:
+        """Check if error is database-related."""
+        # Use collect_error_messages to get all messages from exception chain
+        messages = collect_error_messages(error)
+        return any(is_database_issue(msg) for msg in messages)
+
+    def handle(self, error: Exception, ctx: 'ExecContext', resources: 'Resources') -> Any | None:
+        """Handle database errors and raise appropriate exception."""
+        raise SnowflakeDatabaseException(error)
+
+
+class EngineErrorsHandler(ErrorHandler):
+    """Handle all engine-related errors from Snowflake/RAI."""
+
+    def matches(self, error: Exception, message: str, ctx: 'ExecContext', resources: 'Resources') -> bool:
+        """Check if error is engine-related."""
+        # Use collect_error_messages to get all messages from exception chain
+        messages = collect_error_messages(error)
+        return any(is_engine_issue(msg) for msg in messages)
+
+    def handle(self, error: Exception, ctx: 'ExecContext', resources: 'Resources') -> Any | None:
+        """Handle engine errors and raise appropriate exception."""
+        # Use collect_error_messages to get all messages from exception chain
+        messages = collect_error_messages(error)
+        engine = resources.get_default_engine_name()
+        assert isinstance(engine, str), f"engine must be a string, not {type(engine)}"
+
+        # Check all collected messages for engine error patterns
+        for message in messages:
+            if "engine is in pending" in message or "engine is provisioning" in message:
+                raise EnginePending(engine)
+            elif "engine not found" in message or "no engines found" in message:
+                raise EngineNotFoundException(engine, str(error))
+            elif "engine was deleted" in message:
+                raise EngineNotFoundException(engine, "Engine was deleted")
+            elif "engine is suspended" in message:
+                raise EngineResumeFailed(engine)
+            elif "create/resume" in message:
+                raise EngineProvisioningFailed(engine, error)
+
+        # Generic engine error - use the original error message
+        raise RAIException(str(error))
+
+
+class ServiceNotStartedErrorHandler(ErrorHandler):
+    """Handle RAI service not started errors."""
+
+    def matches(self, error: Exception, message: str, ctx: 'ExecContext', resources: 'Resources') -> bool:
+        messages = [" ".join(msg.split()).lower() for msg in collect_error_messages(error)]
+        return any(
+            (
+                # Native app not activated / service not started
+                "service has not been started" in msg
+                # Native app suspended/deactivated; SPCS control plane is unreachable
+                or "not reachable: service suspended" in msg
+            )
+            for msg in messages
+        )
+
+    def handle(self, error: Exception, ctx: 'ExecContext', resources: 'Resources') -> Any | None:
+        rai_app = resources.config.get("rai_app_name", "")
+        assert isinstance(rai_app, str), f"rai_app_name must be a string, not {type(rai_app)}"
+        raise SnowflakeRaiAppNotStarted(rai_app)
+
+
+class TransactionAbortedErrorHandler(ErrorHandler):
+    """Handle transaction aborted errors with problem details."""
+
+    def matches(self, error: Exception, message: str, ctx: 'ExecContext', resources: 'Resources') -> bool:
+        messages = collect_error_messages(error)
+        return any(re.search(r"state:\s*aborted", msg) for msg in messages)
+
+    def handle(self, error: Exception, ctx: 'ExecContext', resources: 'Resources') -> Any | None:
+        # Use collect_error_messages to get all messages from exception chain
+        messages = collect_error_messages(error)
+        # Check all collected messages for transaction ID
+        for message in messages:
+            txn_id_match = re.search(r"id:\s*([0-9a-f\-]+)", message)
+            if txn_id_match:
+                problems = resources.get_transaction_problems(txn_id_match.group(1))
+                if problems:
+                    # Extract problem details (handle both dict and object formats)
+                    for problem in problems:
+                        if isinstance(problem, dict):
+                            type_field = problem.get('TYPE')
+                            message_field = problem.get('MESSAGE')
+                            report_field = problem.get('REPORT')
+                        else:
+                            type_field = problem.TYPE
+                            message_field = problem.MESSAGE
+                            report_field = problem.REPORT
+                        raise RAIAbortedTransactionError(type_field, message_field, report_field)
+        raise RAIException(str(error))
+
+
+class UseIndexRetryErrorHandler(ErrorHandler):
+    """Handle engine/database errors by polling use_index and retrying the execution.
+
+    Intended for UseIndexResources and subclasses. Register this handler *before*
+    the standard Database/Engine error handlers.
+    """
+
+    def matches(self, error: Exception, message: str, ctx: 'ExecContext', resources: 'Resources') -> bool:
+        if ctx.skip_engine_db_error_retry:
+            return False
+        messages = collect_error_messages(error)
+        return any(msg and (is_database_issue(msg) or is_engine_issue(msg)) for msg in messages)
+
+    def handle(self, error: Exception, ctx: 'ExecContext', resources: 'Resources') -> Any | None:
+        poll_use_index = getattr(resources, "_poll_use_index", None)
+        if not callable(poll_use_index):
+            return None
+
+        engine = resources.get_default_engine_name()
+        engine_size = resources.config.get_default_engine_size()
+        assert isinstance(engine, str), f"engine must be a string, not {type(engine)}"
+
+        model = getattr(resources, "database", "")
+        try:
+            poll_use_index(
+                app_name=resources.get_app_name(),
+                sources=resources.sources,
+                model=model,
+                engine_name=engine,
+                engine_size=engine_size,
+            )
+        except EngineNameValidationException as e:
+            raise EngineNameValidationException(engine) from e
+
+        return ctx.re_execute(resources)
+

relationalai/clients/resources/snowflake/export_procedure.py.jinja

@@ -0,0 +1,249 @@
+from __future__ import annotations
+import json
+import textwrap
+import random
+import string
+import time
+from zoneinfo import ZoneInfo
+from snowflake.snowpark.types import StringType, IntegerType, StructField, StructType, FloatType, MapType, ArrayType, BooleanType, BinaryType, DateType, TimestampType
+import logging
+
+from snowflake.snowpark.exceptions import SnowparkSQLException
+
+class RAIException(Exception):
+    pass
+
+def random_string():
+    chars = string.ascii_letters + string.digits
+    return ''.join(random.choice(chars) for _ in range(32))
+
+def run_sql(session, query, params=None, engine=""):
+    try:
+        return session.sql(query, params).collect()
+    except SnowparkSQLException as e:
+        if "engine not found" in e.message:
+            import sys
+            def exception_handler(exception_type, exception, traceback):
+                print(f"{exception_type.__name__}: {exception}")
+            sys.excepthook = exception_handler
+            raise RAIException(f"RelationalAI engine not found. Please create an engine called `{engine}` using `___RAI_APP___.api.create_engine`.") from None
+        else:
+            raise e
+
+def poll_with_specified_overhead(
+    f,
+    overhead_rate: float,
+    start_time: float | None = None,
+    timeout: int | None = None,
+    max_tries: int | None = None,
+    max_delay: float = 120,
+    min_delay: float = 0.1
+):
+    if overhead_rate < 0:
+        raise ValueError("overhead_rate must be non-negative")
+
+    if start_time is None:
+        start_time = time.time()
+
+    tries = 0
+    max_time = time.time() + timeout if timeout else None
+
+    while True:
+        if f():
+            break
+
+        current_time = time.time()
+
+        if max_tries is not None and tries >= max_tries:
+            raise Exception(f'max tries {max_tries} exhausted')
+
+        if max_time is not None and current_time >= max_time:
+            raise Exception(f'timed out after {timeout} seconds')
+
+        duration = (current_time - start_time) * overhead_rate
+        duration = max(min(duration, max_delay), min_delay, 0)
+
+        time.sleep(duration)
+        tries += 1
+
+def check_ready(session, sql_string, database, engine, wait_for_stream_sync, use_index_id):
+    results = run_sql(session, sql_string, [database, engine, wait_for_stream_sync, use_index_id])
+    # Extract the JSON string from the `USE_INDEX` field
+    use_index_json_str = results[0]["USE_INDEX"]
+
+    # Convert the JSON string to lowercase
+    use_index_json_str = use_index_json_str.lower()
+
+    # Parse the JSON string into a Python dictionary
+    use_index_data = json.loads(use_index_json_str)
+    ready = use_index_data.get("ready", False)
+    errors = use_index_data.get("errors", [])
+
+    if ready:
+        return True
+    elif errors:
+        error_strs = []
+        for error in errors:
+            if error.get("type") == "data":
+                error_strs.append(f"{error.get('message')}, source: {error.get('source')}")
+            elif error.get("type") == "engine":
+                error_strs.append(f"{error.get('message')}")
+        raise Exception("\n".join(error_strs))
+
+def get_engine(session, passed_engine):
+    if passed_engine:
+        return passed_engine
+    try:
+        APP_NAME = {{ APP_NAME }}
+        query = f"select * from {APP_NAME}.api.engines where created_by=current_user() and status = 'READY' and name != 'CDC_MANAGED_ENGINE' order by created_on desc;"
+        row = run_sql(session, query)[0]
+        return row["NAME"]
+    except:
+        return {{ engine }}
+
+def escape(code):
+    return (
+        code
+        .replace("\\", "\\\\\\\\")
+        .replace("'", "\\\'")
+        .replace('"', '\\\\"')
+        .replace("\n", "\\n")
+    )
+def replace_column_names(rejected_rows: list[dict], col_names_map: dict) -> list[dict]:
+    # col_names_map looks like this: {col000: TIMESTAMPS, col002: id}
+    # COLUMN_NAME in rejected_rows looks like this: "OUT23A38988_10B1_4BCB_B566_89F7F3C8D1FB_INTERNAL"["COL000":1]
+
+    parsing_errors = []
+
+    # Convert col_names_map keys to lowercase for case-insensitive matching
+    col_names_map = {key.lower(): value for key, value in col_names_map.items()}
+    for i, row in enumerate(rejected_rows):
+        try:
+            col_name = row['COLUMN_NAME'].split('[')[1].split(':')[0].replace('"', '').strip().lower()
+            if col_name in col_names_map:
+                row['COLUMN_NAME'] = col_names_map[col_name]
+        except Exception as e:
+            parsing_errors.append({"index": i, "message": str(e), "row": str(row)})
+
+    return parsing_errors
+
+def format_rejected_rows(rejected_rows: list[dict], rejected_rows_count: int, col_names_map: dict) -> str:
+    parsing_errors = replace_column_names(rejected_rows, col_names_map)
+    error_indices = {e["index"] for e in parsing_errors}
+    msg = ""
+    grouped_errors = {}
+
+    # Group errors by ROW_NUMBER
+    for i, row in enumerate(rejected_rows):
+        if i in error_indices:
+            continue
+        try:
+            row_number = row['ROW_NUMBER']
+            if row_number not in grouped_errors:
+                grouped_errors[row_number] = {'rejected_record': row['REJECTED_RECORD'], 'errors': []}
+            grouped_errors[row_number]['errors'].append(f"Erroneous column: {row['COLUMN_NAME']}\nError message: {row['ERROR']}")
+        except Exception as e:
+            parsing_errors.append({"index": i, "message": str(e), "row": str(row)})
+
+    msg = f"{rejected_rows_count} rows were skipped due to erroneous data. Here are the first {len(grouped_errors)} rejected rows:\n"
+    for row_number, data in grouped_errors.items():
+        msg += f"""Rejected record: {data['rejected_record']}"""
+        for error in data['errors']:
+            msg += f"- {error}\n"
+        msg += "\n"
+
+    return msg
+
+{% if has_return_hint %}
+def handle(session{{py_inputs}}, passed_engine=""):
+{% else %}
+def handle(session{{py_inputs}}, save_as_table="", passed_engine=""):
+{% endif %}
+    {% if not has_return_hint %}
+    if not save_as_table:
+        raise ValueError("`save_as_table` must be provided.")
+    {% endif %}
+    try:
+        {{clean_inputs}}
+        engine = get_engine(session, passed_engine)
+        logging.debug(f"Using engine: {engine}")
+        rel_code = {{ rel_code }}
+        logging.debug(f"RelationalAI code: {rel_code}")
+        proc_database = {{ proc_database }}
+        logging.debug(f"Procedure database: {proc_database}")
+        APP_NAME = {{ APP_NAME }}
+        logging.debug(f"APP_NAME: {APP_NAME}")
+        database = {{ database }}
+        logging.debug(f"Model database: {database}")
+        sql_out_names = [{{ sql_out_names }}]
+        logging.debug(f"Output column names: {sql_out_names}")
+        table_name = f"{database[:30]}_{random_string()}"
+        temp_table = f"temp_{table_name}"
+        logging.debug(f"Created temporary table name: {temp_table}")
+        nowait_durable = True
+        skip_invalid_data = {{ skip_invalid_data }}
+
+        sources = "{{ source_references }}"
+        logging.debug(f"Using source references: {sources}")
+        wait_for_stream_sync = {{ wait_for_stream_sync }}
+        use_index_id = f"{proc_database}_{random_string()}"
+        sql_string = f"CALL {APP_NAME}.api.use_index([{sources}], {{'model': ?, 'engine': ?, 'wait_for_stream_sync': ?, 'use_index_id': ?}});"
+        logging.debug("Polling for graph index readiness...")
+        poll_with_specified_overhead(lambda: check_ready(session, sql_string, proc_database, engine, wait_for_stream_sync, use_index_id), overhead_rate=0.1, max_delay=1)
+        logging.debug("Graph index is ready")
+        install_code = '''
+        {{ install_code }}
+        '''
+        if install_code:
+            rel_code = install_code + "\n" + rel_code
+        # signature: exec_into_table(model, engine, query, table_name, readonly, nowait_durable)
+        res = run_sql(session, f"call {APP_NAME}.api.exec_into_table(?, ?, ?, ?, ?, ?, ?);", [proc_database, engine, rel_code, table_name, False, nowait_durable, skip_invalid_data], engine=engine)
+        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", [])
+        logging.debug("Executed query into table")
+
+        logging.debug("Sampling output to determine schema...")
+        out_sample = run_sql(session, f"select * from {APP_NAME}.results.{table_name} limit 1;")
+        keys = set()
+        if out_sample:
+            keys = set([k.lower() for k in out_sample[0].as_dict().keys()])
+            logging.debug(f"Found output columns: {keys}")
+        else:
+            logging.debug("No output rows found in sample")
+
+        {% if has_return_hint %}
+        names = ", ".join([f"CAST(col{ix:03} as {type_name}) as \"{name}\"" if f"col{ix:03}" in keys else f"NULL as \"{name}\"" for (ix, (name, type_name)) in enumerate(sql_out_names)])
+        logging.debug(f"Delete result table and output table with columns: {names}")
+        return session.sql(f"""
+        BEGIN
+            let r resultset := (select {names} from {APP_NAME}.results.{table_name});
+            call {APP_NAME}.api.drop_result_table('{table_name}');
+            return table(r);
+        END;
+        """)
+        {% else %}
+        col_names_map = {}
+        for ix, name in enumerate(sql_out_names):
+            col_key = f"col{ix:03}"
+            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()
+        ])
+        logging.debug(f"Creating final output table with columns: {names}")
+        run_sql(session, f"create or replace table {save_as_table} as select {names} from {APP_NAME}.results.{table_name};", [], engine=engine)
+        logging.debug(f"Successfully wrote results to {save_as_table}")
+        if rejected_rows:
+            rejected_rows_formatted_msg = format_rejected_rows(rejected_rows, rejected_rows_count, col_names_map)
+            return f"Results written to {save_as_table} but {rejected_rows_formatted_msg}"
+        else:
+            return f"Results written to {save_as_table}"
+        {% endif %}
+    except Exception as e:
+        logging.debug(f"Error occurred: {str(e)}")
+        msg = str(e).lower()
+        if "no columns returned" in msg or "columns of results could not be determined" in msg:
+            logging.debug("No results returned - creating empty dataframe")
+            return session.createDataFrame([], StructType([{{ py_outs }}]))
+        raise e

relationalai/clients/resources/snowflake/resources_factory.py

@@ -0,0 +1,99 @@
+"""
+Factory function for creating Resources instances based on configuration.
+
+This module provides a factory function that selects the appropriate Resources
+class based on configuration settings (platform, use_direct_access, use_graph_index).
+"""
+from __future__ import annotations
+from typing import Union
+from snowflake.snowpark import Session
+
+from ...config import Config
+from ...local import LocalResources
+from . import Resources, DirectAccessResources, UseIndexResources
+from ....tools.constants import USE_DIRECT_ACCESS, USE_GRAPH_INDEX, Generation
+
+
+def create_resources_instance(
+    config: Config | None = None,
+    profile: str | None = None,
+    dry_run: bool = False,
+    connection: Session | None = None,
+    reset_session: bool = False,
+    language: str = "rel",
+    generation: Generation = Generation.QB,
+) -> Union[LocalResources, DirectAccessResources, UseIndexResources, Resources]:
+    """
+    Factory function that creates the appropriate Resources instance based on config.
+
+    This function selects the Resources class based on:
+    1. Platform (local -> LocalResources)
+    2. use_direct_access flag (DirectAccessResources)
+    3. use_graph_index flag (UseIndexResources)
+    4. Default (base Resources)
+
+    Args:
+        config: Configuration object (optional, will create from profile if not provided)
+        profile: Profile name (optional, used if config is not provided)
+        dry_run: Whether to run in dry-run mode (default: False)
+        connection: Optional Snowflake session connection
+        reset_session: Whether to reset the session (default: False)
+        language: Language for the resources instance (default: "rel")
+        generation: Generation for the resources instance (default: Generation.QB)
+
+    Returns:
+        Appropriate Resources instance based on config settings:
+        - LocalResources if platform is "local"
+        - DirectAccessResources if use_direct_access is enabled
+        - UseIndexResources if use_graph_index is enabled
+        - Resources (base) otherwise
+    """
+    # Create config from profile if not provided
+    if config is None:
+        config = Config(profile)
+
+    platform = config.get("platform", "")
+    if platform == "local":
+        return LocalResources(
+            profile=profile,
+            config=config,
+            dry_run=dry_run,
+            generation=generation,
+            connection=connection,
+            reset_session=reset_session,
+            language=language,
+        )
+
+    if config.get("use_direct_access", USE_DIRECT_ACCESS):
+        return DirectAccessResources(
+            profile=profile,
+            config=config,
+            dry_run=dry_run,
+            generation=generation,
+            connection=connection,
+            reset_session=reset_session,
+            language=language,
+        )
+
+    if config.get("use_graph_index", USE_GRAPH_INDEX):
+        return UseIndexResources(
+            profile=profile,
+            config=config,
+            dry_run=dry_run,
+            generation=generation,
+            connection=connection,
+            reset_session=reset_session,
+            language=language,
+        )
+
+    return Resources(
+        profile=profile,
+        config=config,
+        dry_run=dry_run,
+        generation=generation,
+        connection=connection,
+        reset_session=reset_session,
+        language=language,
+    )
+
+