UncountablePythonSDK 0.0.189__py3-none-any.whl → 0.0.191__py3-none-any.whl

pkgs/type_spec/AGENTS.md

@@ -0,0 +1,13 @@
+<!-- DO NOT EDIT: this file is generated by `docs/tools/compile.py`.
+Edit fragments under `docs/fragments/` or manifests under `docs/manifests/`,
+then run `just docs compile` to regenerate. -->
+
+## Keep feature-specific mechanisms out of the type_spec type system unless universally applicable
+
+Do not extend the type_spec type system with a mechanism scoped to one feature's view of how types
+serialize or behave — for example, custom key conversions added in preparation for LLM/JSON-schema
+tool responses — unless it has been proven universally applicable across type_spec's consumers.
+type_spec is a generic type-definition system shared by the backend, frontend, and SDK, so tying it
+to one view of a feature (such as LLM features) couples the generic types too closely to that single
+consumer. Keep a capability that only one feature needs at that feature's own layer, and promote it
+into the shared type_spec builder/emitter only once multiple consumers demonstrably need it.

pkgs/type_spec/CLAUDE.md

@@ -0,0 +1,5 @@
+<!-- DO NOT EDIT: this file is generated by `docs/tools/compile.py`.
+Edit fragments under `docs/fragments/` or manifests under `docs/manifests/`,
+then run `just docs compile` to regenerate. -->
+
+@AGENTS.md

pkgs/type_spec/emit_typescript_util.py

@@ -175,6 +175,27 @@ def emit_constant_ts(ctx: EmitTypescriptContext, sconst: builder.SpecConstant) -
     ctx.out.write(f"export const {const_name} = {value}\n")
 
 
+def _collect_object_properties(
+    stype: builder.SpecTypeDefnObject,
+) -> dict[str, builder.SpecProperty]:
+    base_properties = (
+        _collect_object_properties(stype.base) if stype.base is not None else {}
+    )
+    return {**base_properties, **(stype.properties or {})}
+
+
+def _type_info_ref_type_ts(*, stype: builder.SpecTypeDefnObject, base_ref: str) -> str:
+    defaulted_names = sorted(
+        ts_property_name(prop)
+        for prop in _collect_object_properties(stype).values()
+        if prop.has_default and not prop.parse_require and prop.default is not None
+    )
+    if len(defaulted_names) == 0:
+        return f"{base_ref}.TypeInfoRef<{stype.name}>"
+    defaulted_keys = " | ".join(f'"{name}"' for name in defaulted_names)
+    return f"{base_ref}.TypeInfoRef<{stype.name}, {base_ref}.WithDefaultsApplied<{stype.name}, {defaulted_keys}>>"
+
+
 def emit_type_info_ref_ts(
     ctx: EmitTypescriptContext,
     stype: builder.SpecTypeDefn,
@@ -184,9 +205,9 @@ def emit_type_info_ref_ts(
     assert ctx.namespace.emit_type_info, (
         f"{ctx.namespace.name}: $emit_type_info_refs requires $emit_type_info"
     )
-    assert (
-        isinstance(stype, builder.SpecTypeDefnObject) and len(stype.parameters) == 0
-    ), (
+    if not isinstance(stype, builder.SpecTypeDefnObject):
+        return
+    assert len(stype.parameters) == 0, (
         f"{ctx.namespace.name}.{stype.name}: type info refs support non-generic objects only"
     )
 
@@ -195,7 +216,7 @@ def emit_type_info_ref_ts(
     ctx.out.write("\n")
     ctx.out.write(MODIFY_NOTICE)
     ctx.out.write(
-        f"export const {stype.name}TypeInfoRef: {base_ref}.TypeInfoRef<{stype.name}> = {{\n"
+        f"export const {stype.name}TypeInfoRef: {_type_info_ref_type_ts(stype=stype, base_ref=base_ref)} = {{\n"
     )
     ctx.out.write(f'{INDENT}typeName: "{ctx.namespace.name}.{stype.name}",\n')
     ctx.out.write("}\n")

pkgs/type_spec/parts/base.ts.prepart

@@ -40,11 +40,16 @@ declare const brand: unique symbol
 export type Branded<TBase, TBrand extends string> = TBase & { readonly [brand]: TBrand }
 
 declare const typeInfoValue: unique symbol
-export interface TypeInfoRef<T> {
+declare const typeInfoParsedValue: unique symbol
+export interface TypeInfoRef<T, TParsed extends T = T> {
   readonly typeName: string
   readonly [typeInfoValue]?: T
+  readonly [typeInfoParsedValue]?: TParsed
 }
 
+export type WithDefaultsApplied<T, TDefaultedKey extends keyof T> = Omit<T, TDefaultedKey> &
+  Required<Pick<T, TDefaultedKey>>
+
 // Ids matching a strict integer number are converted to integers
 export const ID_REGEX = /^-?[1-9][0-9]{0,20}$/
 // Matches ISO 8601 dates. Stricter than `Date.parse`.

uncountable/core/__init__.py

@@ -1,10 +1,11 @@
 from .async_batch import AsyncBatchProcessor
 from .client import Client
 from .file_upload import MediaFileUpload, UploadedFile
-from .types import AuthDetailsApiKey, AuthDetailsOAuth
+from .types import AuthDetailsApiKey, AuthDetailsBearerToken, AuthDetailsOAuth
 
 __all__: list[str] = [
     "AuthDetailsApiKey",
+    "AuthDetailsBearerToken",
     "AuthDetailsOAuth",
     "AsyncBatchProcessor",
     "Client",

uncountable/core/client.py

@@ -31,7 +31,12 @@ from uncountable.types.request_headers_t import (
 )
 
 from .file_upload import FileUpload, FileUploader, UploadedFile
-from .types import AuthDetailsAll, AuthDetailsApiKey, AuthDetailsOAuth
+from .types import (
+    AuthDetailsAll,
+    AuthDetailsApiKey,
+    AuthDetailsBearerToken,
+    AuthDetailsOAuth,
+)
 
 DT = typing.TypeVar("DT")
 
@@ -351,6 +356,8 @@ class Client(ClientMethods):
             case AuthDetailsOAuth():
                 token = self._get_oauth_bearer_token(oauth_details=self._auth_details)
                 return {"Authorization": f"Bearer {token}"}
+            case AuthDetailsBearerToken():
+                return {"Authorization": f"Bearer {self._auth_details.token}"}
         typing.assert_never(self._auth_details)
 
     def _build_http_request(

uncountable/core/types.py

@@ -13,5 +13,10 @@ class AuthDetailsOAuth:
     scope: str = "unc.rnd"
 
 
+@dataclass(kw_only=True)
+class AuthDetailsBearerToken:
+    token: str
+
+
 AuthDetails = AuthDetailsApiKey  # Legacy Mapping
-AuthDetailsAll = AuthDetailsApiKey | AuthDetailsOAuth
+AuthDetailsAll = AuthDetailsApiKey | AuthDetailsOAuth | AuthDetailsBearerToken

uncountable/integration/cli.py

@@ -89,6 +89,8 @@ def register_cancel_queued_job_parser(
                 print(
                     "Job successfully cancelled. The integration server will restart."
                 )
+            case CommandCancelJobStatus.CANCELLED_QUEUED:
+                print("Queued job successfully cancelled.")
             case CommandCancelJobStatus.NO_JOB_FOUND:
                 print("Job not found.")
             case CommandCancelJobStatus.JOB_ALREADY_COMPLETED:

uncountable/integration/job.py

@@ -10,7 +10,6 @@ from pkgs.argument_parser import CachedParser
 from pkgs.serialization_util import serialize_for_api
 from uncountable.core.async_batch import AsyncBatchProcessor
 from uncountable.core.client import Client
-from uncountable.core.environment import get_local_admin_server_port
 from uncountable.core.file_upload import FileUpload
 from uncountable.core.types import AuthDetailsOAuth
 from uncountable.integration.http_server import (
@@ -18,12 +17,7 @@ from uncountable.integration.http_server import (
     GenericHttpResponse,
     HttpException,
 )
-from uncountable.integration.queue_runner.command_server.command_client import (
-    send_job_queue_message,
-)
-from uncountable.integration.queue_runner.command_server.types import (
-    CommandServerException,
-)
+from uncountable.integration.job_queue import JobQueue
 from uncountable.integration.request_context import extract_request_context
 from uncountable.integration.secret_retrieval import retrieve_secret
 from uncountable.integration.telemetry import JobLogger
@@ -111,11 +105,7 @@ class _RequestValidatorClient(Client):
         return self._oauth_bearer_token
 
 
-class CustomHttpJob(Job[GenericHttpRequest]):
-    @property
-    def payload_type(self) -> type[GenericHttpRequest]:
-        return GenericHttpRequest
-
+class HttpJob(ABC):
     @staticmethod
     @abstractmethod
     def validate_request(
@@ -130,6 +120,22 @@ class CustomHttpJob(Job[GenericHttpRequest]):
         """
         ...
 
+    @abstractmethod
+    def respond_to_request(
+        self,
+        *,
+        request: GenericHttpRequest,
+        job_definition: HttpJobDefinitionBase,
+        profile_meta: ProfileMetadata,
+        job_queue: JobQueue,
+    ) -> GenericHttpResponse: ...
+
+
+class CustomHttpJob(Job[GenericHttpRequest], HttpJob):
+    @property
+    def payload_type(self) -> type[GenericHttpRequest]:
+        return GenericHttpRequest
+
     @staticmethod
     def get_validated_oauth_request_user_id(
         *, profile_metadata: ProfileMetadata, request: GenericHttpRequest
@@ -162,6 +168,19 @@ class CustomHttpJob(Job[GenericHttpRequest]):
         """
         ...
 
+    @typing.override
+    def respond_to_request(
+        self,
+        *,
+        request: GenericHttpRequest,
+        job_definition: HttpJobDefinitionBase,
+        profile_meta: ProfileMetadata,
+        job_queue: JobQueue,
+    ) -> GenericHttpResponse:
+        return self.handle_request(
+            request=request, job_definition=job_definition, profile_meta=profile_meta
+        )
+
     def run_outer(self, args: JobArguments) -> JobResult:
         args.logger.log_warning(
             message=f"Unexpected call to run_outer for CustomHttpJob: {args.job_definition.id}"
@@ -189,7 +208,7 @@ def _resolve_webhook_signature_key(
     )
 
 
-class WebhookJob[WPT](Job[webhook_job_t.WebhookEventPayload]):
+class WebhookJob[WPT](Job[webhook_job_t.WebhookEventPayload], HttpJob):
     @property
     def payload_type(self) -> type[webhook_job_t.WebhookEventPayload]:
         return webhook_job_t.WebhookEventPayload
@@ -219,12 +238,14 @@ class WebhookJob[WPT](Job[webhook_job_t.WebhookEventPayload]):
         if not hmac.compare_digest(request_body_signature, passed_signature):
             raise HttpException.payload_failed_signature()
 
-    @staticmethod
-    def handle_request(
+    @typing.override
+    def respond_to_request(
+        self,
         *,
         request: GenericHttpRequest,
-        job_definition: job_definition_t.HttpJobDefinitionBase,
-        profile_meta: ProfileMetadata,  # noqa: ARG004
+        job_definition: HttpJobDefinitionBase,
+        profile_meta: ProfileMetadata,
+        job_queue: JobQueue,
     ) -> GenericHttpResponse:
         try:
             request_body = simplejson.loads(request.body_text)
@@ -234,19 +255,15 @@ class WebhookJob[WPT](Job[webhook_job_t.WebhookEventPayload]):
 
         request_context = extract_request_context(headers=request.headers)
 
-        try:
-            send_job_queue_message(
-                job_ref_name=job_definition.id,
-                payload=queued_job_t.QueuedJobPayload(
-                    invocation_context=queued_job_t.InvocationContextWebhook(
-                        webhook_payload=webhook_payload,
-                        request_context=request_context,
-                    )
-                ),
-                port=get_local_admin_server_port(),
-            )
-        except CommandServerException as e:
-            raise HttpException.unknown_error() from e
+        job_queue.enqueue(
+            job_ref_name=job_definition.id,
+            payload=queued_job_t.QueuedJobPayload(
+                invocation_context=queued_job_t.InvocationContextWebhook(
+                    webhook_payload=webhook_payload,
+                    request_context=request_context,
+                )
+            ),
+        )
 
         return GenericHttpResponse(
             response=simplejson.dumps(serialize_for_api(WebhookResponse())),

uncountable/integration/job_queue.py

@@ -0,0 +1,14 @@
+from uncountable.integration.queue_runner.datastore import Datastore
+from uncountable.types import queued_job_t
+
+
+class JobQueue:
+    def __init__(self, datastore: Datastore) -> None:
+        self._datastore = datastore
+
+    def enqueue(
+        self, *, job_ref_name: str, payload: queued_job_t.QueuedJobPayload
+    ) -> str:
+        return self._datastore.add_job_to_queue(
+            job_payload=payload, job_ref_name=job_ref_name
+        ).queued_job_uuid

uncountable/integration/ownership_lock.py

@@ -0,0 +1,77 @@
+import time
+from collections.abc import Callable, Generator
+from contextlib import contextmanager
+from enum import StrEnum
+
+from opentelemetry.trace import get_current_span
+from sqlalchemy import Connection, text
+from sqlalchemy.engine import Engine
+
+from uncountable.integration.telemetry import Logger
+
+OWNERSHIP_ADVISORY_LOCK_KEY = 0x756E635F696E7467
+OWNERSHIP_LOCK_POLL_INTERVAL_SECS = 1
+
+
+class OwnershipOutcome(StrEnum):
+    ACQUIRED = "acquired"
+    SHUTDOWN_REQUESTED = "shutdown_requested"
+
+
+def _try_acquire_advisory_lock(connection: Connection) -> bool:
+    return bool(
+        connection.execute(
+            text("select pg_try_advisory_lock(:key)"),
+            {"key": OWNERSHIP_ADVISORY_LOCK_KEY},
+        ).scalar_one()
+    )
+
+
+def _release_advisory_lock(connection: Connection) -> None:
+    connection.execute(
+        text("select pg_advisory_unlock(:key)"),
+        {"key": OWNERSHIP_ADVISORY_LOCK_KEY},
+    )
+
+
+def _wait_for_advisory_lock(
+    *, connection: Connection, is_shutdown_requested: Callable[[], bool]
+) -> OwnershipOutcome:
+    logger = Logger(get_current_span())
+    if is_shutdown_requested():
+        return OwnershipOutcome.SHUTDOWN_REQUESTED
+    if _try_acquire_advisory_lock(connection):
+        return OwnershipOutcome.ACQUIRED
+    logger.log_info("waiting for the integration ownership lock held by another pod")
+    while not is_shutdown_requested():
+        time.sleep(OWNERSHIP_LOCK_POLL_INTERVAL_SECS)
+        if is_shutdown_requested():
+            return OwnershipOutcome.SHUTDOWN_REQUESTED
+        if _try_acquire_advisory_lock(connection):
+            logger.log_info("acquired the integration ownership lock")
+            return OwnershipOutcome.ACQUIRED
+    return OwnershipOutcome.SHUTDOWN_REQUESTED
+
+
+@contextmanager
+def hold_integration_ownership(
+    *, engine: Engine, is_shutdown_requested: Callable[[], bool]
+) -> Generator[OwnershipOutcome]:
+    dialect_name = engine.dialect.name
+    match dialect_name:
+        case "sqlite":
+            yield OwnershipOutcome.ACQUIRED
+        case "postgresql":
+            with engine.connect().execution_options(
+                isolation_level="AUTOCOMMIT"
+            ) as connection:
+                outcome = _wait_for_advisory_lock(
+                    connection=connection, is_shutdown_requested=is_shutdown_requested
+                )
+                try:
+                    yield outcome
+                finally:
+                    if outcome == OwnershipOutcome.ACQUIRED:
+                        _release_advisory_lock(connection)
+        case _:
+            raise ValueError(f"unsupported database dialect: {dialect_name!r}")

uncountable/integration/queue_runner/command_server/command_client.py

@@ -85,6 +85,8 @@ def send_job_cancellation_message(
                 return CommandCancelJobStatus.NO_JOB_FOUND
             case CancelJobStatus.CANCELLED_WITH_RESTART:
                 return CommandCancelJobStatus.CANCELLED_WITH_RESTART
+            case CancelJobStatus.CANCELLED_QUEUED:
+                return CommandCancelJobStatus.CANCELLED_QUEUED
             case CancelJobStatus.JOB_ALREADY_COMPLETED:
                 return CommandCancelJobStatus.JOB_ALREADY_COMPLETED
             case _:

uncountable/integration/queue_runner/command_server/command_server.py

@@ -94,6 +94,8 @@ async def serve(command_queue: CommandQueue, datastore: Datastore) -> None:
                     proto_status = CancelJobStatus.NO_JOB_FOUND
                 case CommandCancelJobStatus.CANCELLED_WITH_RESTART:
                     proto_status = CancelJobStatus.CANCELLED_WITH_RESTART
+                case CommandCancelJobStatus.CANCELLED_QUEUED:
+                    proto_status = CancelJobStatus.CANCELLED_QUEUED
                 case CommandCancelJobStatus.JOB_ALREADY_COMPLETED:
                     proto_status = CancelJobStatus.JOB_ALREADY_COMPLETED
                 case _:

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

@@ -71,6 +71,7 @@ enum CancelJobStatus {
 	CANCELLED_WITH_RESTART = 0;
 	NO_JOB_FOUND = 1;
 	JOB_ALREADY_COMPLETED = 2;
+	CANCELLED_QUEUED = 3;
 }
 
 message CancelJobResult {

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

@@ -29,7 +29,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"\xc2\x01\n\x15ListQueuedJobsRequest\x12\x0e\n\x06offset\x18\x01 \x01(\r\x12\r\n\x05limit\x18\x02 \x01(\r\x12\x0e\n\x06status\x18\x03 \x01(\t\x12\x14\n\x0cjob_ref_name\x18\x04 \x01(\t\x12\x32\n\x0esubmitted_from\x18\x05 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x30\n\x0csubmitted_to\x18\x06 \x01(\x0b\x32\x1a.google.protobuf.Timestamp"\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")\n\x17ProvisionWebhookRequest\x12\x0e\n\x06job_id\x18\x01 \x01(\t".\n\x16ProvisionWebhookResult\x12\x14\n\x0cprofile_name\x18\x01 \x01(\t*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\xbd\x03\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\x12G\n\x10ProvisionWebhook\x12\x18.ProvisionWebhookRequest\x1a\x17.ProvisionWebhookResult"\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"\xc2\x01\n\x15ListQueuedJobsRequest\x12\x0e\n\x06offset\x18\x01 \x01(\r\x12\r\n\x05limit\x18\x02 \x01(\r\x12\x0e\n\x06status\x18\x03 \x01(\t\x12\x14\n\x0cjob_ref_name\x18\x04 \x01(\t\x12\x32\n\x0esubmitted_from\x18\x05 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x30\n\x0csubmitted_to\x18\x06 \x01(\x0b\x32\x1a.google.protobuf.Timestamp"\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")\n\x17ProvisionWebhookRequest\x12\x0e\n\x06job_id\x18\x01 \x01(\t".\n\x16ProvisionWebhookResult\x12\x14\n\x0cprofile_name\x18\x01 \x01(\t*p\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\x12\x14\n\x10\x43\x41NCELLED_QUEUED\x10\x03\x32\xbd\x03\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\x12G\n\x10ProvisionWebhook\x12\x18.ProvisionWebhookRequest\x1a\x17.ProvisionWebhookResult"\x00\x62\x06proto3'
 )
 
 _globals = globals()
@@ -42,7 +42,7 @@ _builder.BuildTopDescriptorsAndMessages(
 if not _descriptor._USE_C_DESCRIPTORS:
     DESCRIPTOR._loaded_options = None
     _globals["_CANCELJOBSTATUS"]._serialized_start = 1109
-    _globals["_CANCELJOBSTATUS"]._serialized_end = 1199
+    _globals["_CANCELJOBSTATUS"]._serialized_end = 1221
     _globals["_ENQUEUEJOBREQUEST"]._serialized_start = 118
     _globals["_ENQUEUEJOBREQUEST"]._serialized_end = 187
     _globals["_ENQUEUEJOBRESULT"]._serialized_start = 189
@@ -73,6 +73,6 @@ if not _descriptor._USE_C_DESCRIPTORS:
     _globals["_PROVISIONWEBHOOKREQUEST"]._serialized_end = 1059
     _globals["_PROVISIONWEBHOOKRESULT"]._serialized_start = 1061
     _globals["_PROVISIONWEBHOOKRESULT"]._serialized_end = 1107
-    _globals["_COMMANDSERVER"]._serialized_start = 1202
-    _globals["_COMMANDSERVER"]._serialized_end = 1647
+    _globals["_COMMANDSERVER"]._serialized_start = 1224
+    _globals["_COMMANDSERVER"]._serialized_end = 1669
 # @@protoc_insertion_point(module_scope)

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

@@ -16,10 +16,12 @@ class CancelJobStatus(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
     CANCELLED_WITH_RESTART: _ClassVar[CancelJobStatus]
     NO_JOB_FOUND: _ClassVar[CancelJobStatus]
     JOB_ALREADY_COMPLETED: _ClassVar[CancelJobStatus]
+    CANCELLED_QUEUED: _ClassVar[CancelJobStatus]
 
 CANCELLED_WITH_RESTART: CancelJobStatus
 NO_JOB_FOUND: CancelJobStatus
 JOB_ALREADY_COMPLETED: CancelJobStatus
+CANCELLED_QUEUED: CancelJobStatus
 
 class EnqueueJobRequest(_message.Message):
     __slots__ = ("job_ref_name", "serialized_payload")

uncountable/integration/queue_runner/command_server/types.py

@@ -8,6 +8,7 @@ from uncountable.types import queued_job_t
 
 class CommandCancelJobStatus(StrEnum):
     CANCELLED_WITH_RESTART = "cancelled_with_restart"
+    CANCELLED_QUEUED = "cancelled_queued"
     NO_JOB_FOUND = "no_job_found"
     JOB_ALREADY_COMPLETED = "job_already_completed"
 

uncountable/integration/queue_runner/datastore/datastore_sql_base.py

@@ -204,7 +204,19 @@ class DatastoreSqlBase(Datastore):
 
             return None
 
-    def load_job_queue(self) -> list[queued_job_t.QueuedJob]:
+    def list_queued_job_uuids(self) -> list[str]:
+        with self.session_maker() as session:
+            select_stmt = select(QueuedJob.id).filter(
+                or_(
+                    QueuedJob.status == queued_job_t.JobStatus.QUEUED,
+                    QueuedJob.status.is_(None),
+                )
+            )
+            return [row.id for row in session.execute(select_stmt)]
+
+    def load_queued_jobs(
+        self, *, queued_job_uuids: list[str]
+    ) -> list[queued_job_t.QueuedJob]:
         with self.session_maker() as session:
             select_stmt = (
                 select(
@@ -215,12 +227,7 @@ class DatastoreSqlBase(Datastore):
                     QueuedJob.status,
                     QueuedJob.submitted_at,
                 )
-                .filter(
-                    or_(
-                        QueuedJob.status == queued_job_t.JobStatus.QUEUED,
-                        QueuedJob.status.is_(None),
-                    )
-                )
+                .filter(QueuedJob.id.in_(queued_job_uuids))
                 .order_by(QueuedJob.submitted_at)
             )
 

uncountable/integration/queue_runner/datastore/interface.py

@@ -31,7 +31,12 @@ class Datastore(ABC):
     def increment_num_attempts(self, queued_job_uuid: str) -> int: ...
 
     @abstractmethod
-    def load_job_queue(self) -> list[queued_job_t.QueuedJob]: ...
+    def list_queued_job_uuids(self) -> list[str]: ...
+
+    @abstractmethod
+    def load_queued_jobs(
+        self, *, queued_job_uuids: list[str]
+    ) -> list[queued_job_t.QueuedJob]: ...
 
     @abstractmethod
     def get_next_queued_job_for_ref_name(