UncountablePythonSDK 0.0.24__py3-none-any.whl → 0.0.131__py3-none-any.whl

uncountable/core/client.py

@@ -1,22 +1,33 @@
 import base64
-import json
+import datetime
+import re
 import typing
 from dataclasses import dataclass
+from datetime import UTC, timedelta
 from enum import StrEnum
-from urllib.parse import urljoin
+from io import BytesIO
+from urllib.parse import unquote, urljoin
+from uuid import uuid4
 
 import requests
+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.serialization_helpers import JsonValue
+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 AuthDetails, AuthDetailsApiKey
+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):
@@ -29,47 +40,58 @@ class HTTPRequestBase:
     method: EndpointMethod
     url: str
     headers: dict[str, str]
-    body: typing.Optional[typing.Union[str, dict[str, str]]] = None
-    query_params: typing.Optional[dict[str, str]] = None
 
 
 @dataclass(kw_only=True)
 class HTTPGetRequest(HTTPRequestBase):
-    method: typing.Literal[EndpointMethod.GET]
+    method: EndpointMethod = EndpointMethod.GET
     query_params: dict[str, str]
 
 
 @dataclass(kw_only=True)
 class HTTPPostRequest(HTTPRequestBase):
-    method: typing.Literal[EndpointMethod.POST]
-    body: typing.Union[str, dict[str, str]]
+    method: EndpointMethod = EndpointMethod.POST
+    body: str | dict[str, str]
 
 
 HTTPRequest = HTTPPostRequest | HTTPGetRequest
 
 
-
 @dataclass(kw_only=True)
-class ClientConfig():
-    allow_insecure_tls: bool = False
-      
+class ClientConfig(ClientConfigOptions):
+    transform_request: typing.Callable[[requests.Request], requests.Request] | None = (
+        None
+    )
+    logger: Logger | None = None
+
+
+OAUTH_REFRESH_WINDOW_SECONDS = 60 * 5
 
-class APIResponseError(BaseException):
+
+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
+        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
+        cls,
+        status_code: int,
+        extra_details: dict[str, JsonValue] | None,
+        request_id: str,
     ) -> "APIResponseError":
         message: str
         match status_code:
@@ -92,77 +114,197 @@ class APIResponseError(BaseException):
             case _:
                 message = "unknown error"
         return APIResponseError(
-            status_code=status_code, message=message, extra_details=extra_details
+            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(BaseException):
+class SDKError(Exception):
     message: str
+    request_id: str
 
-    def __init__(self, message: str) -> None:
+    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, please contact Uncountable support: {self.message}"
+        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.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, config: ClientConfig | None = None):
+    def __init__(
+        self,
+        *,
+        base_url: str,
+        auth_details: AuthDetailsAll,
+        config: ClientConfig | None = None,
+    ):
         self._auth_details = auth_details
         self._base_url = base_url
-        self._file_uploader = FileUploader(self._base_url, self._auth_details)
         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 do_request(self, *, api_request: APIRequest, return_type: type[DT]) -> DT:
-        http_request = self._build_http_request(api_request=api_request)
-        match http_request:
-            case HTTPGetRequest():
-                response = requests.get(
-                    http_request.url,
-                    headers=http_request.headers,
-                    params=http_request.query_params,
-                    verify=not self._cfg.allow_insecure_tls
-                )
-            case HTTPPostRequest():
-                response = requests.post(
-                    http_request.url,
-                    headers=http_request.headers,
-                    data=http_request.body,
-                    params=http_request.query_params,
-                    verify=not self._cfg.allow_insecure_tls
-                )
-            case _:
-                typing.assert_never(http_request)
+    @classmethod
+    def _validate_response_status(
+        cls, response: requests.Response, request_id: str
+    ) -> None:
         if response.status_code < 200 or response.status_code > 299:
             extra_details: dict[str, JsonValue] | None = None
             try:
                 data = response.json()
-                if "error" in data:
-                    extra_details = data["error"]
+                extra_details = data
             except JSONDecodeError:
-                pass
+                extra_details = {
+                    "body": response.text,
+                }
             raise APIResponseError.construct_error(
-                status_code=response.status_code, extra_details=extra_details
+                status_code=response.status_code,
+                extra_details=extra_details,
+                request_id=request_id,
+            )
+
+    def _get_response_json(
+        self, response: requests.Response, request_id: str
+    ) -> dict[str, JsonValue]:
+        self._validate_response_status(response, 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, *, timeout: float | None = None
+    ) -> 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, timeout=timeout)
+        return response
+
+    def do_request(self, *, api_request: APIRequest, return_type: type[DT]) -> DT:
+        request_id = str(uuid4())
+        http_request = self._build_http_request(
+            api_request=api_request, request_id=request_id
+        )
+        match http_request:
+            case HTTPGetRequest():
+                request = requests.Request("GET", http_request.url)
+                request.params = http_request.query_params
+            case HTTPPostRequest():
+                request = requests.Request("POST", http_request.url)
+                request.data = http_request.body
+            case _:
+                typing.assert_never(http_request)
+        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)
+            timeout = (
+                api_request.request_options.timeout_secs
+                if api_request.request_options is not None
+                else None
             )
+            response = self._send_request(request, timeout=timeout)
+        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 | JSONDecodeError:
-            raise SDKError("unable to process response")
+        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.datetime.now(tz=UTC)
+            ).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.datetime.now(tz=UTC)
+                + 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():
@@ -170,10 +312,17 @@ class Client(ClientMethods):
                     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()
         data = {"data": json.dumps(serialize_for_api(api_request.args))}
         match method:
@@ -194,6 +343,52 @@ class Client(ClientMethods):
             case _:
                 raise ValueError(f"unsupported request method: {method}")
 
+    def _get_downloaded_filename(self, *, cd: str | None) -> 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)
+        self._validate_response_status(response, request_id)
+
+        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]:

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_http_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

@@ -4,18 +4,21 @@ from dataclasses import dataclass
 from enum import StrEnum
 from io import BytesIO
 from pathlib import Path
-from typing import Generator, Literal, Self
+from typing import Generator, Literal, Self, assert_never
 
 import aiohttp
 import aiotus
 
-from .types import AuthDetails
+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)
@@ -23,10 +26,17 @@ class MediaFileUpload:
     """Upload file from a path on disk"""
 
     path: str
-    type: FileUploadType.MEDIA_FILE_UPLOAD = FileUploadType.MEDIA_FILE_UPLOAD
+    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
+FileUpload = MediaFileUpload | DataFileUpload
 
 
 @dataclass(kw_only=True)
@@ -37,12 +47,14 @@ class FileBytes:
 
 @contextmanager
 def file_upload_data(file_upload: FileUpload) -> Generator[FileBytes, None, None]:
-    match file_upload.type:
-        case FileUploadType.MEDIA_FILE_UPLOAD:
+    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)
@@ -56,15 +68,27 @@ class UploadFailed(Exception):
 
 
 class FileUploader:
-    _auth_details: AuthDetails
+    _auth_details: AuthDetailsAll
     _base_url: str
-
-    def __init__(self: Self, base_url: str, auth_details: AuthDetails) -> None:
+    _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
         )
@@ -73,19 +97,40 @@ class FileUploader:
                 auth=auth, headers={"Origin": self._base_url}
             ) as session,
         ):
-            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,
-                    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])
-                )
+            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:
+                    if file_bytes.bytes_data.read(1) == b"":
+                        raise UploadFailed(
+                            f"Failed to upload empty file: {file_bytes.name}"
+                        )
+                    file_bytes.bytes_data.seek(0)
+                    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]

uncountable/core/types.py

@@ -1,16 +1,4 @@
-import base64
-import json
-import typing
 from dataclasses import dataclass
-from enum import StrEnum
-from urllib.parse import urljoin
-
-import aiohttp
-import requests
-
-from pkgs.argument_parser import CachedParser
-from pkgs.serialization_util import serialize_for_api
-from uncountable.types.client_base import APIRequest, ClientMethods
 
 
 @dataclass(kw_only=True)
@@ -19,4 +7,11 @@ class AuthDetailsApiKey:
     api_secret_key: str
 
 
-AuthDetails = AuthDetailsApiKey
+@dataclass(kw_only=True)
+class AuthDetailsOAuth:
+    refresh_token: str
+    scope: str = "unc.rnd"
+
+
+AuthDetails = AuthDetailsApiKey  # Legacy Mapping
+AuthDetailsAll = AuthDetailsApiKey | AuthDetailsOAuth