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

relationalai/clients/resources/azure/azure.py

@@ -0,0 +1,502 @@
+from __future__ import annotations
+import atexit
+from datetime import datetime, timedelta
+import json
+import textwrap
+import time
+from typing import Any, Dict, Tuple, List, cast
+from urllib.error import HTTPError
+
+from pandas import DataFrame
+
+from .... import debugging
+from ...util import poll_with_specified_overhead
+
+from ....errors import EngineNotFoundException, RAIException, AzureUnsupportedQueryTimeoutException
+from ....rel_utils import assert_no_problems
+from ....loaders.loader import emit_delete_import, import_file, list_available_resources
+from ...config import Config
+from ...types import EngineState, ImportSource, ImportSourceFile, TransactionAsyncResponse
+from ...client import Client, ExportParams, ProviderBase, ResourcesBase
+from .... import dsl, rel, metamodel as m
+from railib import api
+from ... import result_helpers
+
+#--------------------------------------------------
+# Constants
+#--------------------------------------------------
+
+TXN_FIELDS = ["id", "account_name", "state", "created_on", "finished_at", "duration", "database_name", "read_only", "engine_name", "query", "query_size", "tags", "user_agent", "response_format_version"]
+TXN_REPLACE_MAP = {"database_name": "database", "engine_name": "engine", "account_name": "account", "user_agent": "agent"}
+VALID_ENGINE_STATES = ["REQUESTED", "PROVISIONED", "PROVISIONING"]
+ENGINE_SIZES = ["XS", "S", "M", "L", "XL"]
+
+#--------------------------------------------------
+# Resources
+#--------------------------------------------------
+
+class Resources(ResourcesBase):
+    def __init__(self, profile:str|None=None, config:Config|None=None):
+        super().__init__(profile, config=config)
+        self._ctx = None
+        atexit.register(self.cancel_pending_transactions)
+
+    def _api_ctx(self):
+        if not self._ctx:
+            self._ctx = api.Context(**self.config.to_rai_config())
+        return self._ctx
+
+    def reset(self):
+        self._ctx = None
+
+    #--------------------------------------------------
+    # Generic
+    #--------------------------------------------------
+
+    def get_version(self):
+        raise Exception("Azure version not available")
+
+    #--------------------------------------------------
+    # Databases
+    #--------------------------------------------------
+
+    # Note: in contrast to the API definition in `ResourcesBase`, `get_database` in Azure
+    # can return a `List` object instead of a `Dict` object.
+    def get_database(self, database:str):
+        return api.get_database(self._api_ctx(), database)
+
+    # not implemented in Azure
+    def get_installed_packages(self, database: str) -> Dict | None:
+        return super().get_installed_packages(database)
+
+    #--------------------------------------------------
+    # Engines
+    #--------------------------------------------------
+    def get_engine_sizes(self, cloud_provider: str|None=None):
+        return ENGINE_SIZES
+
+    def get_cloud_provider(self) -> str:
+        return "azure"
+
+    def list_engines(
+        self,
+        state: str | None = None,
+        name: str | None = None,
+        type: str | None = None,
+        size: str | None = None,
+        created_by: str | None = None,
+    ):
+        # Azure only supports state filtering at the API level; other filters are ignored.
+        engines = api.list_engines(self._api_ctx(), state)
+        # Ensure EngineState shape includes 'type' for callers/tests.
+        for eng in engines:
+            eng.setdefault("type", "LOGIC")
+        return engines
+
+    def get_engine(self, name: str, type: str) -> EngineState:
+        # type is ignored for Azure as it doesn't support multiple engine types
+        engine = cast(EngineState, api.get_engine(self._api_ctx(), name))
+        engine.setdefault("type", "LOGIC")
+        return engine
+
+    def is_valid_engine_state(self, name:str):
+        return name in VALID_ENGINE_STATES
+
+    def create_engine(
+        self,
+        name: str,
+        type: str | None = None,
+        size: str | None = None,
+        auto_suspend_mins: int | None = None,
+        headers: dict | None = None,
+        settings: dict | None = None,
+    ):
+        # Azure only supports one engine type, so type parameter is ignored
+        if size is None:
+            size = "M"
+        with debugging.span("create_engine", name=name, size=size):
+            return api.create_engine_wait(self._api_ctx(), name, size)
+
+    def delete_engine(self, name:str, type: str):
+        return api.delete_engine(self._api_ctx(), name)
+
+    def suspend_engine(self, name:str, type: str | None = None): # type is ignored for Azure
+        return api.suspend_engine(self._api_ctx(), name)
+
+    def resume_engine(self, name:str, type: str | None = None, headers={}):
+        # type is ignored for Azure as it doesn't support multiple engine types
+        return api.resume_engine_wait(self._api_ctx(), name)
+
+    def resume_engine_async(self, name:str, type: str | None = None, headers: Dict | None = None):
+        return api.resume_engine(self._api_ctx(), name)
+
+    def auto_create_engine_async(self, name: str | None = None, type: str | None = None) -> str:
+        raise Exception("Azure doesn't support auto_create_engine_async")
+
+    def alter_engine_pool(self, size: str | None = None, mins: int | None = None, maxs: int | None = None):
+        raise Exception("Azure doesn't support engine pool alteration")
+
+    #--------------------------------------------------
+    # Graphs
+    #--------------------------------------------------
+
+    def list_graphs(self) -> List[Any]:
+        with debugging.span("list_models"):
+            return api.list_databases(self._api_ctx())
+
+    def get_graph(self, name:str):
+        with debugging.span("get_model", name=name):
+            return api.get_database(self._api_ctx(), name)
+
+    def create_graph(self, name: str):
+        with debugging.span("create_model", name=name):
+            return api.create_database(self._api_ctx(), name)
+
+    def delete_graph(self, name:str):
+        with debugging.span("delete_model", name=name):
+            return api.delete_database(self._api_ctx(), name)
+
+    def clone_graph(self, target_name:str, source_name:str, nowait_durable:bool=False):
+        # not a typo: the argument order is indeed target then source
+        return api.create_database(self._api_ctx(), target_name, source_name)
+
+    #--------------------------------------------------
+    # Models
+    #--------------------------------------------------
+
+    def list_models(self, database: str, engine: str):
+        return api.list_databases(self._api_ctx())
+
+    def create_models(self, database: str, engine: str, models:List[Tuple[str, str]]) -> List[Any]:
+        rel_code = self.create_models_code(models)
+        results = self.exec_raw(database, engine, rel_code, readonly=False)
+        if results.problems:
+            return results.problems
+        return []
+
+    def delete_model(self, database:str, engine:str, name:str):
+        return api.delete_model(self._api_ctx(), database, engine, name)
+
+    def create_models_code(self, models:List[Tuple[str, str]]) -> str:
+        lines = []
+        for (name, code) in models:
+            name = name.replace("\"", "\\\"")
+            assert "\"\"\"\"\"\"\"" not in code, "Code literals must use fewer than 7 quotes."
+
+            lines.append(textwrap.dedent(f"""
+            def delete[:rel, :catalog, :model, "{name}"]: rel[:catalog, :model, "{name}"]
+            def insert[:rel, :catalog, :model, "{name}"]: raw\"\"\"\"\"\"\"
+            """) + code + "\n\"\"\"\"\"\"\"")
+        rel_code = "\n\n".join(lines)
+        return rel_code
+
+    #--------------------------------------------------
+    # Exports
+    #--------------------------------------------------
+
+    def list_exports(self, database: str, engine: str):
+        raise Exception("Azure doesn't support exports")
+
+    def create_export(self, params: ExportParams):
+        if not params.dry_run:
+            raise Exception("Azure doesn't support exports")
+
+    def create_export_table(self, database: str, engine: str, table: str, relation: str, columns: Dict[str, str], code: str, refresh: str|None=None):
+        raise Exception("Azure doesn't support exports")
+
+    def delete_export(self, database: str, engine: str, name: str):
+        raise Exception("Azure doesn't support exports")
+
+    #--------------------------------------------------
+    # Imports
+    #--------------------------------------------------
+
+    def list_imports(self, id:str|None = None, name:str|None = None, model:str|None = None, status:str|None = None, creator:str|None = None):
+        if not model:
+            raise RAIException("Imports can only be listed for a particular model in azure")
+        return [*list_available_resources(self, model, self.get_default_engine_name()).values()]
+
+    def poll_imports(self, sources:List[str], model:str):
+        raise Exception("Azure doesn't support import polling")
+
+    def create_import_stream(self, source:ImportSource, model:str, rate = 1, options: dict|None = None):
+        raise Exception("Azure doesn't support import streams")
+
+    def create_import_snapshot(self, source:ImportSource, model:str, options: dict|None = None):
+        assert isinstance(source, ImportSourceFile), "Azure integration only supports loading from files and URLs right now."
+        import_file(self, model, source, **(options or {}))
+
+    def delete_import(self, import_name: str, model:str, force = False):
+        res = self.exec_raw(model, self.get_default_engine_name(), emit_delete_import(import_name), False)
+        assert_no_problems(res)
+
+    def set_imports_engine_size(self, size: str):
+        raise Exception("Azure doesn't support setting imports engine size")
+
+    def change_imports_status(self, suspend:bool):
+        raise Exception("Azure doesn't support import status changes")
+
+    def get_imports_status(self):
+        return None
+
+    def change_stream_status(self, stream_id: str, model:str, suspend: bool):
+        raise Exception("Azure doesn't support stream status changes")
+
+    def get_import_stream(self, name: str|None, model:str|None):
+        raise Exception("Azure doesn't support get import streams")
+
+    #--------------------------------------------------
+    # Exec
+    #--------------------------------------------------
+
+    def _exec(self, code:str, params:List[Any]|Any|None = None, raw=False, help=True):
+        raise Exception("Azure doesn't support _exec")
+
+    def exec_lqp(self, database: str, engine: str | None, raw_code: bytes, readonly=True, *, inputs: Dict | None = None, nowait_durable=False, headers: Dict | None = None, bypass_index=False, query_timeout_mins: int | None = None):
+        raise Exception("Azure doesn't support exec_lqp")
+
+    def exec_raw(self, database:str, engine:str|None, raw_code:str, readonly=True, *, inputs: Dict | None = None, nowait_durable=False, headers: Dict | None = None, raw_results=True, query_timeout_mins: int | None = None):
+        if query_timeout_mins is not None or self.config.get("query_timeout_mins", None) is not None:
+            config_file_path = getattr(self.config, 'file_path', None)
+            raise AzureUnsupportedQueryTimeoutException(config_file_path=config_file_path)
+        if engine is None:
+            engine = self.get_default_engine_name()
+        try:
+            with debugging.span("transaction") as txn_span:
+                ctx = self._api_ctx()
+                if inputs is None:
+                    inputs = {}
+                with debugging.span("create"):
+                    txn = api.exec_async(ctx, database, engine, raw_code, readonly=readonly, inputs=inputs)
+                txn_id = txn.transaction["id"]
+                txn_span["txn_id"] = txn_id
+                debugging.event("transaction_created", txn_span, txn_id=txn_id)
+
+                # TODO: dedup with SDK
+                rsp = api.TransactionAsyncResponse()
+                txn = api.get_transaction(ctx, txn_id)
+                start_time = time.time()
+
+                def check_done():
+                    with debugging.span("check_status"):
+                        state = api.get_transaction(ctx, txn_id)["state"]
+                        return api.is_txn_term_state(state)
+
+                with debugging.span("wait", txn_id=txn_id):
+                    poll_with_specified_overhead(
+                        check_done,
+                        overhead_rate=0.1,
+                        start_time=start_time,
+                    )
+
+                # TODO: parallelize
+                with debugging.span("fetch"):
+                    rsp.transaction = api.get_transaction(ctx, txn_id)
+                    rsp.metadata = api.get_transaction_metadata(ctx, txn_id)
+                    rsp.problems = api.get_transaction_problems(ctx, txn_id)
+                    with debugging.span("fetch_results"):
+                        rsp.results = api.get_transaction_results(ctx, txn_id)
+
+                return cast(TransactionAsyncResponse, rsp)
+        except HTTPError as err:
+            res = json.loads(err.read().decode())
+            # Grab request id; useful for searching logs
+            request_id = err.headers.get("x-request-id")
+            # RAI API uses a JSON payload in the body to explain why the request failed
+            # This annotates the error with that to make the exception actually useful.
+            if "engine not found" in res.get('message', ''):
+                print("") # the SDK appears to print some stuff before the error message
+                exception = EngineNotFoundException(cast(str, self.config.get('engine', "Unknown")), res.get('message'))
+                raise exception from None
+            raise RAIException("HTTPError", res.get('message', ''), f"details: {res.get('details', '')}; request_id: {request_id}")
+
+    def format_results(self, results, task:m.Task|None=None) -> Tuple[DataFrame, List[Any]]:
+        return result_helpers.format_results(results, task)
+
+    #--------------------------------------------------
+    # Exec format
+    #--------------------------------------------------
+
+    def exec_format(self, database: str, engine: str, raw_code: str, cols:List[str], format:str, inputs: Dict | None = None, readonly: bool = True, nowait_durable=False, skip_invalid_data=False, headers: Dict | None = None, query_timeout_mins: int | None = None) -> Any: # @FIXME: Better type annotation
+        raise Exception("Azure doesn't support alternative formats yet")
+
+    def to_model_type(self, model:dsl.Graph, name: str, source:str):
+        raise Exception("Azure doesn't support import types yet")
+
+    #--------------------------------------------------
+    # Transactions
+    #--------------------------------------------------
+
+    def get_transaction(self, transaction_id):
+        txn = api.get_transaction(self._api_ctx(), transaction_id)
+        if not txn:
+            return None
+        created_on = txn.get("created_on")
+        finished_at = txn.get("finished_at")
+        duration = txn.get("duration")
+        if duration:
+            txn["duration"] = timedelta(milliseconds=duration)
+        elif created_on:
+            txn["duration"] = datetime.now() - datetime.fromtimestamp(created_on / 1000)
+        if created_on:
+            txn["created_on"] = datetime.fromtimestamp(created_on / 1000)
+        if finished_at:
+            txn["finished_at"] = datetime.fromtimestamp(finished_at / 1000)
+        # Remap based on the fields we care about
+        result = {TXN_REPLACE_MAP.get(k, k): v for k, v in txn.items() if k in TXN_FIELDS}
+        return result
+
+    def remap_fields(self, transactions):
+        if not transactions:
+            return []
+        for transaction in transactions:
+            for key in list(transaction.keys()):
+                if key in TXN_REPLACE_MAP:
+                    transaction[TXN_REPLACE_MAP[key]] = transaction.pop(key)
+        return transactions
+
+    def list_transactions(self, **kwargs):
+        TERMINAL_STATES = ["COMPLETED", "ABORTED"]
+        VALID_KEYS = ["id", "state", "engine"]
+
+        state = kwargs.get("state")
+        only_active = kwargs.get("only_active", False)
+        options = {}
+
+        # Azure sdk supports more than just VALID_KEYS as filters but for now we pass through those
+        for k, v in kwargs.items():
+            if k in VALID_KEYS and v is not None:
+                # Only pass state if it is a valid terminal state
+                if k == "state" and v.upper() in TERMINAL_STATES:
+                    options[k] = v.upper()
+                if k != "state":
+                    if k == "engine":
+                        k = "engine_name"
+                    options[k] = v
+        # In Azure we store transactions in cosmos and consul
+        # Cosmos if the state is terminal (COMPLETED or ABORTED) and Consul if the state is not (e.g. "RUNNING")
+        # So we can not filter on active non terminal states via the options passed
+        transactions = api.list_transactions(self._api_ctx(), **options)
+
+        if not transactions:
+            return []
+        # We filter non terminal transactions here
+        if only_active:
+            transactions = [t for t in transactions if t["state"] in ["CREATED", "RUNNING", "PENDING"]]
+        if (isinstance(state, str) and state.upper() not in TERMINAL_STATES):
+            transactions = [t for t in transactions if t["state"] in [state.upper()]]
+        return self.remap_fields(transactions)
+
+    def cancel_transaction(self, transaction_id):
+        return api.cancel_transaction(self._api_ctx(), transaction_id)
+
+    def cancel_pending_transactions(self):
+        # all transactions are executed synchronously against azure
+        pass
+
+    def get_transaction_events(self, transaction_id:str, continuation_token:str):
+        return api._get_resource(
+            self._api_ctx(),
+            f"/transactions/{transaction_id}/events/profiler?continuation_token={continuation_token}",
+        )
+
+    def is_account_flag_set(self, flag: str):
+        raise Exception("Azure doesn't support account flags")
+
+    def is_direct_access_enabled(self) -> bool:
+      raise Exception("Azure doesn't support direct access")
+#--------------------------------------------------
+# Provider
+#--------------------------------------------------
+
+class Provider(ProviderBase):
+
+    def __init__(
+        self,
+        profile: str | None = None,
+        config: Config | None = None,
+        resources: Resources | None = None,
+    ):
+        if resources:
+            self.resources = resources
+        else:
+            self.resources = Resources(profile, config)
+
+#--------------------------------------------------
+# Graph
+#--------------------------------------------------
+
+def Graph(name, *, profile:str|None=None, config:Config, dry_run:bool=False, isolated=True, keep_model:bool=False, format="default"):
+    use_monotype_operators = config.get("compiler.use_monotype_operators", False)
+
+    client = Client(
+        Resources(profile, config),
+        rel.Compiler(config),
+        name,
+        config,
+        dry_run=dry_run,
+        isolated=isolated,
+        keep_model=keep_model,
+    )
+    base_rel = """
+        @inline
+        def make_identity(x..., z):
+            rel_primitive_hash_tuple_uint128(x..., z)
+
+        @inline
+        def pyrel_default({F}, c, k..., v):
+            F(k..., v) or (not F(k..., _) and v = c)
+
+        @inline
+        def pyrel_unwrap(x in UInt128, y): y = x
+
+        @inline
+        def pyrel_dates_period_days(x in Date, y in Date, z in Int):
+            exists((u) | dates_period_days(x, y , u) and u = ^Day[z])
+
+        @inline
+        def pyrel_datetimes_period_milliseconds(x in DateTime, y in DateTime, z in Int):
+            exists((u) | datetimes_period_milliseconds(x, y , u) and u = ^Millisecond[z])
+
+        @inline
+        def pyrel_bool_filter(a, b, {F}, z): { z = if_then_else[F(a, b), boolean_true, boolean_false] }
+
+        @inline
+        def pyrel_strftime(v, fmt, tz in String, s in String):
+            (Date(v) and s = format_date[v, fmt])
+            or (DateTime(v) and s = format_datetime[v, fmt, tz])
+
+        @inline
+        def pyrel_regex_match_all(pattern, string in String, pos in Int, offset in Int, match in String):
+            regex_match_all(pattern, string, offset, match) and offset >= pos
+
+        @inline
+        def pyrel_regex_match(pattern, string in String, pos in Int, offset in Int, match in String):
+            pyrel_regex_match_all(pattern, string, pos, offset, match) and offset = pos
+
+        @inline
+        def pyrel_regex_search(pattern, string in String, pos in Int, offset in Int, match in String):
+            enumerate(pyrel_regex_match_all[pattern, string, pos], 1, offset, match)
+
+        @inline
+        def pyrel_regex_sub(pattern, repl in String, string in String, result in String):
+            string_replace_multiple(string, {(last[regex_match_all[pattern, string]], repl)}, result)
+
+        @inline
+        def pyrel_capture_group(regex in Pattern, string in String, pos in Int, index, match in String):
+            (Integer(index) and capture_group_by_index(regex, string, pos, index, match)) or
+            (String(index) and capture_group_by_name(regex, string, pos, index, match))
+
+        declare __resource
+        declare __compiled_patterns
+    """
+    if use_monotype_operators:
+        base_rel += """
+
+        // use monotyped operators
+        from ::std::monotype import +, -, *, /, <, <=, >, >=
+        """
+    pyrel_base = dsl.build.raw_task(base_rel)
+    debugging.set_source(pyrel_base)
+    client.install("pyrel_base", pyrel_base)
+    return dsl.Graph(client, name, format=format)

relationalai/clients/resources/snowflake/__init__.py

@@ -0,0 +1,20 @@
+"""
+Snowflake resources module.
+"""
+# Import order matters - Resources must be imported first since other classes depend on it
+from .snowflake import Resources, Provider, Graph, SnowflakeClient, APP_NAME, PYREL_ROOT_DB, ExecContext, PrimaryKey, PRINT_TXN_PROGRESS_FLAG
+from .engine_service import EngineType, INTERNAL_ENGINE_SIZES, ENGINE_SIZES_AWS, ENGINE_SIZES_AZURE
+# These imports depend on Resources, so they come after
+from .cli_resources import CLIResources
+from .use_index_resources import UseIndexResources
+from .direct_access_resources import DirectAccessResources
+from .resources_factory import create_resources_instance
+
+__all__ = [
+    'Resources', 'DirectAccessResources', 'Provider', 'Graph', 'SnowflakeClient',
+    'APP_NAME', 'PYREL_ROOT_DB', 'CLIResources', 'UseIndexResources', 'ExecContext', 'EngineType',
+    'INTERNAL_ENGINE_SIZES', 'ENGINE_SIZES_AWS', 'ENGINE_SIZES_AZURE', 'PrimaryKey',
+    'PRINT_TXN_PROGRESS_FLAG', 'create_resources_instance',
+]
+
+

relationalai/clients/resources/snowflake/cli_resources.py

@@ -0,0 +1,98 @@
+"""
+CLI Resources - Resources class for CLI operations.
+Re-raises exceptions without transformation and provides CLI utility methods.
+"""
+from __future__ import annotations
+import json
+from typing import Any
+
+from ....tools.constants import RAI_APP_NAME
+
+# 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, ExecContext
+from .error_handlers import AppMissingErrorHandler, AppFunctionMissingErrorHandler, ServiceNotStartedErrorHandler
+
+
+class CLIResources(Resources):
+    """
+    Resources class for CLI operations.
+    Re-raises exceptions without transformation and provides CLI utility methods.
+    """
+
+    def _handle_standard_exec_errors(self, e: Exception, ctx: ExecContext) -> Any | None:
+        """
+        For CLI resources, keep exceptions raw except for few specific cases.
+        """
+        message = str(e).lower()
+        for handler in (
+            ServiceNotStartedErrorHandler(),
+            AppMissingErrorHandler(),
+            AppFunctionMissingErrorHandler(),
+        ):
+            if handler.matches(e, message, ctx, self):
+                handler.handle(e, ctx, self)
+        raise e
+
+    def list_warehouses(self):
+        """List all warehouses in the Snowflake account."""
+        results = self._exec("SHOW WAREHOUSES")
+        if not results:
+            return []
+        return [{"name": name}
+                for (name, *rest) in results]
+
+    def list_compute_pools(self):
+        """List all compute pools in the Snowflake account."""
+        results = self._exec("SHOW COMPUTE POOLS")
+        if not results:
+            return []
+        return [{"name": name, "status": status, "min_nodes": min_nodes, "max_nodes": max_nodes, "instance_family": instance_family}
+                for (name, status, min_nodes, max_nodes, instance_family, *rest) in results]
+
+    def list_roles(self):
+        """List all available roles in the Snowflake account."""
+        results = self._exec("SELECT CURRENT_AVAILABLE_ROLES()")
+        if not results:
+            return []
+        # the response is a single row with a single column containing
+        # a stringified JSON array of role names:
+        row = results[0]
+        if not row:
+            return []
+        return [{"name": name} for name in json.loads(row[0])]
+
+    def list_apps(self):
+        """List all applications in the Snowflake account."""
+        all_apps = self._exec(f"SHOW APPLICATIONS LIKE '{RAI_APP_NAME}'")
+        if not all_apps:
+            all_apps = self._exec("SHOW APPLICATIONS")
+            if not all_apps:
+                return []
+        return [{"name": name}
+                for (time, name, *rest) in all_apps]
+
+    def list_databases(self):
+        """List all databases in the Snowflake account."""
+        results = self._exec("SHOW DATABASES")
+        if not results:
+            return []
+        return [{"name": name}
+                for (time, name, *rest) in results]
+
+    def list_sf_schemas(self, database: str):
+        """List all schemas in a given database."""
+        results = self._exec(f"SHOW SCHEMAS IN {database}")
+        if not results:
+            return []
+        return [{"name": name}
+                for (time, name, *rest) in results]
+
+    def list_tables(self, database: str, schema: str):
+        """List all tables and views in a given schema."""
+        results = self._exec(f"SHOW OBJECTS IN {database}.{schema}")
+        items = []
+        if results:
+            for (time, name, db_name, schema_name, kind, *rest) in results:
+                items.append({"name": name, "kind": kind.lower()})
+        return items