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

uncountable/integration/queue_runner/command_server/command_client.py

@@ -0,0 +1,133 @@
+from contextlib import contextmanager
+from typing import Generator
+
+import grpc
+import simplejson as json
+
+from pkgs.serialization_util import serialize_for_api
+from uncountable.integration.queue_runner.command_server.protocol.command_server_pb2 import (
+    CheckHealthRequest,
+    CheckHealthResult,
+    EnqueueJobRequest,
+    EnqueueJobResult,
+    ListQueuedJobsRequest,
+    ListQueuedJobsResult,
+    RetryJobRequest,
+    RetryJobResult,
+    VaccuumQueuedJobsRequest,
+    VaccuumQueuedJobsResult,
+)
+from uncountable.integration.queue_runner.command_server.types import (
+    CommandServerBadResponse,
+    CommandServerTimeout,
+)
+from uncountable.types import queued_job_t
+
+from .protocol.command_server_pb2_grpc import CommandServerStub
+
+_LOCAL_RPC_HOST = "localhost"
+_DEFAULT_MESSAGE_TIMEOUT_SECS = 2
+
+
+@contextmanager
+def command_server_connection(
+    host: str, port: int
+) -> Generator[CommandServerStub, None, None]:
+    try:
+        with grpc.insecure_channel(f"{host}:{port}") as channel:
+            stub = CommandServerStub(channel)
+            yield stub
+    except grpc._channel._InactiveRpcError as e:
+        raise CommandServerTimeout() from e
+
+
+def send_job_queue_message(
+    *,
+    job_ref_name: str,
+    payload: queued_job_t.QueuedJobPayload,
+    host: str = "localhost",
+    port: int,
+) -> str:
+    with command_server_connection(host=host, port=port) as stub:
+        request = EnqueueJobRequest(
+            job_ref_name=job_ref_name,
+            serialized_payload=json.dumps(serialize_for_api(payload)),
+        )
+
+        response = stub.EnqueueJob(request, timeout=_DEFAULT_MESSAGE_TIMEOUT_SECS)
+
+        assert isinstance(response, EnqueueJobResult)
+        if not response.successfully_queued:
+            raise CommandServerBadResponse("queue operation was not successful")
+
+        return response.queued_job_uuid
+
+
+def send_retry_job_message(
+    *,
+    job_uuid: str,
+    host: str = "localhost",
+    port: int,
+) -> str:
+    with command_server_connection(host=host, port=port) as stub:
+        request = RetryJobRequest(uuid=job_uuid)
+
+        try:
+            response = stub.RetryJob(request, timeout=_DEFAULT_MESSAGE_TIMEOUT_SECS)
+            assert isinstance(response, RetryJobResult)
+            if not response.successfully_queued:
+                raise CommandServerBadResponse("queue operation was not successful")
+
+            return response.queued_job_uuid
+        except grpc.RpcError as e:
+            raise ValueError(e.details())  # type: ignore
+
+
+def check_health(*, host: str = _LOCAL_RPC_HOST, port: int) -> bool:
+    with command_server_connection(host=host, port=port) as stub:
+        request = CheckHealthRequest()
+
+        response = stub.CheckHealth(request, timeout=_DEFAULT_MESSAGE_TIMEOUT_SECS)
+
+        assert isinstance(response, CheckHealthResult)
+
+        return response.success
+
+
+def send_list_queued_jobs_message(
+    *,
+    offset: int,
+    limit: int,
+    host: str = "localhost",
+    port: int,
+) -> list[ListQueuedJobsResult.ListQueuedJobsResultItem]:
+    with command_server_connection(host=host, port=port) as stub:
+        request = ListQueuedJobsRequest(
+            offset=offset,
+            limit=limit,
+        )
+
+        try:
+            response = stub.ListQueuedJobs(
+                request, timeout=_DEFAULT_MESSAGE_TIMEOUT_SECS
+            )
+        except grpc.RpcError as e:
+            raise ValueError(e.details())  # type: ignore
+
+        assert isinstance(response, ListQueuedJobsResult)
+        return list(response.queued_jobs)
+
+
+def send_vaccuum_queued_jobs_message(*, host: str = "localhost", port: int) -> None:
+    with command_server_connection(host=host, port=port) as stub:
+        request = VaccuumQueuedJobsRequest()
+
+        try:
+            response = stub.VaccuumQueuedJobs(
+                request, timeout=_DEFAULT_MESSAGE_TIMEOUT_SECS
+            )
+        except grpc.RpcError as e:
+            raise ValueError(e.details())  # type: ignore
+
+        assert isinstance(response, VaccuumQueuedJobsResult)
+        return None

uncountable/integration/queue_runner/command_server/command_server.py

@@ -0,0 +1,142 @@
+import asyncio
+
+import grpc.aio as grpc_aio
+import simplejson as json
+from google.protobuf.timestamp_pb2 import Timestamp
+from grpc import StatusCode
+
+from pkgs.argument_parser import CachedParser
+from uncountable.core.environment import get_local_admin_server_port
+from uncountable.integration.queue_runner.command_server.protocol.command_server_pb2 import (
+    CheckHealthRequest,
+    CheckHealthResult,
+    EnqueueJobRequest,
+    EnqueueJobResult,
+    ListQueuedJobsRequest,
+    ListQueuedJobsResult,
+    RetryJobRequest,
+    RetryJobResult,
+    VaccuumQueuedJobsRequest,
+    VaccuumQueuedJobsResult,
+)
+from uncountable.integration.queue_runner.command_server.types import (
+    CommandEnqueueJob,
+    CommandEnqueueJobResponse,
+    CommandQueue,
+    CommandRetryJob,
+    CommandRetryJobResponse,
+    CommandVaccuumQueuedJobs,
+    CommandVaccuumQueuedJobsResponse,
+)
+from uncountable.integration.queue_runner.datastore import DatastoreSqlite
+from uncountable.types import queued_job_t
+
+from .constants import ListQueuedJobsConstants
+from .protocol.command_server_pb2_grpc import (
+    CommandServerServicer,
+    add_CommandServerServicer_to_server,
+)
+
+queued_job_payload_parser = CachedParser(queued_job_t.QueuedJobPayload)
+
+
+async def serve(command_queue: CommandQueue, datastore: DatastoreSqlite) -> None:
+    server = grpc_aio.server()
+
+    class CommandServerHandler(CommandServerServicer):
+        async def EnqueueJob(
+            self, request: EnqueueJobRequest, context: grpc_aio.ServicerContext
+        ) -> EnqueueJobResult:
+            payload_json = json.loads(request.serialized_payload)
+            payload = queued_job_payload_parser.parse_api(payload_json)
+            response_queue: asyncio.Queue[CommandEnqueueJobResponse] = asyncio.Queue()
+            await command_queue.put(
+                CommandEnqueueJob(
+                    job_ref_name=request.job_ref_name,
+                    payload=payload,
+                    response_queue=response_queue,
+                )
+            )
+            response = await response_queue.get()
+            result = EnqueueJobResult(
+                successfully_queued=True, queued_job_uuid=response.queued_job_uuid
+            )
+            return result
+
+        async def RetryJob(
+            self, request: RetryJobRequest, context: grpc_aio.ServicerContext
+        ) -> RetryJobResult:
+            response_queue: asyncio.Queue[CommandRetryJobResponse] = asyncio.Queue()
+            await command_queue.put(
+                CommandRetryJob(
+                    queued_job_uuid=request.uuid, response_queue=response_queue
+                )
+            )
+            response = await response_queue.get()
+            if response.queued_job_uuid is not None:
+                return RetryJobResult(
+                    successfully_queued=True, queued_job_uuid=response.queued_job_uuid
+                )
+            else:
+                return RetryJobResult(successfully_queued=False, queued_job_uuid="")
+
+        async def CheckHealth(
+            self, request: CheckHealthRequest, context: grpc_aio.ServicerContext
+        ) -> CheckHealthResult:
+            return CheckHealthResult(success=True)
+
+        async def ListQueuedJobs(
+            self, request: ListQueuedJobsRequest, context: grpc_aio.ServicerContext
+        ) -> ListQueuedJobsResult:
+            if (
+                request.limit < ListQueuedJobsConstants.LIMIT_MIN
+                or request.limit > ListQueuedJobsConstants.LIMIT_MAX
+            ):
+                await context.abort(
+                    StatusCode.INVALID_ARGUMENT, "Limit must be between 1 and 100."
+                )
+
+            if request.offset < ListQueuedJobsConstants.OFFSET_MIN:
+                await context.abort(
+                    StatusCode.INVALID_ARGUMENT, "Offset cannot be negative."
+                )
+
+            queued_job_metadata = datastore.list_queued_job_metadata(
+                offset=request.offset, limit=request.limit
+            )
+
+            response_list: list[ListQueuedJobsResult.ListQueuedJobsResultItem] = []
+            for item in queued_job_metadata:
+                proto_timestamp = Timestamp()
+                proto_timestamp.FromDatetime(item.submitted_at)
+
+                response_list.append(
+                    ListQueuedJobsResult.ListQueuedJobsResultItem(
+                        uuid=item.queued_job_uuid,
+                        job_ref_name=item.job_ref_name,
+                        num_attempts=item.num_attempts,
+                        submitted_at=proto_timestamp,
+                        status=item.status,
+                    )
+                )
+            return ListQueuedJobsResult(queued_jobs=response_list)
+
+        async def VaccuumQueuedJobs(
+            self, request: VaccuumQueuedJobsRequest, context: grpc_aio.ServicerContext
+        ) -> VaccuumQueuedJobsResult:
+            response_queue: asyncio.Queue[CommandVaccuumQueuedJobsResponse] = (
+                asyncio.Queue()
+            )
+            await command_queue.put(
+                CommandVaccuumQueuedJobs(response_queue=response_queue)
+            )
+            return VaccuumQueuedJobsResult()
+
+    add_CommandServerServicer_to_server(CommandServerHandler(), server)
+
+    listen_addr = f"[::]:{get_local_admin_server_port()}"
+
+    server.add_insecure_port(listen_addr)
+
+    await server.start()
+    await server.wait_for_termination()

uncountable/integration/queue_runner/command_server/constants.py

@@ -0,0 +1,4 @@
+class ListQueuedJobsConstants:
+    LIMIT_MIN = 1
+    LIMIT_MAX = 100
+    OFFSET_MIN = 0

uncountable/integration/queue_runner/command_server/protocol/command_server.proto

@@ -0,0 +1,58 @@
+syntax = "proto3";
+import "google/protobuf/timestamp.proto";
+
+service CommandServer {
+	rpc EnqueueJob(EnqueueJobRequest) returns (EnqueueJobResult) {}
+	rpc RetryJob(RetryJobRequest) returns (RetryJobResult) {}
+	rpc CheckHealth(CheckHealthRequest) returns (CheckHealthResult) {}
+	rpc ListQueuedJobs(ListQueuedJobsRequest) returns (ListQueuedJobsResult) {}
+	rpc VaccuumQueuedJobs(VaccuumQueuedJobsRequest) returns (VaccuumQueuedJobsResult) {}
+}
+
+message EnqueueJobRequest {
+	string job_ref_name = 1;
+	string serialized_payload = 2;
+}
+
+message EnqueueJobResult {
+	bool successfully_queued = 1;
+	string queued_job_uuid = 2;
+}
+
+message RetryJobRequest {
+	string uuid = 1;
+}
+
+message RetryJobResult {
+	bool successfully_queued = 1;
+	string queued_job_uuid = 2;
+}
+
+message VaccuumQueuedJobsRequest {
+}
+
+message VaccuumQueuedJobsResult {
+}
+
+message CheckHealthRequest {}
+
+message CheckHealthResult {
+	bool success = 1;
+}
+
+message ListQueuedJobsRequest {
+	uint32 offset = 1;
+	uint32 limit = 2;
+}
+
+message ListQueuedJobsResult {
+	message ListQueuedJobsResultItem {
+	string uuid = 1;
+	string job_ref_name = 2;
+	int64 num_attempts = 3;
+	google.protobuf.Timestamp submitted_at = 4;
+	string status = 5;
+	}
+
+	repeated ListQueuedJobsResultItem queued_jobs  = 1;
+}

uncountable/integration/queue_runner/command_server/protocol/command_server_pb2.py

@@ -0,0 +1,57 @@
+# ruff: noqa
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: uncountable/integration/queue_runner/command_server/protocol/command_server.proto
+# Protobuf Python Version: 4.25.1
+"""Generated protocol buffer code."""
+
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+from google.protobuf import timestamp_pb2 as google_dot_protobuf_dot_timestamp__pb2
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(
+    b'\nQuncountable/integration/queue_runner/command_server/protocol/command_server.proto\x1a\x1fgoogle/protobuf/timestamp.proto"E\n\x11\x45nqueueJobRequest\x12\x14\n\x0cjob_ref_name\x18\x01 \x01(\t\x12\x1a\n\x12serialized_payload\x18\x02 \x01(\t"H\n\x10\x45nqueueJobResult\x12\x1b\n\x13successfully_queued\x18\x01 \x01(\x08\x12\x17\n\x0fqueued_job_uuid\x18\x02 \x01(\t"\x1f\n\x0fRetryJobRequest\x12\x0c\n\x04uuid\x18\x01 \x01(\t"F\n\x0eRetryJobResult\x12\x1b\n\x13successfully_queued\x18\x01 \x01(\x08\x12\x17\n\x0fqueued_job_uuid\x18\x02 \x01(\t"\x1a\n\x18VaccuumQueuedJobsRequest"\x19\n\x17VaccuumQueuedJobsResult"\x14\n\x12\x43heckHealthRequest"$\n\x11\x43heckHealthResult\x12\x0f\n\x07success\x18\x01 \x01(\x08"6\n\x15ListQueuedJobsRequest\x12\x0e\n\x06offset\x18\x01 \x01(\r\x12\r\n\x05limit\x18\x02 \x01(\r"\xf4\x01\n\x14ListQueuedJobsResult\x12\x43\n\x0bqueued_jobs\x18\x01 \x03(\x0b\x32..ListQueuedJobsResult.ListQueuedJobsResultItem\x1a\x96\x01\n\x18ListQueuedJobsResultItem\x12\x0c\n\x04uuid\x18\x01 \x01(\t\x12\x14\n\x0cjob_ref_name\x18\x02 \x01(\t\x12\x14\n\x0cnum_attempts\x18\x03 \x01(\x03\x12\x30\n\x0csubmitted_at\x18\x04 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x0e\n\x06status\x18\x05 \x01(\t2\xc0\x02\n\rCommandServer\x12\x35\n\nEnqueueJob\x12\x12.EnqueueJobRequest\x1a\x11.EnqueueJobResult"\x00\x12/\n\x08RetryJob\x12\x10.RetryJobRequest\x1a\x0f.RetryJobResult"\x00\x12\x38\n\x0b\x43heckHealth\x12\x13.CheckHealthRequest\x1a\x12.CheckHealthResult"\x00\x12\x41\n\x0eListQueuedJobs\x12\x16.ListQueuedJobsRequest\x1a\x15.ListQueuedJobsResult"\x00\x12J\n\x11VaccuumQueuedJobs\x12\x19.VaccuumQueuedJobsRequest\x1a\x18.VaccuumQueuedJobsResult"\x00\x62\x06proto3'
+)
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(
+    DESCRIPTOR,
+    "uncountable.integration.queue_runner.command_server.protocol.command_server_pb2",
+    _globals,
+)
+if _descriptor._USE_C_DESCRIPTORS == False:
+    DESCRIPTOR._options = None
+    _globals["_ENQUEUEJOBREQUEST"]._serialized_start = 118
+    _globals["_ENQUEUEJOBREQUEST"]._serialized_end = 187
+    _globals["_ENQUEUEJOBRESULT"]._serialized_start = 189
+    _globals["_ENQUEUEJOBRESULT"]._serialized_end = 261
+    _globals["_RETRYJOBREQUEST"]._serialized_start = 263
+    _globals["_RETRYJOBREQUEST"]._serialized_end = 294
+    _globals["_RETRYJOBRESULT"]._serialized_start = 296
+    _globals["_RETRYJOBRESULT"]._serialized_end = 366
+    _globals["_VACCUUMQUEUEDJOBSREQUEST"]._serialized_start = 368
+    _globals["_VACCUUMQUEUEDJOBSREQUEST"]._serialized_end = 394
+    _globals["_VACCUUMQUEUEDJOBSRESULT"]._serialized_start = 396
+    _globals["_VACCUUMQUEUEDJOBSRESULT"]._serialized_end = 421
+    _globals["_CHECKHEALTHREQUEST"]._serialized_start = 423
+    _globals["_CHECKHEALTHREQUEST"]._serialized_end = 443
+    _globals["_CHECKHEALTHRESULT"]._serialized_start = 445
+    _globals["_CHECKHEALTHRESULT"]._serialized_end = 481
+    _globals["_LISTQUEUEDJOBSREQUEST"]._serialized_start = 483
+    _globals["_LISTQUEUEDJOBSREQUEST"]._serialized_end = 537
+    _globals["_LISTQUEUEDJOBSRESULT"]._serialized_start = 540
+    _globals["_LISTQUEUEDJOBSRESULT"]._serialized_end = 784
+    _globals["_LISTQUEUEDJOBSRESULT_LISTQUEUEDJOBSRESULTITEM"]._serialized_start = 634
+    _globals["_LISTQUEUEDJOBSRESULT_LISTQUEUEDJOBSRESULTITEM"]._serialized_end = 784
+    _globals["_COMMANDSERVER"]._serialized_start = 787
+    _globals["_COMMANDSERVER"]._serialized_end = 1107
+# @@protoc_insertion_point(module_scope)

uncountable/integration/queue_runner/command_server/protocol/command_server_pb2.pyi

@@ -0,0 +1,114 @@
+# ruff: noqa
+from google.protobuf import timestamp_pb2 as _timestamp_pb2
+from google.protobuf.internal import containers as _containers
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import message as _message
+from typing import (
+    ClassVar as _ClassVar,
+    Iterable as _Iterable,
+    Mapping as _Mapping,
+    Optional as _Optional,
+    Union as _Union,
+)
+
+DESCRIPTOR: _descriptor.FileDescriptor
+
+class EnqueueJobRequest(_message.Message):
+    __slots__ = ("job_ref_name", "serialized_payload")
+    JOB_REF_NAME_FIELD_NUMBER: _ClassVar[int]
+    SERIALIZED_PAYLOAD_FIELD_NUMBER: _ClassVar[int]
+    job_ref_name: str
+    serialized_payload: str
+    def __init__(
+        self,
+        job_ref_name: _Optional[str] = ...,
+        serialized_payload: _Optional[str] = ...,
+    ) -> None: ...
+
+class EnqueueJobResult(_message.Message):
+    __slots__ = ("successfully_queued", "queued_job_uuid")
+    SUCCESSFULLY_QUEUED_FIELD_NUMBER: _ClassVar[int]
+    QUEUED_JOB_UUID_FIELD_NUMBER: _ClassVar[int]
+    successfully_queued: bool
+    queued_job_uuid: str
+    def __init__(
+        self, successfully_queued: bool = ..., queued_job_uuid: _Optional[str] = ...
+    ) -> None: ...
+
+class RetryJobRequest(_message.Message):
+    __slots__ = ("uuid",)
+    UUID_FIELD_NUMBER: _ClassVar[int]
+    uuid: str
+    def __init__(self, uuid: _Optional[str] = ...) -> None: ...
+
+class RetryJobResult(_message.Message):
+    __slots__ = ("successfully_queued", "queued_job_uuid")
+    SUCCESSFULLY_QUEUED_FIELD_NUMBER: _ClassVar[int]
+    QUEUED_JOB_UUID_FIELD_NUMBER: _ClassVar[int]
+    successfully_queued: bool
+    queued_job_uuid: str
+    def __init__(
+        self, successfully_queued: bool = ..., queued_job_uuid: _Optional[str] = ...
+    ) -> None: ...
+
+class VaccuumQueuedJobsRequest(_message.Message):
+    __slots__ = ()
+    def __init__(self) -> None: ...
+
+class VaccuumQueuedJobsResult(_message.Message):
+    __slots__ = ()
+    def __init__(self) -> None: ...
+
+class CheckHealthRequest(_message.Message):
+    __slots__ = ()
+    def __init__(self) -> None: ...
+
+class CheckHealthResult(_message.Message):
+    __slots__ = ("success",)
+    SUCCESS_FIELD_NUMBER: _ClassVar[int]
+    success: bool
+    def __init__(self, success: bool = ...) -> None: ...
+
+class ListQueuedJobsRequest(_message.Message):
+    __slots__ = ("offset", "limit")
+    OFFSET_FIELD_NUMBER: _ClassVar[int]
+    LIMIT_FIELD_NUMBER: _ClassVar[int]
+    offset: int
+    limit: int
+    def __init__(
+        self, offset: _Optional[int] = ..., limit: _Optional[int] = ...
+    ) -> None: ...
+
+class ListQueuedJobsResult(_message.Message):
+    __slots__ = ("queued_jobs",)
+    class ListQueuedJobsResultItem(_message.Message):
+        __slots__ = ("uuid", "job_ref_name", "num_attempts", "submitted_at", "status")
+        UUID_FIELD_NUMBER: _ClassVar[int]
+        JOB_REF_NAME_FIELD_NUMBER: _ClassVar[int]
+        NUM_ATTEMPTS_FIELD_NUMBER: _ClassVar[int]
+        SUBMITTED_AT_FIELD_NUMBER: _ClassVar[int]
+        STATUS_FIELD_NUMBER: _ClassVar[int]
+        uuid: str
+        job_ref_name: str
+        num_attempts: int
+        submitted_at: _timestamp_pb2.Timestamp
+        status: str
+        def __init__(
+            self,
+            uuid: _Optional[str] = ...,
+            job_ref_name: _Optional[str] = ...,
+            num_attempts: _Optional[int] = ...,
+            submitted_at: _Optional[_Union[_timestamp_pb2.Timestamp, _Mapping]] = ...,
+            status: _Optional[str] = ...,
+        ) -> None: ...
+
+    QUEUED_JOBS_FIELD_NUMBER: _ClassVar[int]
+    queued_jobs: _containers.RepeatedCompositeFieldContainer[
+        ListQueuedJobsResult.ListQueuedJobsResultItem
+    ]
+    def __init__(
+        self,
+        queued_jobs: _Optional[
+            _Iterable[_Union[ListQueuedJobsResult.ListQueuedJobsResultItem, _Mapping]]
+        ] = ...,
+    ) -> None: ...