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

v0/relationalai/clients/local.py

@@ -1,571 +0,0 @@
-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 v0.relationalai.clients.client import ResourcesBase, ProviderBase
-from v0.relationalai.clients.config import Config
-from v0.relationalai.clients.types import TransactionAsyncResponse
-from v0.relationalai.clients.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):
-        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.")

v0/relationalai/clients/profile_polling.py

@@ -1,73 +0,0 @@
-from dataclasses import dataclass
-import logging
-from queue import Queue
-import threading
-import time
-import traceback
-from typing import Union
-
-from v0.relationalai import debugging
-from v0.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()