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

relationalai/clients/local.py

@@ -0,0 +1,571 @@
+from __future__ import annotations
+
+import base64
+import json
+from urllib.parse import quote, urlencode
+import pyarrow as pa
+import requests
+from email import message_from_bytes, policy
+from email.message import EmailMessage
+
+from dataclasses import dataclass
+from typing import Any, Dict, List, Iterable, Literal, Optional, Tuple, Union
+from requests.adapters import HTTPAdapter
+from urllib3.util.retry import Retry
+
+from .client import ResourcesBase, ProviderBase
+from .config import Config
+from .types import TransactionAsyncResponse
+from .util import get_pyrel_version
+from ..errors import ResponseStatusException
+from .. import debugging
+
+@dataclass
+class Endpoint:
+    method: str
+    endpoint: str
+
+TXN_ABORT_REASON_TIMEOUT = "QUERY_TIMEOUT_EXCEEDED"
+
+class LocalClient:
+
+    def __init__(self, host: str, port: int):
+        self.service_endpoint = f"http://{host}:{port}"
+        self.endpoints: Dict[str, Endpoint] = {
+            "create_txn": Endpoint(method="POST", endpoint="/transactions"), # API for creating a (query) transaction
+            "create_db": Endpoint(method="POST", endpoint="/v2/<local_org>/<local_account>/databases"), # API for creating a database
+        }
+        self.http_session = self._create_retry_session()
+
+    def _create_retry_session(self) -> requests.Session:
+        http_session = requests.Session()
+        retries = Retry(
+            total=3,
+            backoff_factor=0.3,
+            status_forcelist=[500, 502, 503, 504],
+            allowed_methods=frozenset({"GET", "POST", "PUT", "DELETE"}),
+            raise_on_status=False
+        )
+        adapter = HTTPAdapter(max_retries=retries)
+        http_session.mount("http://", adapter)
+        http_session.mount("https://", adapter)
+        http_session.headers.update({
+            "Connection": "keep-alive",
+            "Accept": "application/json",
+            "user-agent": get_pyrel_version(None)
+        })
+        return http_session
+
+    def request(
+        self,
+        endpoint: str,
+        payload: Dict[str, Any] | None = None,
+        headers: Dict[str, str] | None = None,
+        path_params: Dict[str, str] | None = None,
+        query_params: Dict[str, str] | None = None,
+    ) -> requests.Response:
+        with debugging.span("local_request"):
+            url, method = self._prepare_url(endpoint, path_params, query_params)
+            request_headers = self._prepare_headers(headers)
+            return self.http_session.request(method, url, json=payload, headers=request_headers)
+
+    def _prepare_url(self, endpoint: str, path_params: Dict[str, str] | None = None, query_params: Dict[str, str] | None = None) -> Tuple[str, str]:
+        try:
+            ep = self.endpoints[endpoint]
+        except KeyError:
+            raise ValueError(f"Invalid endpoint: {endpoint}. Available endpoints: {list(self.endpoints.keys())}")
+        url = f"{self.service_endpoint}{ep.endpoint}"
+        if path_params:
+            escaped_path_params = {k: quote(v, safe='') for k, v in path_params.items()}
+            url = url.format(**escaped_path_params)
+        if query_params:
+            url += '?' + urlencode(query_params)
+        return url, ep.method
+
+    def _prepare_headers(self, headers: Dict[str, str] | None) -> Dict[str, str]:
+        request_headers = {}
+        if headers:
+            request_headers.update(headers)
+        return request_headers
+
+class LocalResources(ResourcesBase):
+
+    def __init__(
+        self,
+        profile: str | None = None,
+        config: Config | None = None,
+        connection: Any | None = None,
+        dry_run: bool = False,
+        reset_session: bool = False,
+        generation: Any | None = None,
+        language: str = "lqp",
+    ):
+        super().__init__(config=config)
+        self.language = language
+
+        host = self.config.get("host", None) or "localhost"
+        port = int(self.config.get("port", None) or 8010)
+        self.client = LocalClient(host, port)
+
+        # Keep track of whether we already created the local database to run on.
+        self._db_created = False
+
+    def reset(self):
+        raise NotImplementedError("reset not supported in local mode")
+    
+    #--------------------------------------------------
+    # Check direct access is enabled (0 implemented)
+    #--------------------------------------------------
+
+    def is_direct_access_enabled(self) -> bool:
+        raise NotImplementedError("is_direct_access_enabled not supported in local mode")
+
+    #--------------------------------------------------
+    # Snowflake Account Flags (0 implemented)
+    #--------------------------------------------------
+
+    def is_account_flag_set(self, flag: str) -> bool:
+        raise NotImplementedError("is_account_flag_set not supported in local mode")
+
+    #--------------------------------------------------
+    # Databases (0 implemented)
+    #--------------------------------------------------
+
+    def get_database(self, database: str):
+        raise NotImplementedError("get_database not supported in local mode")
+
+    def get_installed_packages(self, database: str):
+        raise NotImplementedError("get_installed_packages not supported in local mode")
+
+    def _create_database(self, name:str):
+        with debugging.span("create_database"):
+            response = self.client.request(
+                endpoint="create_db",
+                payload = {
+                    "language": "lqp",
+                    "readonly": False,
+                    "database_id": name,
+                    "transaction_id": "mock_txn_id:create-db",
+                }
+            )
+            if response.status_code != 200:
+                raise ResponseStatusException("Failed to create database in local mode.", response)
+
+            parsed_response = self._parse_txn_response(response)
+
+            state = parsed_response["state"]
+            if state != "COMPLETED":
+                # NOTE for now we just ignore it when the database creation transaction
+                # aborts. In most cases where this happens, it is because the database
+                # already exists. If it doesn't, we'll fail later on anyways.
+                pass
+
+            return parsed_response
+
+    #--------------------------------------------------
+    # Engines (0 implemented)
+    #--------------------------------------------------
+
+    def get_version(self):
+        raise NotImplementedError("get_version not supported in local mode")
+
+    def get_engine_sizes(self, cloud_provider: str | None = None):
+        raise NotImplementedError("get_engine_sizes not supported in local mode")
+
+    def list_engines(self, state: str | None = None):
+        raise NotImplementedError("list_engines not supported in local mode")
+
+    def get_engine(self, name: str):
+        raise NotImplementedError("get_engine not supported in local mode")
+
+    def get_cloud_provider(self) -> str:
+        raise NotImplementedError("get_cloud_provider not supported in local mode")
+
+    def is_valid_engine_state(self, name: str) -> bool:
+        raise NotImplementedError("is_valid_engine_state not supported in local mode")
+
+    def create_engine(self, name: str, size:str|None=None, auto_suspend_mins:int|None=None):
+        raise NotImplementedError("create_engine not supported in local mode")
+
+    def delete_engine(self, name: str, force: bool = False):
+        raise NotImplementedError("delete_engine not supported in local mode")
+
+    def suspend_engine(self, name: str):
+        raise NotImplementedError("suspend_engine not supported in local mode")
+
+    def resume_engine(self, name: str, headers: Dict | None = None):
+        raise NotImplementedError("resume_engine not supported in local mode")
+
+    def resume_engine_async(self, name: str):
+        raise NotImplementedError("resume_engine_async not supported in local mode")
+
+    def alter_engine_pool(self, size: str | None = None, mins: int | None = None, maxs: int | None = None):
+        raise NotImplementedError("alter_engine_pool not supported in local mode")
+
+    def auto_create_engine_async(self, name: str | None = None) -> str:
+        raise NotImplementedError("auto_create_engine_async not supported in local mode")
+
+    #--------------------------------------------------
+    # Graphs (0 implemented)
+    #--------------------------------------------------
+
+    def list_graphs(self):
+        raise NotImplementedError("list_graphs not supported in local mode")
+
+    def get_graph(self, name: str):
+        raise NotImplementedError("get_graph not supported in local mode")
+
+    def create_graph(self, name: str):
+        raise NotImplementedError("create_graph not supported in local mode")
+
+    def delete_graph(self, name: str, force=False, language: str = "lqp"):
+        # NOTE There is no way to delete a local graph at the moment, which is documented in
+        # the README. This function will get called in regular usage, so for now we just
+        # silently do nothing.
+        return None
+
+    def clone_graph(self, target_name: str, source_name: str, nowait_durable: bool = True):
+        raise NotImplementedError("clone_graph not supported in local mode")
+
+    def maybe_poll_use_index(
+        self,
+        app_name: str,
+        sources: Iterable[str],
+        model: str,
+        engine_name: str,
+        engine_size: str | None = None,
+        program_span_id: str | None = None,
+        headers: Dict | None = None,
+    ):
+        # Snowflake sources are not supported in local mode.
+        if len(list(sources)) != 0:
+            raise NotImplementedError("Snowflake sources are not supported in local mode")
+
+        # Nothing to do for local mode if the database was already created.
+        if self._db_created:
+            return
+
+        self._create_database(model)
+        self._db_created = True
+
+    #--------------------------------------------------
+    # Models (0 implemented)
+    #--------------------------------------------------
+
+    def list_models(self, database: str, engine: str):
+        raise NotImplementedError("list_models not supported in local mode")
+
+    def create_models(self, database: str, engine: str, models):
+        raise NotImplementedError("create_models not supported in local mode")
+
+    def delete_model(self, database: str, engine: str, name: str):
+        raise NotImplementedError("delete_model not supported in local mode")
+
+    def create_models_code(self, models):
+        raise NotImplementedError("create_models_code not supported in local mode")
+
+    #--------------------------------------------------
+    # Exports (0 implemented)
+    #--------------------------------------------------
+
+    def list_exports(self, database: str, engine: str):
+        raise NotImplementedError("list_exports not supported in local mode")
+
+    def create_export(self, params):
+        raise NotImplementedError("create_export not supported in local mode")
+
+    def create_export_table(self, database: str, engine: str, table: str, relation: str, columns: Dict[str, str], code: str, refresh: str | None = None):
+        raise NotImplementedError("create_export_table not supported in local mode")
+
+    def delete_export(self, database: str, engine: str, name: str):
+        raise NotImplementedError("delete_export not supported in local mode")
+
+    #--------------------------------------------------
+    # Imports (0 implemented)
+    #--------------------------------------------------
+
+    def get_import_stream(self, name: str | None, model: str | None):
+        raise NotImplementedError("get_import_stream not supported in local mode")
+
+    def change_stream_status(self, stream_id: str, model: str, suspend: bool):
+        raise NotImplementedError("change_stream_status not supported in local mode")
+
+    def change_imports_status(self, suspend: bool):
+        raise NotImplementedError("change_imports_status not supported in local mode")
+
+    def get_imports_status(self):
+        raise NotImplementedError("get_imports_status not supported in local mode")
+
+    def set_imports_engine_size(self, size: str):
+        raise NotImplementedError("set_imports_engine_size not supported in local mode")
+
+    def list_imports(self, id: str | None = None, name: str | None = None, model: str | None = None, status: str | None = None, creator: str | None = None):
+        raise NotImplementedError("list_imports not supported in local mode")
+
+    def poll_imports(self, sources, model: str):
+        raise NotImplementedError("poll_imports not supported in local mode")
+
+    def create_import_stream(self, source, model: str, rate: int | None = None, options: dict | None = None):
+        raise NotImplementedError("create_import_stream not supported in local mode")
+
+    def create_import_snapshot(self, source, model: str, options: dict | None):
+        raise NotImplementedError("create_import_snapshot not supported in local mode")
+
+    def delete_import(self, import_name: str, model: str, force=False):
+        raise NotImplementedError("delete_import not supported in local mode")
+
+    #--------------------------------------------------
+    # Exec Async
+    #--------------------------------------------------
+    
+    def _parse_multipart_response(self, response: requests.Response) -> Dict[str, Any]:
+        response_map = {}
+        response_map['results'] = {}
+
+        content_type = response.headers.get("Content-Type")
+        if not content_type:
+            raise ValueError("Missing Content-Type header")
+
+        # Parse using Python's email module
+        # Construct a proper MIME message with headers
+        mime_content = f"Content-Type: {content_type}\r\n\r\n".encode('utf-8') + response.content
+        msg: EmailMessage = message_from_bytes(mime_content, policy=policy.default)  # type: ignore[assignment]
+
+        if not msg.is_multipart():
+            raise ValueError("Response is not multipart")
+
+        for part in msg.iter_parts():
+            part_content_type = part.get_content_type() or 'application/octet-stream'
+            part_content: bytes = part.get_payload(decode=True)  # type: ignore[assignment]
+
+            # Extract disposition info
+            disposition = part.get('Content-Disposition', '')
+            part_name = None
+            part_filename = None
+            if disposition:
+                params = {}
+                for param_str in disposition.split(';')[1:]:
+                    if '=' in param_str:
+                        key, val = param_str.strip().split('=', 1)
+                        params[key] = val.strip('"')
+                part_name = params.get('name')
+                part_filename = params.get('filename')
+
+            if 'json' in part_content_type.lower() and not part_filename:
+                # JSON part (no filename)
+                try:
+                    response_map[part_name] = json.loads(part_content.decode('utf-8'))
+                except json.JSONDecodeError as e:
+                    raise ValueError(f"Failed to decode JSON part: {e}")
+
+            elif part_filename:
+                # File/attachment part
+                response_map['results'][part_filename] = {
+                    'name': part_name,
+                    'filename': part_filename,
+                    'content': part_content,
+                    'content_type': part_content_type
+                }
+
+        return response_map
+
+    def _parse_txn_response(self, get_txn_response: requests.Response) -> Dict[str, Any]:
+        with debugging.span("parse_txn_response"):
+            response_content_type = get_txn_response.headers.get("Content-Type", "")
+            if "application/json" in response_content_type:
+                parsed_response = get_txn_response.json()
+            elif "multipart/form-data" in response_content_type:
+                parsed_response = self._parse_multipart_response(get_txn_response)
+                state = parsed_response.get('transaction', {}).get('state', None)
+                parsed_response['state'] = state
+                txn_id = parsed_response.get('transaction', {}).get('id', None)
+                parsed_response['id'] = txn_id
+            else:
+                raise Exception(f"Unsupported Response Content Type:, {response_content_type}")
+            return parsed_response
+
+    def _create_transaction(
+        self,
+        target_endpoint: str,
+        payload: Dict[str, Any] | None = None,
+        headers: Dict[str, str] | None = None,
+        path_params: Dict[str, str] | None = None,
+        query_params: Dict[str, str] | None = None
+    ) -> Dict[str, Any]:
+        with debugging.span("transaction"):
+            response = self.client.request(
+                target_endpoint,
+                path_params=path_params,
+                query_params=query_params,
+                payload=payload,
+                headers=headers,
+            )
+            assert response, "No results from create transaction request."
+            if response.status_code not in [200, 201]:
+                raise ResponseStatusException("Failed to create transaction.", response)
+
+            return self._parse_txn_response(response)
+
+    def _convert_txn_response(self, parsed_response: Dict[str, Any]) -> TransactionAsyncResponse:
+        results = []
+        for (file_name, file_info) in parsed_response.get("results", {}).items():
+            if not file_name.endswith(".arrow"):
+                continue
+            with pa.ipc.open_stream(file_info["content"]) as reader:
+                schema = reader.schema
+                batches = [batch for batch in reader]
+                table = pa.Table.from_batches(batches=batches, schema=schema)
+                results.append({"relationId": file_info["name"], "table": table})
+
+        rsp = TransactionAsyncResponse()
+        rsp.transaction = {
+            "id": parsed_response["id"],
+            "state": parsed_response["state"]
+        }
+        rsp.results = results
+        rsp.metadata = parsed_response.get("metadata", [])
+        rsp.problems = parsed_response.get("problems", [])
+
+        return rsp
+
+    def _exec_async_v2(
+        self,
+        database: str,
+        engine: Union[str, None],
+        raw_code: str,
+        inputs: Dict | None = None,
+        readonly=True,
+        nowait_durable=False,
+        headers: Dict[str, str] | None = None,
+        bypass_index=False,
+        language: str = "lqp",
+        query_timeout_mins: int | None = None,
+    ):
+        payload = {
+            "engine_name": engine,
+            "dbname": database,
+            "query": raw_code,
+            "language": "lqp",
+            "readonly": readonly,
+        }
+
+        parsed_response = self._create_transaction(
+            target_endpoint="create_txn",
+            payload=payload,
+            headers=headers
+        )
+
+        state = parsed_response["state"]
+        if state not in ["COMPLETED", "ABORTED"]:
+            # Local mode only supports synchronous transactions for now.
+            txn_id = parsed_response["id"]
+            raise Exception(f"Transaction {txn_id} did not complete or abort synchronously (state = {state}).")
+
+        return self._convert_txn_response(parsed_response)
+
+    #--------------------------------------------------
+    # Exec (1 implemented, others are stubs)
+    #--------------------------------------------------
+
+    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,
+    ):
+        raw_code_b64 = base64.b64encode(raw_code).decode("utf-8")
+        return self._exec_async_v2(
+            database, engine, raw_code_b64, inputs, readonly, nowait_durable,
+            headers=headers, bypass_index=bypass_index, language='lqp',
+            query_timeout_mins=query_timeout_mins,
+        )
+
+    def exec_raw(
+        self,
+        database: str,
+        engine: str | None,
+        raw_code: str,
+        readonly: bool = True,
+        *,
+        inputs: Dict | None = None,
+        nowait_durable: bool = False,
+        headers: Dict | None = None,
+        query_timeout_mins: Optional[int] = None,
+    ):
+        raise NotImplementedError("exec_raw not supported in local mode - use exec_lqp instead")
+
+    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:
+        raise NotImplementedError("exec_format not supported in local mode")
+
+    def format_results(self, results, task=None):
+        raise NotImplementedError("format_results not supported in local mode")
+
+    def to_model_type(self, model, name: str, source: str):
+        raise NotImplementedError("to_model_type not supported in local mode")
+
+    def _exec(self, code: str, params=None, raw=False, help=True):
+        raise NotImplementedError("_exec not supported in local mode")
+
+    #--------------------------------------------------
+    # Transactions (0 implemented)
+    #--------------------------------------------------
+
+    def get_transaction(self, transaction_id: str):
+        raise NotImplementedError("get_transaction not supported in local mode")
+
+    def list_transactions(self, *, limit: int, only_active=False, **kwargs):
+        raise NotImplementedError("list_transactions not supported in local mode")
+
+    def cancel_transaction(self, transaction_id: str):
+        raise NotImplementedError("cancel_transaction not supported in local mode")
+
+    def cancel_pending_transactions(self):
+        raise NotImplementedError("cancel_pending_transactions not supported in local mode")
+
+    def get_transaction_events(self, transaction_id: str, continuation_token: str):
+        raise NotImplementedError("get_transaction_events not supported in local mode")
+
+class LocalProvider(ProviderBase):
+
+    def __init__(
+        self,
+        host: str = "localhost",
+        port: int = 8010,
+        profile: str | None = None,
+        config: Config | None = None,
+        resources: LocalResources | None = None,
+    ):
+        if resources:
+            self.resources = resources
+        else:
+            self.resources = LocalResources(
+                profile=profile,
+                config=config,
+                connection=None,
+                dry_run=False,
+                reset_session=False,
+                generation=None,
+                language="lqp",
+            )
+        self.config = self.resources.config
+
+    # A mock sql function that simulates SQL execution for the LocalProvider. This is only
+    # here to make the test suite pass.
+    def sql(self, query: str, params: List[Any] = [], format: Literal["list", "pandas", "polars", "lazy"] = "list"):
+        if format == "list":
+            # Used by "test_snapshot_abstract.py"
+            import re
+            pattern = r"\s*SELECT EXISTS\(\s*SELECT 1\s*FROM (\w+)\.INFORMATION_SCHEMA\.TABLES\s*WHERE TABLE_SCHEMA = '(\w+)'\s*AND TABLE_NAME = '(\w+)'\s*\) AS TABLE_EXISTS;\s*"
+            query_cleaned = query.replace("\n", " ").replace("  ", " ")
+            match_obj = re.match(pattern, query_cleaned, re.DOTALL)
+            if match_obj:
+                return [{"TABLE_EXISTS": False}]
+        else:
+            raise NotImplementedError("Only specific SQL queries are defined in local mode.")

relationalai/clients/profile_polling.py

@@ -0,0 +1,73 @@
+from dataclasses import dataclass
+import logging
+from queue import Queue
+import threading
+import time
+import traceback
+from typing import Union
+
+from relationalai import debugging
+from relationalai.clients.client import ResourcesBase
+
+PROFILE_POLL_SLEEP_S = 1.5
+
+@dataclass
+class ProfilePollingItem:
+    txn_id: str
+    continuation_token: str
+
+# using None as a sentinel value in the queue to indicate
+# that it's shutting down
+SHUTTING_DOWN = None
+
+class ProfilePollerHandler(logging.Handler):
+    """
+    A handler which sees transaction start events come over the log stream,
+    starts polling for profile events for each transaction, and puts them into
+    the same stream.
+    """
+    
+    def __init__(self, resources: ResourcesBase):
+        super().__init__()
+        self.resources = resources
+        self.queue = Queue[Union[ProfilePollingItem,None]]()
+
+        # daemon thread so it doesn't keep the process alive forever
+        self.thread = threading.Thread(target=self._polling_loop, daemon=True)
+        self.thread.start()
+    
+    def emit(self, record):
+        if (isinstance(record.msg, dict) and
+            record.msg["event"] == "span_start"):
+            span = record.msg["span"]
+            if span.type == "wait" and span.parent.type == "transaction":
+                txn_id = record.msg["span"].attrs["txn_id"]
+                self.queue.put(ProfilePollingItem(txn_id, ""))
+    
+    def _polling_loop(self):
+        while True:
+            try:
+                item = self.queue.get()
+                if item is SHUTTING_DOWN:
+                    return
+                
+                assert item is not None
+                resp = self.resources.get_transaction_events(item.txn_id, item.continuation_token)
+                    
+                debugging.event('profile_events', txn_id=item.txn_id, profile_events=resp['events'])
+                
+                continuation_token = resp['continuation_token']
+                if continuation_token != '':
+                    self.queue.put(ProfilePollingItem(item.txn_id, continuation_token))
+            except Exception as e:
+                debugging.error(e)
+                print('Error polling profile events:', e)
+                traceback.print_exc()
+            
+            # If multiple txns are running concurrently, a given one of them may be polled
+            # at a longer interval than PROFILE_POLL_SLEEP_S, but that's ok.
+            time.sleep(PROFILE_POLL_SLEEP_S)
+    
+    def shut_down(self):
+        self.queue.put(SHUTTING_DOWN)
+        self.thread.join()

relationalai/clients/resources/__init__.py

@@ -0,0 +1,8 @@
+"""
+Resources module for RelationalAI clients.
+"""
+
+# This package contains resource implementations for different platforms:
+# - snowflake: Snowflake-specific resources
+# - azure: Azure-specific resources (optional, may not be available in all environments)
+