UncountablePythonSDK 0.0.126__py3-none-any.whl → 0.0.142.dev0__py3-none-any.whl

uncountable/integration/cli.py

@@ -1,4 +1,6 @@
 import argparse
+import json
+from typing import assert_never
 
 from dateutil import tz
 from opentelemetry.trace import get_current_span
@@ -6,8 +8,13 @@ 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_cancellation_message,
     send_job_queue_message,
     send_list_queued_jobs_message,
+    send_retry_job_message,
+)
+from uncountable.integration.queue_runner.command_server.types import (
+    CommandCancelJobStatus,
 )
 from uncountable.integration.telemetry import Logger
 from uncountable.types import queued_job_t
@@ -24,12 +31,28 @@ def register_enqueue_job_parser(
         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")
+    run_parser.add_argument(
+        "--payload", type=str, help="JSON payload for webhook invocation context"
+    )
 
     def _handle_enqueue_job(args: argparse.Namespace) -> None:
+        invocation_context: queued_job_t.InvocationContext
+
+        if args.payload is not None:
+            try:
+                webhook_payload = json.loads(args.payload)
+                invocation_context = queued_job_t.InvocationContextWebhook(
+                    webhook_payload=webhook_payload
+                )
+            except json.JSONDecodeError as e:
+                raise ValueError(f"Invalid JSON payload: {e}")
+        else:
+            invocation_context = queued_job_t.InvocationContextManual()
+
         send_job_queue_message(
             job_ref_name=args.job_id,
             payload=queued_job_t.QueuedJobPayload(
-                invocation_context=queued_job_t.InvocationContextManual()
+                invocation_context=invocation_context
             ),
             host=args.host,
             port=get_local_admin_server_port(),
@@ -38,6 +61,42 @@ def register_enqueue_job_parser(
     run_parser.set_defaults(func=_handle_enqueue_job)
 
 
+def register_cancel_queued_job_parser(
+    sub_parser_manager: argparse._SubParsersAction,
+    parents: list[argparse.ArgumentParser],
+) -> None:
+    cancel_parser = sub_parser_manager.add_parser(
+        "cancel",
+        parents=parents,
+        help="Cancel a queued job with a given host and queued job UUID",
+        description="Cancel a job with a given host and queued job UUID",
+    )
+    cancel_parser.add_argument(
+        "uuid", type=str, help="The UUID of the queued job to cancel"
+    )
+
+    def _handle_cancel_queued_job(args: argparse.Namespace) -> None:
+        resp = send_job_cancellation_message(
+            queued_job_uuid=args.uuid,
+            host=args.host,
+            port=get_local_admin_server_port(),
+        )
+
+        match resp:
+            case CommandCancelJobStatus.CANCELLED_WITH_RESTART:
+                print(
+                    "Job successfully cancelled. The integration server will restart."
+                )
+            case CommandCancelJobStatus.NO_JOB_FOUND:
+                print("Job not found.")
+            case CommandCancelJobStatus.JOB_ALREADY_COMPLETED:
+                print("Job already completed.")
+            case _:
+                assert_never(resp)
+
+    cancel_parser.set_defaults(func=_handle_cancel_queued_job)
+
+
 def register_list_queued_jobs(
     sub_parser_manager: argparse._SubParsersAction,
     parents: list[argparse.ArgumentParser],
@@ -70,12 +129,13 @@ def register_list_queued_jobs(
             port=get_local_admin_server_port(),
         )
 
-        headers = ["UUID", "Job Ref Name", "Attempts", "Submitted At"]
+        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
@@ -85,6 +145,31 @@ def register_list_queued_jobs(
     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"
+    )
+
+    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)
+
+
 def main() -> None:
     logger = Logger(get_current_span())
 
@@ -104,7 +189,9 @@ def main() -> None:
     )
 
     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])
+    register_cancel_queued_job_parser(subparser_action, parents=[base_parser])
 
     args = main_parser.parse_args()
     with logger.push_scope(args.command):

uncountable/integration/executors/executors.py

@@ -88,7 +88,6 @@ def execute_job(
     job_definition: job_definition_t.JobDefinition,
     profile_metadata: job_definition_t.ProfileMetadata,
     args: JobArguments,
-    job_uuid: str,
 ) -> job_definition_t.JobResult:
     with args.logger.push_scope(job_definition.name) as job_logger:
         job = resolve_executor(job_definition.executor, profile_metadata)
@@ -104,7 +103,7 @@ def execute_job(
                 run_entity = _create_run_entity(
                     client=args.client,
                     logging_settings=job_definition.logging_settings,
-                    job_uuid=job_uuid,
+                    job_uuid=args.job_uuid,
                 )
             result = job.run_outer(args=args)
         except Exception as e:

uncountable/integration/executors/generic_upload_executor.py

@@ -41,7 +41,7 @@ def _get_extension(filename: str) -> str | None:
 
 def _run_keyword_detection(data: io.BytesIO, keyword: str) -> bool:
     try:
-        text = io.TextIOWrapper(data)
+        text = io.TextIOWrapper(data, encoding="utf-8")
         for line in text:
             if (
                 keyword in line

uncountable/integration/job.py

@@ -28,7 +28,6 @@ from uncountable.integration.secret_retrieval.retrieve_secret import retrieve_se
 from uncountable.integration.telemetry import JobLogger
 from uncountable.types import (
     base_t,
-    entity_t,
     job_definition_t,
     queued_job_t,
     webhook_job_t,
@@ -49,6 +48,7 @@ class JobArguments:
     batch_processor: AsyncBatchProcessor
     logger: JobLogger
     payload: base_t.JsonValue
+    job_uuid: str
 
 
 # only for compatibility:
@@ -254,13 +254,13 @@ class RunsheetWebhookJob(WebhookJob[webhook_job_t.RunsheetWebhookPayload]):
         self,
         *,
         args: JobArguments,
-        entities: list[entity_t.Entity],
+        payload: webhook_job_t.RunsheetWebhookPayload,
     ) -> FileUpload: ...
 
     def run(
         self, args: JobArguments, payload: webhook_job_t.RunsheetWebhookPayload
     ) -> JobResult:
-        runsheet = self.build_runsheet(args=args, entities=payload.entities)
+        runsheet = self.build_runsheet(args=args, payload=payload)
 
         files = args.client.upload_files(file_uploads=[runsheet])
         args.client.complete_async_upload(

uncountable/integration/queue_runner/command_server/__init__.py

@@ -4,6 +4,8 @@ from .types import (
     CommandEnqueueJob,
     CommandEnqueueJobResponse,
     CommandQueue,
+    CommandRetryJob,
+    CommandRetryJobResponse,
     CommandServerBadResponse,
     CommandServerException,
     CommandServerTimeout,
@@ -16,6 +18,8 @@ __all__: list[str] = [
     "send_job_queue_message",
     "CommandEnqueueJob",
     "CommandEnqueueJobResponse",
+    "CommandRetryJob",
+    "CommandRetryJobResponse",
     "CommandTask",
     "CommandQueue",
     "CommandServerTimeout",

uncountable/integration/queue_runner/command_server/command_client.py

@@ -6,14 +6,22 @@ import simplejson as json
 
 from pkgs.serialization_util import serialize_for_api
 from uncountable.integration.queue_runner.command_server.protocol.command_server_pb2 import (
+    CancelJobRequest,
+    CancelJobResult,
+    CancelJobStatus,
     CheckHealthRequest,
     CheckHealthResult,
     EnqueueJobRequest,
     EnqueueJobResult,
     ListQueuedJobsRequest,
     ListQueuedJobsResult,
+    RetryJobRequest,
+    RetryJobResult,
+    VaccuumQueuedJobsRequest,
+    VaccuumQueuedJobsResult,
 )
 from uncountable.integration.queue_runner.command_server.types import (
+    CommandCancelJobStatus,
     CommandServerBadResponse,
     CommandServerTimeout,
 )
@@ -59,6 +67,46 @@ def send_job_queue_message(
         return response.queued_job_uuid
 
 
+def send_job_cancellation_message(
+    *, queued_job_uuid: str, host: str = "localhost", port: int
+) -> CommandCancelJobStatus:
+    with command_server_connection(host=host, port=port) as stub:
+        request = CancelJobRequest(job_uuid=queued_job_uuid)
+
+        response = stub.CancelJob(request, timeout=_DEFAULT_MESSAGE_TIMEOUT_SECS)
+
+        assert isinstance(response, CancelJobResult)
+        match response.status:
+            case CancelJobStatus.NO_JOB_FOUND:
+                return CommandCancelJobStatus.NO_JOB_FOUND
+            case CancelJobStatus.CANCELLED_WITH_RESTART:
+                return CommandCancelJobStatus.CANCELLED_WITH_RESTART
+            case CancelJobStatus.JOB_ALREADY_COMPLETED:
+                return CommandCancelJobStatus.JOB_ALREADY_COMPLETED
+            case _:
+                raise CommandServerBadResponse(f"unknown status: {response.status}")
+
+
+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()
@@ -92,3 +140,18 @@ def send_list_queued_jobs_message(
 
         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

@@ -1,23 +1,39 @@
 import asyncio
+from typing import assert_never
 
+import grpc.aio as grpc_aio
 import simplejson as json
 from google.protobuf.timestamp_pb2 import Timestamp
-from grpc import StatusCode, aio
+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 (
+    CancelJobRequest,
+    CancelJobResult,
+    CancelJobStatus,
     CheckHealthRequest,
     CheckHealthResult,
     EnqueueJobRequest,
     EnqueueJobResult,
     ListQueuedJobsRequest,
     ListQueuedJobsResult,
+    RetryJobRequest,
+    RetryJobResult,
+    VaccuumQueuedJobsRequest,
+    VaccuumQueuedJobsResult,
 )
 from uncountable.integration.queue_runner.command_server.types import (
+    CommandCancelJob,
+    CommandCancelJobResponse,
+    CommandCancelJobStatus,
     CommandEnqueueJob,
     CommandEnqueueJobResponse,
     CommandQueue,
+    CommandRetryJob,
+    CommandRetryJobResponse,
+    CommandVaccuumQueuedJobs,
+    CommandVaccuumQueuedJobsResponse,
 )
 from uncountable.integration.queue_runner.datastore import DatastoreSqlite
 from uncountable.types import queued_job_t
@@ -32,11 +48,11 @@ queued_job_payload_parser = CachedParser(queued_job_t.QueuedJobPayload)
 
 
 async def serve(command_queue: CommandQueue, datastore: DatastoreSqlite) -> None:
-    server = aio.server()
+    server = grpc_aio.server()
 
     class CommandServerHandler(CommandServerServicer):
         async def EnqueueJob(
-            self, request: EnqueueJobRequest, context: aio.ServicerContext
+            self, request: EnqueueJobRequest, context: grpc_aio.ServicerContext
         ) -> EnqueueJobResult:
             payload_json = json.loads(request.serialized_payload)
             payload = queued_job_payload_parser.parse_api(payload_json)
@@ -54,13 +70,57 @@ async def serve(command_queue: CommandQueue, datastore: DatastoreSqlite) -> None
             )
             return result
 
+        async def CancelJob(
+            self, request: CancelJobRequest, context: grpc_aio.ServicerContext
+        ) -> CancelJobResult:
+            response_queue: asyncio.Queue[CommandCancelJobResponse] = asyncio.Queue()
+            await command_queue.put(
+                CommandCancelJob(
+                    queued_job_uuid=request.job_uuid,
+                    response_queue=response_queue,
+                )
+            )
+
+            response = await response_queue.get()
+
+            proto_status: CancelJobStatus
+            match response.status:
+                case CommandCancelJobStatus.NO_JOB_FOUND:
+                    proto_status = CancelJobStatus.NO_JOB_FOUND
+                case CommandCancelJobStatus.CANCELLED_WITH_RESTART:
+                    proto_status = CancelJobStatus.CANCELLED_WITH_RESTART
+                case CommandCancelJobStatus.JOB_ALREADY_COMPLETED:
+                    proto_status = CancelJobStatus.JOB_ALREADY_COMPLETED
+                case _:
+                    assert_never(response.status)
+
+            result = CancelJobResult(status=proto_status)
+            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: aio.ServicerContext
+            self, request: CheckHealthRequest, context: grpc_aio.ServicerContext
         ) -> CheckHealthResult:
             return CheckHealthResult(success=True)
 
         async def ListQueuedJobs(
-            self, request: ListQueuedJobsRequest, context: aio.ServicerContext
+            self, request: ListQueuedJobsRequest, context: grpc_aio.ServicerContext
         ) -> ListQueuedJobsResult:
             if (
                 request.limit < ListQueuedJobsConstants.LIMIT_MIN
@@ -90,10 +150,22 @@ async def serve(command_queue: CommandQueue, datastore: DatastoreSqlite) -> None
                         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()}"

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

@@ -3,8 +3,11 @@ 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) {}
+	rpc CancelJob(CancelJobRequest) returns (CancelJobResult) {}
 }
 
 message EnqueueJobRequest {
@@ -17,6 +20,21 @@ message EnqueueJobResult {
 	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 {
@@ -34,7 +52,22 @@ message ListQueuedJobsResult {
 	string job_ref_name = 2;
 	int64 num_attempts = 3;
 	google.protobuf.Timestamp submitted_at = 4;
+	string status = 5;
 	}
 
 	repeated ListQueuedJobsResultItem queued_jobs  = 1;
 }
+
+message CancelJobRequest {
+        string job_uuid = 1;
+}
+
+enum CancelJobStatus {
+	CANCELLED_WITH_RESTART = 0;
+	NO_JOB_FOUND = 1;
+	JOB_ALREADY_COMPLETED = 2;
+}
+
+message CancelJobResult {
+	CancelJobStatus status = 1;
+}

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

@@ -18,7 +18,7 @@ from google.protobuf import timestamp_pb2 as google_dot_protobuf_dot_timestamp__
 
 
 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"\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"\xe4\x01\n\x14ListQueuedJobsResult\x12\x43\n\x0bqueued_jobs\x18\x01 \x03(\x0b\x32..ListQueuedJobsResult.ListQueuedJobsResultItem\x1a\x86\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.Timestamp2\xc3\x01\n\rCommandServer\x12\x35\n\nEnqueueJob\x12\x12.EnqueueJobRequest\x1a\x11.EnqueueJobResult"\x00\x12\x38\n\x0b\x43heckHealth\x12\x13.CheckHealthRequest\x1a\x12.CheckHealthResult"\x00\x12\x41\n\x0eListQueuedJobs\x12\x16.ListQueuedJobsRequest\x1a\x15.ListQueuedJobsResult"\x00\x62\x06proto3'
+    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(\t"$\n\x10\x43\x61ncelJobRequest\x12\x10\n\x08job_uuid\x18\x01 \x01(\t"3\n\x0f\x43\x61ncelJobResult\x12 \n\x06status\x18\x01 \x01(\x0e\x32\x10.CancelJobStatus*Z\n\x0f\x43\x61ncelJobStatus\x12\x1a\n\x16\x43\x41NCELLED_WITH_RESTART\x10\x00\x12\x10\n\x0cNO_JOB_FOUND\x10\x01\x12\x19\n\x15JOB_ALREADY_COMPLETED\x10\x02\x32\xf4\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\x12\x32\n\tCancelJob\x12\x11.CancelJobRequest\x1a\x10.CancelJobResult"\x00\x62\x06proto3'
 )
 
 _globals = globals()
@@ -30,20 +30,34 @@ _builder.BuildTopDescriptorsAndMessages(
 )
 if _descriptor._USE_C_DESCRIPTORS == False:
     DESCRIPTOR._options = None
+    _globals["_CANCELJOBSTATUS"]._serialized_start = 877
+    _globals["_CANCELJOBSTATUS"]._serialized_end = 967
     _globals["_ENQUEUEJOBREQUEST"]._serialized_start = 118
     _globals["_ENQUEUEJOBREQUEST"]._serialized_end = 187
     _globals["_ENQUEUEJOBRESULT"]._serialized_start = 189
     _globals["_ENQUEUEJOBRESULT"]._serialized_end = 261
-    _globals["_CHECKHEALTHREQUEST"]._serialized_start = 263
-    _globals["_CHECKHEALTHREQUEST"]._serialized_end = 283
-    _globals["_CHECKHEALTHRESULT"]._serialized_start = 285
-    _globals["_CHECKHEALTHRESULT"]._serialized_end = 321
-    _globals["_LISTQUEUEDJOBSREQUEST"]._serialized_start = 323
-    _globals["_LISTQUEUEDJOBSREQUEST"]._serialized_end = 377
-    _globals["_LISTQUEUEDJOBSRESULT"]._serialized_start = 380
-    _globals["_LISTQUEUEDJOBSRESULT"]._serialized_end = 608
-    _globals["_LISTQUEUEDJOBSRESULT_LISTQUEUEDJOBSRESULTITEM"]._serialized_start = 474
-    _globals["_LISTQUEUEDJOBSRESULT_LISTQUEUEDJOBSRESULTITEM"]._serialized_end = 608
-    _globals["_COMMANDSERVER"]._serialized_start = 611
-    _globals["_COMMANDSERVER"]._serialized_end = 806
+    _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["_CANCELJOBREQUEST"]._serialized_start = 786
+    _globals["_CANCELJOBREQUEST"]._serialized_end = 822
+    _globals["_CANCELJOBRESULT"]._serialized_start = 824
+    _globals["_CANCELJOBRESULT"]._serialized_end = 875
+    _globals["_COMMANDSERVER"]._serialized_start = 970
+    _globals["_COMMANDSERVER"]._serialized_end = 1342
 # @@protoc_insertion_point(module_scope)

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

@@ -1,6 +1,7 @@
 # ruff: noqa
 from google.protobuf import timestamp_pb2 as _timestamp_pb2
 from google.protobuf.internal import containers as _containers
+from google.protobuf.internal import enum_type_wrapper as _enum_type_wrapper
 from google.protobuf import descriptor as _descriptor
 from google.protobuf import message as _message
 from typing import (
@@ -13,6 +14,16 @@ from typing import (
 
 DESCRIPTOR: _descriptor.FileDescriptor
 
+class CancelJobStatus(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
+    __slots__ = ()
+    CANCELLED_WITH_RESTART: _ClassVar[CancelJobStatus]
+    NO_JOB_FOUND: _ClassVar[CancelJobStatus]
+    JOB_ALREADY_COMPLETED: _ClassVar[CancelJobStatus]
+
+CANCELLED_WITH_RESTART: CancelJobStatus
+NO_JOB_FOUND: CancelJobStatus
+JOB_ALREADY_COMPLETED: CancelJobStatus
+
 class EnqueueJobRequest(_message.Message):
     __slots__ = ("job_ref_name", "serialized_payload")
     JOB_REF_NAME_FIELD_NUMBER: _ClassVar[int]
@@ -35,6 +46,30 @@ class EnqueueJobResult(_message.Message):
         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: ...
@@ -58,21 +93,24 @@ class ListQueuedJobsRequest(_message.Message):
 class ListQueuedJobsResult(_message.Message):
     __slots__ = ("queued_jobs",)
     class ListQueuedJobsResultItem(_message.Message):
-        __slots__ = ("uuid", "job_ref_name", "num_attempts", "submitted_at")
+        __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]
@@ -85,3 +123,17 @@ class ListQueuedJobsResult(_message.Message):
             _Iterable[_Union[ListQueuedJobsResult.ListQueuedJobsResultItem, _Mapping]]
         ] = ...,
     ) -> None: ...
+
+class CancelJobRequest(_message.Message):
+    __slots__ = ("job_uuid",)
+    JOB_UUID_FIELD_NUMBER: _ClassVar[int]
+    job_uuid: str
+    def __init__(self, job_uuid: _Optional[str] = ...) -> None: ...
+
+class CancelJobResult(_message.Message):
+    __slots__ = ("status",)
+    STATUS_FIELD_NUMBER: _ClassVar[int]
+    status: CancelJobStatus
+    def __init__(
+        self, status: _Optional[_Union[CancelJobStatus, str]] = ...
+    ) -> None: ...