UncountablePythonSDK 0.0.8__py3-none-any.whl → 0.0.92__py3-none-any.whl

uncountable/core/client.py

@@ -1,17 +1,32 @@
-from dataclasses import dataclass
-from enum import StrEnum
-import json
-from urllib.parse import urljoin
 import base64
+import re
 import typing
+from dataclasses import dataclass
+from datetime import datetime, timedelta
+from enum import StrEnum
+from io import BytesIO
+from urllib.parse import unquote, urljoin
+from uuid import uuid4
+
 import requests
-from uncountable.types.base import JsonValue
-from uncountable.types.client_base import APIRequest, ClientMethods
+import simplejson as json
+from opentelemetry.sdk.resources import Attributes
+from requests.exceptions import JSONDecodeError
+
 from pkgs.argument_parser import CachedParser
-from pkgs.serialization_util import serialize_for_api
+from pkgs.serialization_util import JsonValue, serialize_for_api
+from uncountable.core.environment import get_version
+from uncountable.integration.telemetry import Logger, push_scope_optional
+from uncountable.types import download_file_t
+from uncountable.types.client_base import APIRequest, ClientMethods
+from uncountable.types.client_config import ClientConfigOptions
 
+from .file_upload import FileUpload, FileUploader, UploadedFile
+from .types import AuthDetailsAll, AuthDetailsApiKey, AuthDetailsOAuth
 
 DT = typing.TypeVar("DT")
+UNC_REQUEST_ID_HEADER = "X-UNC-REQUEST-ID"
+UNC_SDK_VERSION_HEADER = "X-UNC-SDK-VERSION"
 
 
 class EndpointMethod(StrEnum):
@@ -24,105 +39,340 @@ class HTTPRequestBase:
     method: EndpointMethod
     url: str
     headers: dict[str, str]
-    body: typing.Optional[str] = None
-    query_params: typing.Optional[dict[str, str]]
 
 
 @dataclass(kw_only=True)
 class HTTPGetRequest(HTTPRequestBase):
-    method: typing.Literal[EndpointMethod.GET]
+    method = EndpointMethod.GET
     query_params: dict[str, str]
 
 
 @dataclass(kw_only=True)
 class HTTPPostRequest(HTTPRequestBase):
-    method: typing.Literal[EndpointMethod.POST]
+    method = EndpointMethod.POST
+    body: typing.Union[str, dict[str, str]]
 
 
 HTTPRequest = HTTPPostRequest | HTTPGetRequest
 
 
 @dataclass(kw_only=True)
-class AuthDetailsApiKey:
-    api_id: str
-    api_secret_key: str
+class ClientConfig(ClientConfigOptions):
+    transform_request: typing.Callable[[requests.Request], requests.Request] | None = (
+        None
+    )
+    logger: typing.Optional[Logger] = None
+
 
+OAUTH_REFRESH_WINDOW_SECONDS = 60 * 5
 
-AuthDetails = AuthDetailsApiKey
+
+class APIResponseError(Exception):
+    status_code: int
+    message: str
+    extra_details: dict[str, JsonValue] | None
+
+    def __init__(
+        self,
+        status_code: int,
+        message: str,
+        extra_details: dict[str, JsonValue] | None,
+        request_id: str,
+    ) -> None:
+        super().__init__(status_code, message, extra_details)
+        self.status_code = status_code
+        self.message = message
+        self.extra_details = extra_details
+        self.request_id = request_id
+
+    @classmethod
+    def construct_error(
+        cls,
+        status_code: int,
+        extra_details: dict[str, JsonValue] | None,
+        request_id: str,
+    ) -> "APIResponseError":
+        message: str
+        match status_code:
+            case 403:
+                message = "unexpected: unauthorized"
+            case 410:
+                message = "unexpected: not found"
+            case 400:
+                message = "unexpected: bad arguments"
+            case 501:
+                message = "unexpected: unimplemented"
+            case 504:
+                message = "unexpected: timeout"
+            case 404:
+                message = "not found"
+            case 409:
+                message = "bad arguments"
+            case 422:
+                message = "unprocessable"
+            case _:
+                message = "unknown error"
+        return APIResponseError(
+            status_code=status_code,
+            message=message,
+            extra_details=extra_details,
+            request_id=request_id,
+        )
+
+    def __str__(self) -> str:
+        details_obj = {
+            "request_id": self.request_id,
+            "status_code": self.status_code,
+            "extra_details": self.extra_details,
+        }
+        details = json.dumps(details_obj)
+        return f"API response error ({self.status_code}): '{self.message}'. Details: {details}"
+
+
+class SDKError(Exception):
+    message: str
+    request_id: str
+
+    def __init__(self, message: str, *, request_id: str) -> None:
+        super().__init__(message)
+        self.message = message
+        self.request_id = request_id
+
+    def __str__(self) -> str:
+        return f"internal SDK error (request id {self.request_id}), please contact Uncountable support: {self.message}"
+
+
+@dataclass(kw_only=True)
+class OAuthBearerTokenCache:
+    token: str
+    expires_at: datetime
+
+
+@dataclass(kw_only=True)
+class GetOauthBearerTokenData:
+    access_token: str
+    expires_in: int
+    token_type: str
+    scope: str
+
+
+oauth_bearer_token_data_parser = CachedParser(GetOauthBearerTokenData)
+
+
+@dataclass
+class DownloadedFile:
+    name: str
+    size: int
+    data: BytesIO
+
+
+DownloadedFiles = list[DownloadedFile]
 
 
 class Client(ClientMethods):
     _parser_map: dict[type, CachedParser] = {}
-    _auth_details: AuthDetails
+    _auth_details: AuthDetailsAll
     _base_url: str
+    _file_uploader: FileUploader
+    _cfg: ClientConfig
+    _oauth_bearer_token_cache: OAuthBearerTokenCache | None = None
+    _session: requests.Session
 
-    def __init__(self, *, base_url: str, auth_details: AuthDetails):
+    def __init__(
+        self,
+        *,
+        base_url: str,
+        auth_details: AuthDetailsAll,
+        config: ClientConfig | None = None,
+    ):
         self._auth_details = auth_details
         self._base_url = base_url
+        self._cfg = config or ClientConfig()
+        self._session = requests.Session()
+        self._session.verify = not self._cfg.allow_insecure_tls
+        self._file_uploader = FileUploader(
+            self._base_url,
+            self._auth_details,
+            self._cfg.allow_insecure_tls,
+            logger=self._cfg.logger,
+        )
+
+    def _get_response_json(
+        self, response: requests.Response, request_id: str
+    ) -> dict[str, JsonValue]:
+        if response.status_code < 200 or response.status_code > 299:
+            extra_details: dict[str, JsonValue] | None = None
+            try:
+                data = response.json()
+                extra_details = data
+            except JSONDecodeError:
+                extra_details = {
+                    "body": response.text,
+                }
+            raise APIResponseError.construct_error(
+                status_code=response.status_code,
+                extra_details=extra_details,
+                request_id=request_id,
+            )
+        try:
+            return typing.cast(dict[str, JsonValue], response.json())
+        except JSONDecodeError as e:
+            raise SDKError("unable to process response", request_id=request_id) from e
+
+    def _send_request(self, request: requests.Request) -> requests.Response:
+        if self._cfg.extra_headers is not None:
+            request.headers = {**request.headers, **self._cfg.extra_headers}
+        if self._cfg.transform_request is not None:
+            request = self._cfg.transform_request(request)
+        prepared_request = request.prepare()
+        response = self._session.send(prepared_request)
+        return response
 
     def do_request(self, *, api_request: APIRequest, return_type: type[DT]) -> DT:
-        http_request = self._build_http_request(api_request=api_request)
+        request_id = str(uuid4())
+        http_request = self._build_http_request(
+            api_request=api_request, request_id=request_id
+        )
         match http_request:
             case HTTPGetRequest():
-                response = requests.get(
-                    http_request.url,
-                    headers=http_request.headers,
-                    params=http_request.query_params,
-                )
+                request = requests.Request("GET", http_request.url)
+                request.params = http_request.query_params
             case HTTPPostRequest():
-                response = requests.post(
-                    http_request.url,
-                    headers=http_request.headers,
-                    data=http_request.body,
-                    params=http_request.query_params,
-                )
+                request = requests.Request("POST", http_request.url)
+                request.data = http_request.body
             case _:
                 typing.assert_never(http_request)
-        if response.status_code < 200 or response.status_code > 299:
-            # TODO: handle_error
-            pass
+        request.headers = http_request.headers
+        attributes: Attributes = {
+            "method": http_request.method,
+            "endpoint": api_request.endpoint,
+        }
+        with push_scope_optional(self._cfg.logger, "api_call", attributes=attributes):
+            if self._cfg.logger is not None:
+                self._cfg.logger.log_info(api_request.endpoint, attributes=attributes)
+            response = self._send_request(request)
+        response_data = self._get_response_json(response, request_id=request_id)
         cached_parser = self._get_cached_parser(return_type)
         try:
-            data = response.json()["data"]
+            data = response_data["data"]
             return cached_parser.parse_api(data)
-        except ValueError as err:
-            # TODO: handle parse error
-            raise err
+        except (ValueError, JSONDecodeError, KeyError) as e:
+            raise SDKError("unable to process response", request_id=request_id) from e
 
     def _get_cached_parser(self, data_type: type[DT]) -> CachedParser[DT]:
         if data_type not in self._parser_map:
             self._parser_map[data_type] = CachedParser(data_type)
         return self._parser_map[data_type]
 
+    def _get_oauth_bearer_token(self, *, oauth_details: AuthDetailsOAuth) -> str:
+        if (
+            self._oauth_bearer_token_cache is None
+            or (
+                self._oauth_bearer_token_cache.expires_at - datetime.now()
+            ).total_seconds()
+            < OAUTH_REFRESH_WINDOW_SECONDS
+        ):
+            refresh_url = urljoin(self._base_url, "/token/get_bearer_token")
+            request = requests.Request("POST", refresh_url)
+            request.data = {
+                "client_secret": oauth_details.refresh_token,
+                "scope": oauth_details.scope,
+                "grant_type": "client_credentials",
+            }
+            response = self._send_request(request)
+            data = self._get_response_json(response, request_id=str(uuid4()))
+            token_data = oauth_bearer_token_data_parser.parse_storage(data)
+            self._oauth_bearer_token_cache = OAuthBearerTokenCache(
+                token=token_data.access_token,
+                expires_at=datetime.now() + timedelta(seconds=token_data.expires_in),
+            )
+
+        return self._oauth_bearer_token_cache.token
+
     def _build_auth_headers(self) -> dict[str, str]:
         match self._auth_details:
             case AuthDetailsApiKey():
                 encoded = base64.standard_b64encode(
-                    f"{self._auth_details.api_id}:{self._auth_details.api_secret_key}".encode(
-                        "utf-8"
-                    )
+                    f"{self._auth_details.api_id}:{self._auth_details.api_secret_key}".encode()
                 ).decode("utf-8")
                 return {"Authorization": f"Basic {encoded}"}
+            case AuthDetailsOAuth():
+                token = self._get_oauth_bearer_token(oauth_details=self._auth_details)
+                return {"Authorization": f"Bearer {token}"}
         typing.assert_never(self._auth_details)
 
-    def _build_http_request(self, *, api_request: APIRequest) -> HTTPRequest:
+    def _build_http_request(
+        self, *, api_request: APIRequest, request_id: str
+    ) -> HTTPRequest:
         headers = self._build_auth_headers()
+        headers[UNC_REQUEST_ID_HEADER] = request_id
+        headers[UNC_SDK_VERSION_HEADER] = get_version()
         method = api_request.method.lower()
-        query_params = {"data": json.dumps(serialize_for_api(api_request.args))}
+        data = {"data": json.dumps(serialize_for_api(api_request.args))}
         match method:
             case "get":
                 return HTTPGetRequest(
                     method=EndpointMethod.GET,
                     url=urljoin(self._base_url, api_request.endpoint),
-                    query_params=query_params,
+                    query_params=data,
                     headers=headers,
                 )
             case "post":
                 return HTTPPostRequest(
                     method=EndpointMethod.POST,
                     url=urljoin(self._base_url, api_request.endpoint),
+                    body=data,
                     headers=headers,
-                    query_params=query_params
                 )
             case _:
                 raise ValueError(f"unsupported request method: {method}")
+
+    def _get_downloaded_filename(self, *, cd: typing.Optional[str]) -> str:
+        if not cd:
+            return "Unknown"
+
+        fname = re.findall(r"filename\*=UTF-8''(.+)", cd)
+        if fname:
+            return unquote(fname[0])
+
+        fname = re.findall(r'filename="?(.+)"?', cd)
+        if fname:
+            return str(fname[0].strip('"'))
+
+        return "Unknown"
+
+    def download_files(
+        self, *, file_query: download_file_t.FileDownloadQuery
+    ) -> DownloadedFiles:
+        """Download a file from uncountable."""
+        request_id = str(uuid4())
+        api_request = APIRequest(
+            method=download_file_t.ENDPOINT_METHOD,
+            endpoint=download_file_t.ENDPOINT_PATH,
+            args=download_file_t.Arguments(
+                file_query=file_query,
+            ),
+        )
+        http_request = self._build_http_request(
+            api_request=api_request, request_id=request_id
+        )
+        request = requests.Request(http_request.method.value, http_request.url)
+        request.headers = http_request.headers
+        assert isinstance(http_request, HTTPGetRequest)
+        request.params = http_request.query_params
+        response = self._send_request(request)
+        content = response.content
+        content_disposition = response.headers.get("Content-Disposition", None)
+        return [
+            DownloadedFile(
+                name=self._get_downloaded_filename(cd=content_disposition),
+                size=len(content),
+                data=BytesIO(content),
+            )
+        ]
+
+    def upload_files(
+        self: typing.Self, *, file_uploads: list[FileUpload]
+    ) -> list[UploadedFile]:
+        """Upload files to uncountable, returning file ids that are usable with other SDK operations."""
+        return self._file_uploader.upload_files(file_uploads=file_uploads)

uncountable/core/environment.py

@@ -0,0 +1,41 @@
+import functools
+import os
+from importlib.metadata import PackageNotFoundError, version
+
+from uncountable.types import integration_server_t
+
+
+@functools.cache
+def get_version() -> str:
+    try:
+        version_str = version("UncountablePythonSDK")
+    except PackageNotFoundError:
+        version_str = "unknown"
+    return version_str
+
+
+def get_server_env() -> str | None:
+    return os.environ.get("UNC_SERVER_ENV")
+
+
+def get_webhook_server_port() -> int:
+    return int(os.environ.get("UNC_WEBHOOK_SERVER_PORT", "5001"))
+
+
+def get_local_admin_server_port() -> int:
+    return int(os.environ.get("UNC_ADMIN_SERVER_PORT", "50051"))
+
+
+def get_otel_enabled() -> bool:
+    return os.environ.get("UNC_OTEL_ENABLED") == "true"
+
+
+def get_profiles_module() -> str:
+    return os.environ["UNC_PROFILES_MODULE"]
+
+
+def get_integration_envs() -> list[integration_server_t.IntegrationEnvironment]:
+    return [
+        integration_server_t.IntegrationEnvironment(env)
+        for env in os.environ.get("UNC_INTEGRATION_ENVS", "prod").split(",")
+    ]

uncountable/core/file_upload.py

@@ -0,0 +1,135 @@
+import asyncio
+from contextlib import contextmanager
+from dataclasses import dataclass
+from enum import StrEnum
+from io import BytesIO
+from pathlib import Path
+from typing import Generator, Literal, Self, assert_never
+
+import aiohttp
+import aiotus
+
+from uncountable.integration.telemetry import Logger, push_scope_optional
+
+from .types import AuthDetailsAll, AuthDetailsApiKey
+
+_CHUNK_SIZE = 5 * 1024 * 1024  # s3 requires 5MiB minimum
+
+
+class FileUploadType(StrEnum):
+    MEDIA_FILE_UPLOAD = "MEDIA_FILE_UPLOAD"
+    DATA_FILE_UPLOAD = "DATA_FILE_UPLOAD"
+
+
+@dataclass(kw_only=True)
+class MediaFileUpload:
+    """Upload file from a path on disk"""
+
+    path: str
+    type: Literal[FileUploadType.MEDIA_FILE_UPLOAD] = FileUploadType.MEDIA_FILE_UPLOAD
+
+
+@dataclass(kw_only=True)
+class DataFileUpload:
+    data: BytesIO
+    name: str
+    type: Literal[FileUploadType.DATA_FILE_UPLOAD] = FileUploadType.DATA_FILE_UPLOAD
+
+
+FileUpload = MediaFileUpload | DataFileUpload
+
+
+@dataclass(kw_only=True)
+class FileBytes:
+    name: str
+    bytes_data: BytesIO
+
+
+@contextmanager
+def file_upload_data(file_upload: FileUpload) -> Generator[FileBytes, None, None]:
+    match file_upload:
+        case MediaFileUpload():
+            with open(file_upload.path, "rb") as f:
+                yield FileBytes(
+                    name=Path(file_upload.path).name, bytes_data=BytesIO(f.read())
+                )
+        case DataFileUpload():
+            yield FileBytes(name=file_upload.name, bytes_data=file_upload.data)
+
+
+@dataclass(kw_only=True)
+class UploadedFile:
+    name: str
+    file_id: int
+
+
+class UploadFailed(Exception):
+    pass
+
+
+class FileUploader:
+    _auth_details: AuthDetailsAll
+    _base_url: str
+    _allow_insecure_tls: bool
+
+    def __init__(
+        self: Self,
+        base_url: str,
+        auth_details: AuthDetailsAll,
+        allow_insecure_tls: bool = False,
+        logger: Logger | None = None,
+    ) -> None:
+        self._base_url = base_url
+        self._auth_details = auth_details
+        self._allow_insecure_tls = allow_insecure_tls
+        self._logger = logger
+
+    async def _upload_file(self: Self, file_upload: FileUpload) -> UploadedFile:
+        creation_url = f"{self._base_url}/api/external/file_upload/files"
+        if not isinstance(self._auth_details, AuthDetailsApiKey):
+            raise NotImplementedError("Unsupported authentication method.")
+
+        auth = aiohttp.BasicAuth(
+            self._auth_details.api_id, self._auth_details.api_secret_key
+        )
+        async with (
+            aiohttp.ClientSession(
+                auth=auth, headers={"Origin": self._base_url}
+            ) as session,
+        ):
+            attributes = {}
+            match file_upload:
+                case MediaFileUpload():
+                    attributes["file_path"] = file_upload.path
+                case DataFileUpload():
+                    attributes["file_name"] = file_upload.name
+                case _:
+                    assert_never(file_upload)
+            with push_scope_optional(
+                self._logger, "upload_file", attributes=attributes
+            ):
+                if self._logger is not None:
+                    self._logger.log_info("Uploading file", attributes=attributes)
+                with file_upload_data(file_upload) as file_bytes:
+                    location = await aiotus.upload(
+                        creation_url,
+                        file_bytes.bytes_data,
+                        {"filename": file_bytes.name.encode()},
+                        client_session=session,
+                        config=aiotus.RetryConfiguration(
+                            ssl=not self._allow_insecure_tls
+                        ),
+                        chunksize=_CHUNK_SIZE,
+                    )
+                    if location is None:
+                        raise UploadFailed(f"Failed to upload: {file_bytes.name}")
+                    return UploadedFile(
+                        name=file_bytes.name, file_id=int(location.path.split("/")[-1])
+                    )
+
+    def upload_files(
+        self: Self, *, file_uploads: list[FileUpload]
+    ) -> list[UploadedFile]:
+        return [
+            asyncio.run(self._upload_file(file_upload)) for file_upload in file_uploads
+        ]

uncountable/core/types.py

@@ -0,0 +1,17 @@
+from dataclasses import dataclass
+
+
+@dataclass(kw_only=True)
+class AuthDetailsApiKey:
+    api_id: str
+    api_secret_key: str
+
+
+@dataclass(kw_only=True)
+class AuthDetailsOAuth:
+    refresh_token: str
+    scope: str = "unc.rnd"
+
+
+AuthDetails = AuthDetailsApiKey  # Legacy Mapping
+AuthDetailsAll = AuthDetailsApiKey | AuthDetailsOAuth

uncountable/integration/cli.py

@@ -0,0 +1,49 @@
+import argparse
+
+from opentelemetry.trace import get_current_span
+
+from uncountable.core.environment import get_local_admin_server_port
+from uncountable.integration.queue_runner.command_server.command_client import (
+    send_job_queue_message,
+)
+from uncountable.integration.telemetry import Logger
+from uncountable.types import queued_job_t
+
+
+def main() -> None:
+    logger = Logger(get_current_span())
+
+    parser = argparse.ArgumentParser(
+        description="Process a job with a given command and job ID."
+    )
+
+    parser.add_argument(
+        "command",
+        type=str,
+        choices=["run"],
+        help="The command to execute (e.g., 'run')",
+    )
+
+    parser.add_argument("job_id", type=str, help="The ID of the job to process")
+
+    parser.add_argument(
+        "--host", type=str, default="localhost", nargs="?", help="The host to run on"
+    )
+
+    args = parser.parse_args()
+
+    with logger.push_scope(args.command):
+        if args.command == "run":
+            send_job_queue_message(
+                job_ref_name=args.job_id,
+                payload=queued_job_t.QueuedJobPayload(
+                    invocation_context=queued_job_t.InvocationContextManual()
+                ),
+                host=args.host,
+                port=get_local_admin_server_port(),
+            )
+        else:
+            parser.print_usage()
+
+
+main()