UncountablePythonSDK 0.0.83__py3-none-any.whl → 0.0.132__py3-none-any.whl

pkgs/type_spec/value_spec/convert_type.py

@@ -31,6 +31,11 @@ TYPE_MAP = {
     # not part of type_spec's types now
     "Symbol": MappedType(base_type=value_spec_t.BaseType.SYMBOL),
     "Any": MappedType(base_type=value_spec_t.BaseType.ANY),
+    "None": MappedType(base_type=value_spec_t.BaseType.NONE),
+    "Tuple": MappedType(
+        base_type=value_spec_t.BaseType.TUPLE, variable_param_count=True
+    ),
+    "Never": MappedType(base_type=value_spec_t.BaseType.NEVER),
 }
 
 
@@ -56,3 +61,16 @@ def convert_to_value_spec_type(parsed: ParsedTypePath) -> value_spec_t.ValueType
     return value_spec_t.ValueType(base_type=mapped.base_type, parameters=parameters)
 
     # Our formatter was duplicating the previous line for an unknown reason, this comment blocks that
+
+
+def convert_from_value_spec_type(
+    base_type: value_spec_t.BaseType,
+) -> str:
+    for type_spec_type, mapped_type in TYPE_MAP.items():
+        if (
+            mapped_type.base_type == base_type
+            and mapped_type.param_count == 0
+            and mapped_type.variable_param_count is False
+        ):
+            return type_spec_type
+    raise ValueError(f"invalid value spec type {base_type}")

pkgs/type_spec/value_spec/emit_python.py

@@ -70,14 +70,19 @@ def _emit_function_wrapper(function: value_spec_t.Function) -> str:
         else:
             python_type = _emit_python_type(argument.type)
             if (
-                argument.pass_null
+                argument.on_null == value_spec_t.OnNull.PASS
                 or argument.extant == value_spec_t.ArgumentExtant.MISSING
             ):
                 python_type += " | None"
                 any_pass_null = True
+
+            if python_type.startswith("base_t.ExtJsonValue"):
+                return_statement = f"self._extract({index})"
+            else:
+                return_statement = f"cast({python_type}, self._extract({index}))"
             out.write(
                 f"""{INDENT}def get_{argument.ref_name}(self) -> {python_type}:
-{INDENT}{INDENT}return cast({python_type}, self._extract({index}))
+{INDENT}{INDENT}return {return_statement}
 """
             )
         out.write("\n")
@@ -164,6 +169,8 @@ def _emit_function(function: value_spec_t.Function, indent: str) -> str:
         f"{sub_indent}description={encode_common_string(function.description)},\n"
     )
     out.write(f"{sub_indent}brief={encode_common_string(function.brief)},\n")
+    if function.draft:
+        out.write(f"{sub_indent}draft={function.draft},\n")
     out.write(
         f"{sub_indent}return_value={_emit_function_return(function.return_value, sub_indent)},\n"
     )
@@ -189,7 +196,10 @@ def _emit_argument(argument: value_spec_t.FunctionArgument, indent: str) -> str:
     out.write(
         f"{sub_indent}description={encode_common_string(argument.description)},\n"
     )
-    out.write(f"{sub_indent}pass_null={str(argument.pass_null)},\n")
+    # Quick enum emit since we have only one such type here
+    out.write(
+        f"{sub_indent}on_null=value_spec_t.OnNull.{str(argument.on_null).upper()},\n"
+    )
     out.write(
         f"{sub_indent}extant=value_spec_t.ArgumentExtant.{argument.extant.name},\n"
     )

pkgs/type_spec/value_spec/types.py

@@ -8,7 +8,7 @@ from ..util import ParsedTypePath
 @dataclass(kw_only=True, frozen=True)
 class ParsedFunctionArgument:
     ref_name: str
-    pass_null: bool
+    on_null: value_spec_t.OnNull
     extant: value_spec_t.ArgumentExtant
     type_path: ParsedTypePath
 

uncountable/core/async_batch.py

@@ -4,7 +4,7 @@ from uncountable.types.async_batch import AsyncBatchRequest
 from uncountable.types.async_batch_processor import AsyncBatchProcessorBase
 
 
-class AsyncBatchSubmissionError(BaseException):
+class AsyncBatchSubmissionError(Exception):
     pass
 
 

uncountable/core/client.py

@@ -1,9 +1,12 @@
 import base64
+import datetime
+import re
 import typing
 from dataclasses import dataclass
-from datetime import datetime, timedelta
+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
@@ -12,10 +15,10 @@ 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 JobLogger
+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
 
@@ -41,14 +44,14 @@ class HTTPRequestBase:
 
 @dataclass(kw_only=True)
 class HTTPGetRequest(HTTPRequestBase):
-    method = EndpointMethod.GET
+    method: EndpointMethod = EndpointMethod.GET
     query_params: dict[str, str]
 
 
 @dataclass(kw_only=True)
 class HTTPPostRequest(HTTPRequestBase):
-    method = EndpointMethod.POST
-    body: typing.Union[str, dict[str, str]]
+    method: EndpointMethod = EndpointMethod.POST
+    body: str | dict[str, str]
 
 
 HTTPRequest = HTTPPostRequest | HTTPGetRequest
@@ -59,7 +62,7 @@ class ClientConfig(ClientConfigOptions):
     transform_request: typing.Callable[[requests.Request], requests.Request] | None = (
         None
     )
-    job_logger: typing.Optional[JobLogger] = None
+    logger: Logger | None = None
 
 
 OAUTH_REFRESH_WINDOW_SECONDS = 60 * 5
@@ -71,16 +74,24 @@ class APIResponseError(Exception):
     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:
@@ -103,9 +114,21 @@ class APIResponseError(Exception):
             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(Exception):
     message: str
@@ -123,7 +146,7 @@ class SDKError(Exception):
 @dataclass(kw_only=True)
 class OAuthBearerTokenCache:
     token: str
-    expires_at: datetime
+    expires_at: datetime.datetime
 
 
 @dataclass(kw_only=True)
@@ -137,6 +160,16 @@ class GetOauthBearerTokenData:
 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: AuthDetailsAll
@@ -159,35 +192,49 @@ class Client(ClientMethods):
         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
+            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]:
+    @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
+                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) -> requests.Response:
+    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)
+        response = self._session.send(prepared_request, timeout=timeout)
         return response
 
     def do_request(self, *, api_request: APIRequest, return_type: type[DT]) -> DT:
@@ -205,15 +252,19 @@ class Client(ClientMethods):
             case _:
                 typing.assert_never(http_request)
         request.headers = http_request.headers
-        if self._cfg.job_logger is not None:
-            attributes: Attributes = {
-                "method": http_request.method,
-                "endpoint": api_request.endpoint,
-            }
-            with self._cfg.job_logger.push_scope("api_call", attributes=attributes):
-                response = self._send_request(request)
-        else:
-            response = self._send_request(request)
+        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:
@@ -231,7 +282,8 @@ class Client(ClientMethods):
         if (
             self._oauth_bearer_token_cache is None
             or (
-                self._oauth_bearer_token_cache.expires_at - datetime.now()
+                self._oauth_bearer_token_cache.expires_at
+                - datetime.datetime.now(tz=UTC)
             ).total_seconds()
             < OAUTH_REFRESH_WINDOW_SECONDS
         ):
@@ -247,7 +299,8 @@ class Client(ClientMethods):
             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),
+                expires_at=datetime.datetime.now(tz=UTC)
+                + timedelta(seconds=token_data.expires_in),
             )
 
         return self._oauth_bearer_token_cache.token
@@ -290,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

@@ -18,12 +18,12 @@ 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_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))
+    return int(os.environ.get("UNC_ADMIN_SERVER_PORT", "50051"))
 
 
 def get_otel_enabled() -> bool:

uncountable/core/file_upload.py

@@ -4,11 +4,13 @@ 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 uncountable.integration.telemetry import Logger, push_scope_optional
+
 from .types import AuthDetailsAll, AuthDetailsApiKey
 
 _CHUNK_SIZE = 5 * 1024 * 1024  # s3 requires 5MiB minimum
@@ -75,10 +77,12 @@ class FileUploader:
         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"
@@ -93,20 +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,
-                    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])
-                )
+            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/integration/cli.py

@@ -1,49 +1,142 @@
 import argparse
 
+from dateutil import tz
 from opentelemetry.trace import get_current_span
+from tabulate import tabulate
 
 from uncountable.core.environment import get_local_admin_server_port
 from uncountable.integration.queue_runner.command_server.command_client import (
     send_job_queue_message,
+    send_list_queued_jobs_message,
+    send_retry_job_message,
 )
 from uncountable.integration.telemetry import Logger
 from uncountable.types import queued_job_t
 
 
-def main() -> None:
-    logger = Logger(get_current_span())
+def register_enqueue_job_parser(
+    sub_parser_manager: argparse._SubParsersAction,
+    parents: list[argparse.ArgumentParser],
+) -> None:
+    run_parser = sub_parser_manager.add_parser(
+        "run",
+        parents=parents,
+        help="Process a job with a given host and job ID",
+        description="Process a job with a given host and job ID",
+    )
+    run_parser.add_argument("job_id", type=str, help="The ID of the job to process")
+
+    def _handle_enqueue_job(args: argparse.Namespace) -> None:
+        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(),
+        )
 
-    parser = argparse.ArgumentParser(
-        description="Process a job with a given command and job ID."
+    run_parser.set_defaults(func=_handle_enqueue_job)
+
+
+def register_list_queued_jobs(
+    sub_parser_manager: argparse._SubParsersAction,
+    parents: list[argparse.ArgumentParser],
+) -> None:
+    list_queued_jobs_parser = sub_parser_manager.add_parser(
+        "list-queued-jobs",
+        parents=parents,
+        help="List all jobs queued on the integration server",
+        description="List all jobs queued on the integration server",
     )
 
-    parser.add_argument(
-        "command",
-        type=str,
-        choices=["run"],
-        help="The command to execute (e.g., 'run')",
+    list_queued_jobs_parser.add_argument(
+        "--offset",
+        type=int,
+        default=0,
+        help="Number of jobs to skip. Should be non-negative.",
+    )
+    list_queued_jobs_parser.add_argument(
+        "--limit",
+        type=int,
+        default=100,
+        help="A number between 1 and 100 specifying the number of jobs to return in the result set.",
+    )
+
+    def _handle_list_queued_jobs(args: argparse.Namespace) -> None:
+        queued_jobs = send_list_queued_jobs_message(
+            offset=args.offset,
+            limit=args.limit,
+            host=args.host,
+            port=get_local_admin_server_port(),
+        )
+
+        headers = ["UUID", "Job Ref Name", "Attempts", "Status", "Submitted At"]
+        rows = [
+            [
+                job.uuid,
+                job.job_ref_name,
+                job.num_attempts,
+                job.status,
+                job.submitted_at.ToDatetime(tz.UTC).astimezone(tz.tzlocal()),
+            ]
+            for job in queued_jobs
+        ]
+        print(tabulate(rows, headers=headers, tablefmt="grid"))
+
+    list_queued_jobs_parser.set_defaults(func=_handle_list_queued_jobs)
+
+
+def register_retry_job_parser(
+    sub_parser_manager: argparse._SubParsersAction,
+    parents: list[argparse.ArgumentParser],
+) -> None:
+    retry_failed_jobs_parser = sub_parser_manager.add_parser(
+        "retry-job",
+        parents=parents,
+        help="Retry failed job on the integration server",
+        description="Retry failed job on the integration server",
+    )
+
+    retry_failed_jobs_parser.add_argument(
+        "job_uuid", type=str, help="The uuid of the job to retry"
     )
 
-    parser.add_argument("job_id", type=str, help="The ID of the job to process")
+    def _handle_retry_job(args: argparse.Namespace) -> None:
+        send_retry_job_message(
+            job_uuid=args.job_uuid,
+            host=args.host,
+            port=get_local_admin_server_port(),
+        )
+
+    retry_failed_jobs_parser.set_defaults(func=_handle_retry_job)
 
-    parser.add_argument(
+
+def main() -> None:
+    logger = Logger(get_current_span())
+
+    main_parser = argparse.ArgumentParser(
+        description="Execute a given integrations server command."
+    )
+
+    base_parser = argparse.ArgumentParser(add_help=False)
+    base_parser.add_argument(
         "--host", type=str, default="localhost", nargs="?", help="The host to run on"
     )
 
-    args = parser.parse_args()
+    subparser_action = main_parser.add_subparsers(
+        dest="command",
+        required=True,
+        help="The command to execute (e.g., 'run')",
+    )
+
+    register_enqueue_job_parser(subparser_action, parents=[base_parser])
+    register_retry_job_parser(subparser_action, parents=[base_parser])
+    register_list_queued_jobs(subparser_action, parents=[base_parser])
 
+    args = main_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()
+        args.func(args)
 
 
 main()

uncountable/integration/construct_client.py

@@ -37,15 +37,15 @@ def _construct_client_config(
     return ClientConfig(
         allow_insecure_tls=profile_meta.client_options.allow_insecure_tls,
         extra_headers=profile_meta.client_options.extra_headers,
-        job_logger=job_logger,
+        logger=job_logger,
     )
 
 
 def construct_uncountable_client(
-    profile_meta: ProfileMetadata, job_logger: JobLogger
+    profile_meta: ProfileMetadata, logger: JobLogger
 ) -> Client:
     return Client(
         base_url=profile_meta.base_url,
         auth_details=_construct_auth_details(profile_meta),
-        config=_construct_client_config(profile_meta, job_logger),
+        config=_construct_client_config(profile_meta, logger),
     )