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

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,
-    )
-
-