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

relationalai/clients/direct_access_client.py

@@ -0,0 +1,118 @@
+from __future__ import annotations
+
+import requests
+from dataclasses import dataclass
+from urllib.parse import urlencode, quote
+from requests.adapters import HTTPAdapter
+from urllib3.util.retry import Retry
+from typing import Any, Dict, Optional, Tuple
+
+from relationalai.auth.token_handler import TokenHandler
+from relationalai.clients.config import Config
+from relationalai.clients.util import get_pyrel_version
+from relationalai import debugging
+from relationalai.tools.constants import Generation
+from relationalai.environments import runtime_env, SnowbookEnvironment
+
+@dataclass
+class Endpoint:
+    method: str
+    endpoint: str
+
+class DirectAccessClient:
+    """
+    DirectAccessClient is a client for direct service access without service function calls.
+    """
+
+    def __init__(self, config: Config, token_handler: TokenHandler, service_endpoint: str, generation: Optional[Generation] = None):
+        self._config: Config = config
+        self._token_handler: TokenHandler = token_handler
+        self.service_endpoint: str = service_endpoint
+        self.generation: Optional[Generation] = generation
+        self._is_snowflake_notebook = isinstance(runtime_env, SnowbookEnvironment)
+        self.endpoints: Dict[str, Endpoint] = {
+            "create_txn": Endpoint(method="POST", endpoint="/v1alpha1/transactions"),
+            "get_txn": Endpoint(method="GET", endpoint="/v1alpha1/transactions/{txn_id}"),
+            "get_txn_artifacts": Endpoint(method="GET", endpoint="/v1alpha1/transactions/{txn_id}/artifacts"),
+            "get_txn_problems": Endpoint(method="GET", endpoint="/v1alpha1/transactions/{txn_id}/problems"),
+            "get_txn_events": Endpoint(method="GET", endpoint="/v1alpha1/transactions/{txn_id}/events/{stream_name}"),
+            "get_package_versions": Endpoint(method="GET", endpoint="/v1alpha1/databases/{db_name}/package_versions"),
+            "get_model_package_versions": Endpoint(method="POST", endpoint="/v1alpha1/models/get_package_versions"),
+            "create_db": Endpoint(method="POST", endpoint="/v1alpha1/databases"),
+            "get_db": Endpoint(method="GET", endpoint="/v1alpha1/databases"),
+            "delete_db": Endpoint(method="DELETE", endpoint="/v1alpha1/databases/{db_name}"),
+            "release_index": Endpoint(method="POST", endpoint="/v1alpha1/index/release"),
+            "list_engines": Endpoint(method="GET", endpoint="/v1alpha1/engines"),
+            "get_engine": Endpoint(method="GET", endpoint="/v1alpha1/engines/{engine_type}/{engine_name}"),
+            "create_engine": Endpoint(method="POST", endpoint="/v1alpha1/engines/{engine_type}"),
+            "delete_engine": Endpoint(method="DELETE", endpoint="/v1alpha1/engines/{engine_type}/{engine_name}"),
+            "suspend_engine": Endpoint(method="POST", endpoint="/v1alpha1/engines/{engine_type}/{engine_name}/suspend"),
+            "resume_engine": Endpoint(method="POST", endpoint="/v1alpha1/engines/{engine_type}/{engine_name}/resume_async"),
+            "prepare_index": Endpoint(method="POST", endpoint="/v1alpha1/index/prepare"),
+            "get_job": Endpoint(method="GET", endpoint="/v1alpha1/jobs/{job_type}/{job_id}"),
+            "list_jobs": Endpoint(method="GET", endpoint="/v1alpha1/jobs"),
+            "get_job_events": Endpoint(method="GET", endpoint="/v1alpha1/jobs/{job_type}/{job_id}/events/{stream_name}"),
+            "create_job": Endpoint(method="POST", endpoint="/v1alpha1/jobs"),
+            "cancel_job": Endpoint(method="POST", endpoint="/v1alpha1/jobs/{job_type}/{job_id}/cancel"),
+        }
+        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", "PATCH", "OPTIONS", "HEAD"}),
+            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"})
+        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:
+        """
+        Send a request to the service endpoint.
+        """
+        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)
+        # Authorization tokens are not needed in a snowflake notebook environment
+        if not self._is_snowflake_notebook:
+            request_headers["Authorization"] = f'Snowflake Token="{self._token_handler.get_ingress_token(self.service_endpoint)}"'
+            # needed for oauth, does no harm for other authentication methods
+            request_headers["X-SF-SPCS-Authentication-Method"] = 'OAUTH'
+            request_headers["Content-Type"] = 'application/x-www-form-urlencoded'
+        request_headers["Accept"] = "application/json"
+
+        request_headers["user-agent"] = get_pyrel_version(self.generation)
+        request_headers["pyrel_program_id"] = debugging.get_program_span_id() or ""
+
+        return debugging.add_current_propagation_headers(request_headers)

relationalai/clients/exec_txn_poller.py

@@ -0,0 +1,91 @@
+from __future__ import annotations
+
+import time
+from typing import Dict, Optional, TYPE_CHECKING
+
+from relationalai import debugging
+from relationalai.clients.util import poll_with_specified_overhead
+from relationalai.tools.cli_controls import create_progress
+from relationalai.util.format import format_duration
+
+if TYPE_CHECKING:
+    from relationalai.clients.resources.snowflake import Resources
+
+# Polling behavior constants
+POLL_OVERHEAD_RATE = 0.1  # Overhead rate for exponential backoff
+
+# Text color constants
+GREEN_COLOR = '\033[92m'
+GRAY_COLOR = '\033[90m'
+ENDC = '\033[0m'
+
+
+class ExecTxnPoller:
+    """
+    Encapsulates the polling logic for exec_async transaction completion.
+    """
+
+    def __init__(
+        self,
+        resource: "Resources",
+        txn_id: str,
+        headers: Optional[Dict] = None,
+        txn_start_time: Optional[float] = None,
+    ):
+        self.res = resource
+        self.txn_id = txn_id
+        self.headers = headers or {}
+        self.txn_start_time = txn_start_time or time.time()
+
+    def poll(self) -> bool:
+        """
+        Poll for transaction completion with interactive progress display.
+
+        Returns:
+            True if transaction completed successfully, False otherwise
+        """
+
+        # Don't show duration summary - we handle our own completion message
+        with create_progress(
+            description="Evaluating Query...",
+            success_message="",  # We'll handle this ourselves
+            leading_newline=False,
+            trailing_newline=False,
+            show_duration_summary=False,
+        ) as progress:
+            def check_status() -> bool:
+                """Check if transaction is complete."""
+                elapsed = time.time() - self.txn_start_time
+                # Update the main status with elapsed time
+                progress.update_main_status(
+                    query_progress_message(self.txn_id, elapsed)
+                )
+                return self.res._check_exec_async_status(self.txn_id, headers=self.headers)
+
+            with debugging.span("wait", txn_id=self.txn_id):
+                poll_with_specified_overhead(check_status, overhead_rate=POLL_OVERHEAD_RATE)
+
+            # Calculate final duration
+            total_duration = time.time() - self.txn_start_time
+
+            # Update to success message with duration
+            progress.update_main_status(
+                query_complete_message(self.txn_id, total_duration)
+            )
+
+        return True
+
+def query_progress_message(id: str, duration: float) -> str:
+    return (
+        # Print with whitespace to align with the end of the transaction ID
+        f"Evaluating Query... {format_duration(duration):>18}\n" +
+        f"{GRAY_COLOR}Query: {id}{ENDC}"
+    )
+
+def query_complete_message(id: str, duration: float, status_header: bool = False) -> str:
+    return (
+        (f"{GREEN_COLOR}✅ " if status_header else "") +
+        # Print with whitespace to align with the end of the transaction ID
+        f"Query Complete: {format_duration(duration):>24}\n" +
+        f"{GRAY_COLOR}Query: {id}{ENDC}"
+    )

relationalai/clients/hash_util.py

@@ -0,0 +1,31 @@
+import hashlib
+
+def hash_string(s: str, digits=32):
+    return hashlib.md5(s.encode()).hexdigest()[:digits]
+
+def database_name_from_sproc_name(sproc_name: str) -> str:
+    # Import here to avoid circular import issues
+    from relationalai.clients.util import IdentityParser, ParseError
+    from relationalai.rel_utils import sanitize_identifier
+
+    err_text = f"Invalid stored procedure name: {sproc_name}"
+    try:
+        parser = IdentityParser(sproc_name, require_all_parts=True)
+        if not (parser.identity and parser.db and parser.schema and parser.entity):
+            raise ValueError(err_text)
+
+        # Remove all special characters from the name and replace with '_'
+        sanitized_name = sanitize_identifier(parser.identity)
+        hashed_sanitized_name = hash_string(sanitized_name)
+
+        # If the name is less than 32 characters, append the hash of the name
+        if len(sanitized_name) < 32:
+            return f"{sanitized_name}_{hashed_sanitized_name}"
+
+        # If the name is more than 32 characters, truncate it to 31 characters and sanitize each part
+        truncated_database = sanitize_identifier(parser.db[:10])
+        truncated_schema = sanitize_identifier(parser.schema[:10])
+        truncated_proc = sanitize_identifier(parser.entity[:(31 - len(truncated_database) - len(truncated_schema))])
+        return f"{truncated_database}_{truncated_schema}_{truncated_proc}_{hashed_sanitized_name}"
+    except ParseError:
+        raise ValueError(err_text)