relationalai 0.13.5__py3-none-any.whl → 1.0.0a1__py3-none-any.whl

relationalai/clients/resources/snowflake/error_handlers.py

@@ -1,240 +0,0 @@
-"""
-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

@@ -1,249 +0,0 @@
-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

@@ -1,99 +0,0 @@
-"""
-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,
-    )
-
-