UncountablePythonSDK 0.0.133__py3-none-any.whl → 0.0.136.dev0__py3-none-any.whl

examples/integration-server/jobs/materials_auto/example_cron.py

@@ -17,5 +17,5 @@ class MyCronJob(CronJob):
             if field_val.field_ref_name == "name":
                 name = field_val.value
         args.logger.log_info(f"material family found with name: {name}")
-        time.sleep(1.5)
+        time.sleep(20)
         return JobResult(success=True)

examples/integration-server/jobs/materials_auto/example_runsheet_wh.py

@@ -1,13 +1,19 @@
 from io import BytesIO
 
+from openpyxl import Workbook, load_workbook
 from uncountable.core.file_upload import DataFileUpload, FileUpload
 from uncountable.integration.job import JobArguments, RunsheetWebhookJob, register_job
 from uncountable.types import (
     download_file_t,
     entity_t,
+    export_default_runsheet_t,
     identifier_t,
     webhook_job_t,
 )
+from uncountable.types.client_base import APIRequest
+
+RUNSHEET_REF_NAME = "recipe_export_runsheet"
+RUNSHEET_REF_NAME_2 = "recipe_export_runsheet_2"
 
 
 @register_job
@@ -18,22 +24,53 @@ class StandardRunsheetGenerator(RunsheetWebhookJob):
         args: JobArguments,
         payload: webhook_job_t.RunsheetWebhookPayload,
     ) -> FileUpload:
-        args.logger.log_info("Retrieving pre-exported runsheet file from async job")
-
-        file_query = download_file_t.FileDownloadQueryEntityField(
-            entity=entity_t.EntityIdentifier(
-                type=entity_t.EntityType.ASYNC_JOB,
-                identifier_key=identifier_t.IdentifierKeyId(id=payload.async_job_id),
-            ),
-            field_key=identifier_t.IdentifierKeyRefName(
-                ref_name="unc_async_job_export_runsheet_recipe_export"
-            ),
+        args.logger.log_info("Exporting default runsheets")
+
+        entity_identifiers: list[identifier_t.IdentifierKey] = [
+            identifier_t.IdentifierKeyId(id=entity.id) for entity in payload.entities
+        ]
+
+        combined_wb = Workbook()
+        combined_sheet = combined_wb.active or combined_wb.create_sheet(
+            title="Combined Runsheet"
         )
+        combined_sheet.title = "Combined Runsheet"
+
+        for ref_name in [RUNSHEET_REF_NAME, RUNSHEET_REF_NAME_2]:
+            api_request = APIRequest(
+                method=export_default_runsheet_t.ENDPOINT_METHOD,
+                endpoint=export_default_runsheet_t.ENDPOINT_PATH,
+                args=export_default_runsheet_t.Arguments(
+                    entities=entity_identifiers,
+                    runsheet_key=identifier_t.IdentifierKeyRefName(ref_name=ref_name),
+                    entity_type=payload.entities[0].type
+                    if payload.entities
+                    else entity_t.EntityType.RECIPE,
+                ),
+            )
+
+            response = args.client.do_request(
+                api_request=api_request,
+                return_type=export_default_runsheet_t.Data,
+            )
+
+            file_query = download_file_t.FileDownloadQueryTextDocumentId(
+                text_document_id=response.text_document_id,
+            )
+
+            downloaded_files = args.client.download_files(file_query=file_query)
+            file_data = downloaded_files[0].data.read()
+
+            wb = load_workbook(filename=BytesIO(file_data))
+            for sheet_name in wb.sheetnames:
+                for row in wb[sheet_name].iter_rows(values_only=True):
+                    combined_sheet.append(row)
 
-        downloaded_files = args.client.download_files(file_query=file_query)
+        output = BytesIO()
+        combined_wb.save(output)
+        output.seek(0)
 
-        file_data = downloaded_files[0].data.read()
         return DataFileUpload(
-            data=BytesIO(file_data),
-            name=downloaded_files[0].name,
+            data=output,
+            name="combined_runsheet.xlsx",
         )

pkgs/type_spec/builder.py

@@ -423,7 +423,16 @@ class SpecTypeDefn(SpecType):
             parse_require = False
             literal = unwrap_literal_type(ptype)
             if literal is not None:
-                default = literal.value
+                if isinstance(
+                    literal.value_type, SpecTypeDefnStringEnum
+                ) and isinstance(literal.value, str):
+                    resolved_value = literal.value_type.values.get(literal.value)
+                    assert resolved_value is not None, (
+                        f"Value {literal.value} not found in enum"
+                    )
+                    default = resolved_value.value
+                else:
+                    default = literal.value
                 has_default = True
                 parse_require = True
 
@@ -1587,6 +1596,11 @@ class SpecBuilder:
                 f"'examples' in example files are expected to be a list, endpoint_path={path_details.resolved_path}"
             )
         for example in examples_data:
+            if not isinstance(example, dict):
+                raise Exception(
+                    f"each example in example file is expected to be a dict, endpoint_path={path_details.resolved_path}"
+                )
+
             arguments = example["arguments"]
             data_example = example["data"]
             if not isinstance(arguments, dict) or not isinstance(data_example, dict):

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,10 +8,14 @@ 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
 
@@ -25,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(),
@@ -39,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],
@@ -133,6 +191,7 @@ 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/queue_runner/command_server/command_client.py

@@ -6,6 +6,9 @@ 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,
@@ -18,6 +21,7 @@ from uncountable.integration.queue_runner.command_server.protocol.command_server
     VaccuumQueuedJobsResult,
 )
 from uncountable.integration.queue_runner.command_server.types import (
+    CommandCancelJobStatus,
     CommandServerBadResponse,
     CommandServerTimeout,
 )
@@ -63,6 +67,26 @@ 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,

uncountable/integration/queue_runner/command_server/command_server.py

@@ -1,4 +1,5 @@
 import asyncio
+from typing import assert_never
 
 import grpc.aio as grpc_aio
 import simplejson as json
@@ -8,6 +9,9 @@ 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,
@@ -20,6 +24,9 @@ from uncountable.integration.queue_runner.command_server.protocol.command_server
     VaccuumQueuedJobsResult,
 )
 from uncountable.integration.queue_runner.command_server.types import (
+    CommandCancelJob,
+    CommandCancelJobResponse,
+    CommandCancelJobStatus,
     CommandEnqueueJob,
     CommandEnqueueJobResponse,
     CommandQueue,
@@ -63,6 +70,33 @@ 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:

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

@@ -7,6 +7,7 @@ service CommandServer {
 	rpc CheckHealth(CheckHealthRequest) returns (CheckHealthResult) {}
 	rpc ListQueuedJobs(ListQueuedJobsRequest) returns (ListQueuedJobsResult) {}
 	rpc VaccuumQueuedJobs(VaccuumQueuedJobsRequest) returns (VaccuumQueuedJobsResult) {}
+	rpc CancelJob(CancelJobRequest) returns (CancelJobResult) {}
 }
 
 message EnqueueJobRequest {
@@ -56,3 +57,17 @@ message ListQueuedJobsResult {
 
 	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"\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'
+    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,6 +30,8 @@ _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
@@ -52,6 +54,10 @@ if _descriptor._USE_C_DESCRIPTORS == False:
     _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
+    _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]
@@ -112,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: ...

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

@@ -44,6 +44,11 @@ class CommandServerStub(object):
             request_serializer=uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.VaccuumQueuedJobsRequest.SerializeToString,
             response_deserializer=uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.VaccuumQueuedJobsResult.FromString,
         )
+        self.CancelJob = channel.unary_unary(
+            "/CommandServer/CancelJob",
+            request_serializer=uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.CancelJobRequest.SerializeToString,
+            response_deserializer=uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.CancelJobResult.FromString,
+        )
 
 
 class CommandServerServicer(object):
@@ -79,6 +84,12 @@ class CommandServerServicer(object):
         context.set_details("Method not implemented!")
         raise NotImplementedError("Method not implemented!")
 
+    def CancelJob(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details("Method not implemented!")
+        raise NotImplementedError("Method not implemented!")
+
 
 def add_CommandServerServicer_to_server(servicer, server):
     rpc_method_handlers = {
@@ -107,6 +118,11 @@ def add_CommandServerServicer_to_server(servicer, server):
             request_deserializer=uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.VaccuumQueuedJobsRequest.FromString,
             response_serializer=uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.VaccuumQueuedJobsResult.SerializeToString,
         ),
+        "CancelJob": grpc.unary_unary_rpc_method_handler(
+            servicer.CancelJob,
+            request_deserializer=uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.CancelJobRequest.FromString,
+            response_serializer=uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.CancelJobResult.SerializeToString,
+        ),
     }
     generic_handler = grpc.method_handlers_generic_handler(
         "CommandServer", rpc_method_handlers
@@ -262,3 +278,32 @@ class CommandServer(object):
             timeout,
             metadata,
         )
+
+    @staticmethod
+    def CancelJob(
+        request,
+        target,
+        options=(),
+        channel_credentials=None,
+        call_credentials=None,
+        insecure=False,
+        compression=None,
+        wait_for_ready=None,
+        timeout=None,
+        metadata=None,
+    ):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            "/CommandServer/CancelJob",
+            uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.CancelJobRequest.SerializeToString,
+            uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.CancelJobResult.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+        )

uncountable/integration/queue_runner/command_server/types.py

@@ -6,10 +6,17 @@ from enum import StrEnum
 from uncountable.types import queued_job_t
 
 
+class CommandCancelJobStatus(StrEnum):
+    CANCELLED_WITH_RESTART = "cancelled_with_restart"
+    NO_JOB_FOUND = "no_job_found"
+    JOB_ALREADY_COMPLETED = "job_already_completed"
+
+
 class CommandType(StrEnum):
     ENQUEUE_JOB = "enqueue_job"
     RETRY_JOB = "retry_job"
     VACCUUM_QUEUED_JOBS = "vaccuum_queued_jobs"
+    CANCEL_JOB = "cancel_job"
 
 
 RT = typing.TypeVar("RT")
@@ -55,7 +62,20 @@ class CommandVaccuumQueuedJobs(CommandBase[CommandVaccuumQueuedJobsResponse]):
     type: CommandType = CommandType.VACCUUM_QUEUED_JOBS
 
 
-_Command = CommandEnqueueJob | CommandRetryJob | CommandVaccuumQueuedJobs
+@dataclass(kw_only=True)
+class CommandCancelJobResponse:
+    status: CommandCancelJobStatus
+
+
+@dataclass(kw_only=True)
+class CommandCancelJob(CommandBase[CommandCancelJobResponse]):
+    type: CommandType = CommandType.CANCEL_JOB
+    queued_job_uuid: str
+
+
+_Command = (
+    CommandEnqueueJob | CommandRetryJob | CommandVaccuumQueuedJobs | CommandCancelJob
+)
 
 
 CommandQueue = asyncio.Queue[_Command]

uncountable/integration/queue_runner/datastore/datastore_sqlite.py

@@ -234,6 +234,31 @@ class DatastoreSqlite(Datastore):
 
             return queued_jobs
 
+    def get_queued_job(self, *, uuid: str) -> queued_job_t.QueuedJob | None:
+        with self.session_maker() as session:
+            select_stmt = select(
+                QueuedJob.id,
+                QueuedJob.payload,
+                QueuedJob.num_attempts,
+                QueuedJob.job_ref_name,
+                QueuedJob.status,
+                QueuedJob.submitted_at,
+            ).filter(QueuedJob.id == uuid)
+
+            row = session.execute(select_stmt).one_or_none()
+            return (
+                queued_job_t.QueuedJob(
+                    queued_job_uuid=row.id,
+                    job_ref_name=row.job_ref_name,
+                    num_attempts=row.num_attempts,
+                    status=row.status or queued_job_t.JobStatus.QUEUED,
+                    submitted_at=row.submitted_at,
+                    payload=queued_job_payload_parser.parse_storage(row.payload),
+                )
+                if row is not None
+                else None
+            )
+
     def vaccuum_queued_jobs(self) -> None:
         with self.session_maker() as session:
             delete_stmt = (

uncountable/integration/queue_runner/datastore/interface.py

@@ -27,3 +27,6 @@ class Datastore(ABC):
     def list_queued_job_metadata(
         self, offset: int, limit: int | None
     ) -> list[queued_job_t.QueuedJobMetadata]: ...
+
+    @abstractmethod
+    def get_queued_job(self, *, uuid: str) -> queued_job_t.QueuedJob | None: ...

uncountable/integration/queue_runner/job_scheduler.py

@@ -1,5 +1,7 @@
 import asyncio
+import os
 import sys
+import threading
 import typing
 from concurrent.futures import ProcessPoolExecutor
 from dataclasses import dataclass
@@ -15,6 +17,9 @@ from uncountable.integration.queue_runner.command_server import (
     CommandTask,
 )
 from uncountable.integration.queue_runner.command_server.types import (
+    CommandCancelJob,
+    CommandCancelJobResponse,
+    CommandCancelJobStatus,
     CommandVaccuumQueuedJobs,
 )
 from uncountable.integration.queue_runner.datastore import DatastoreSqlite
@@ -24,7 +29,7 @@ from uncountable.integration.scan_profiles import load_profiles
 from uncountable.integration.telemetry import Logger
 from uncountable.types import job_definition_t, queued_job_t
 
-from .types import ResultQueue, ResultTask
+from .types import RESTART_EXIT_CODE, ResultQueue, ResultTask
 
 _MAX_JOB_WORKERS = 5
 
@@ -142,6 +147,34 @@ async def start_scheduler(
                 CommandEnqueueJobResponse(queued_job_uuid=queued_job_uuid)
             )
 
+        async def _handle_cancel_job_command(command: CommandCancelJob) -> None:
+            queued_job = datastore.get_queued_job(uuid=command.queued_job_uuid)
+            if queued_job is None:
+                await command.response_queue.put(
+                    CommandCancelJobResponse(status=CommandCancelJobStatus.NO_JOB_FOUND)
+                )
+                return
+
+            if queued_job.status == queued_job_t.JobStatus.QUEUED:
+                datastore.remove_job_from_queue(command.queued_job_uuid)
+                await command.response_queue.put(
+                    CommandCancelJobResponse(
+                        status=CommandCancelJobStatus.CANCELLED_WITH_RESTART
+                    )
+                )
+
+                def delayed_exit() -> None:
+                    os._exit(RESTART_EXIT_CODE)
+
+                threading.Timer(interval=5, function=delayed_exit).start()
+
+            else:
+                await command.response_queue.put(
+                    CommandCancelJobResponse(
+                        status=CommandCancelJobStatus.JOB_ALREADY_COMPLETED
+                    )
+                )
+
         async def _handle_retry_job_command(command: CommandRetryJob) -> None:
             queued_job = datastore.retry_job(command.queued_job_uuid)
             if queued_job is None:
@@ -181,6 +214,8 @@ async def start_scheduler(
                             await _handle_retry_job_command(command=command)
                         case CommandVaccuumQueuedJobs():
                             _handle_vaccuum_queued_jobs_command(command=command)
+                        case CommandCancelJob():
+                            await _handle_cancel_job_command(command=command)
                         case _:
                             typing.assert_never(command)
                     command_task = asyncio.create_task(command_queue.get())

uncountable/integration/queue_runner/types.py

@@ -5,3 +5,5 @@ from uncountable.types import queued_job_t
 ListenQueue = Queue[queued_job_t.QueuedJob]
 ResultQueue = Queue[queued_job_t.QueuedJobResult]
 ResultTask = Task[queued_job_t.QueuedJobResult]
+
+RESTART_EXIT_CODE = 147