UncountablePythonSDK 0.0.190__tar.gz → 0.0.191__tar.gz

{uncountablepythonsdk-0.0.190/UncountablePythonSDK.egg-info → uncountablepythonsdk-0.0.191}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: UncountablePythonSDK
-Version: 0.0.190
+Version: 0.0.191
 Summary: Uncountable SDK
 Project-URL: Homepage, https://github.com/uncountableinc/uncountable-python-sdk
 Project-URL: Repository, https://github.com/uncountableinc/uncountable-python-sdk.git

{uncountablepythonsdk-0.0.190 → uncountablepythonsdk-0.0.191/UncountablePythonSDK.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: UncountablePythonSDK
-Version: 0.0.190
+Version: 0.0.191
 Summary: Uncountable SDK
 Project-URL: Homepage, https://github.com/uncountableinc/uncountable-python-sdk
 Project-URL: Repository, https://github.com/uncountableinc/uncountable-python-sdk.git

{uncountablepythonsdk-0.0.190 → uncountablepythonsdk-0.0.191}/UncountablePythonSDK.egg-info/SOURCES.txt

@@ -155,6 +155,8 @@ uncountable/integration/construct_client.py
 uncountable/integration/cron.py
 uncountable/integration/entrypoint.py
 uncountable/integration/job.py
+uncountable/integration/job_queue.py
+uncountable/integration/ownership_lock.py
 uncountable/integration/request_context.py
 uncountable/integration/scan_profiles.py
 uncountable/integration/scheduler.py

{uncountablepythonsdk-0.0.190 → uncountablepythonsdk-0.0.191}/UncountablePythonSDK.egg-info/scm_file_list.json

@@ -159,6 +159,8 @@
     "uncountable/integration/http_server/__init__.py",
     "uncountable/integration/http_server/types.py",
     "uncountable/integration/job.py",
+    "uncountable/integration/job_queue.py",
+    "uncountable/integration/ownership_lock.py",
     "uncountable/integration/queue_runner/__init__.py",
     "uncountable/integration/queue_runner/command_server/__init__.py",
     "uncountable/integration/queue_runner/command_server/command_client.py",

{uncountablepythonsdk-0.0.190 → uncountablepythonsdk-0.0.191}/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")

{uncountablepythonsdk-0.0.190 → uncountablepythonsdk-0.0.191}/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`.

{uncountablepythonsdk-0.0.190 → uncountablepythonsdk-0.0.191}/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:

{uncountablepythonsdk-0.0.190 → uncountablepythonsdk-0.0.191}/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())),

uncountablepythonsdk-0.0.191/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

uncountablepythonsdk-0.0.191/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}")

{uncountablepythonsdk-0.0.190 → uncountablepythonsdk-0.0.191}/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 _:

{uncountablepythonsdk-0.0.190 → uncountablepythonsdk-0.0.191}/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 _:

{uncountablepythonsdk-0.0.190 → uncountablepythonsdk-0.0.191}/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 {

{uncountablepythonsdk-0.0.190 → uncountablepythonsdk-0.0.191}/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)

{uncountablepythonsdk-0.0.190 → uncountablepythonsdk-0.0.191}/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")

{uncountablepythonsdk-0.0.190 → uncountablepythonsdk-0.0.191}/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"
 

{uncountablepythonsdk-0.0.190 → uncountablepythonsdk-0.0.191}/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)
             )
 

{uncountablepythonsdk-0.0.190 → uncountablepythonsdk-0.0.191}/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(

{uncountablepythonsdk-0.0.190 → uncountablepythonsdk-0.0.191}/uncountable/integration/queue_runner/job_scheduler.py

@@ -31,6 +31,9 @@ from uncountable.types import job_definition_t, queued_job_t
 
 from .types import RESTART_EXIT_CODE, ResultQueue, ResultTask
 
+QUEUED_JOB_POLL_INTERVAL_SECS = 1
+RESTART_DELAY_SECS = 5
+
 
 @dataclass(kw_only=True, frozen=True)
 class JobListenerKey:
@@ -65,10 +68,12 @@ def on_worker_crash(
 
 
 def _start_workers(
+    *,
     process_pool: ProcessPoolExecutor,
     result_queue: ResultQueue,
     datastore: Datastore,
     worker_tasks: list[asyncio.Task],
+    running_job_uuids: set[str],
 ) -> dict[str, Worker]:
     profiles = load_profiles()
     job_queue_worker_lookup: dict[JobListenerKey, Worker] = {}
@@ -84,6 +89,7 @@ def _start_workers(
                     listen_queue=asyncio.Queue(),
                     result_queue=result_queue,
                     datastore=datastore,
+                    running_job_uuids=running_job_uuids,
                 )
                 task = asyncio.create_task(worker.run_worker_loop())
                 task.add_done_callback(on_worker_crash(job_worker_key))
@@ -108,11 +114,16 @@ async def start_scheduler(
     with ProcessPoolExecutor(
         max_workers=get_max_job_workers(), max_tasks_per_child=1
     ) as process_pool:
+        running_job_uuids: set[str] = set()
         job_worker_lookup = _start_workers(
-            process_pool, result_queue, datastore=datastore, worker_tasks=worker_tasks
+            process_pool=process_pool,
+            result_queue=result_queue,
+            datastore=datastore,
+            worker_tasks=worker_tasks,
+            running_job_uuids=running_job_uuids,
         )
 
-        queued_jobs = datastore.load_job_queue()
+        dispatched_job_uuids: set[str] = set()
 
         async def enqueue_queued_job(queued_job: queued_job_t.QueuedJob) -> None:
             try:
@@ -123,8 +134,22 @@ async def start_scheduler(
                     queued_job.queued_job_uuid, queued_job_t.JobStatus.FAILED
                 )
                 return
+            dispatched_job_uuids.add(queued_job.queued_job_uuid)
             await worker.listen_queue.put(queued_job)
 
+        async def _dispatch_undispatched_queued_jobs() -> None:
+            undispatched_job_uuids = [
+                queued_job_uuid
+                for queued_job_uuid in datastore.list_queued_job_uuids()
+                if queued_job_uuid not in dispatched_job_uuids
+            ]
+            if len(undispatched_job_uuids) == 0:
+                return
+            for queued_job in datastore.load_queued_jobs(
+                queued_job_uuids=undispatched_job_uuids
+            ):
+                await enqueue_queued_job(queued_job)
+
         async def _enqueue_or_deduplicate_job(
             job_ref_name: str,
             payload: queued_job_t.QueuedJobPayload,
@@ -172,30 +197,36 @@ async def start_scheduler(
                 )
                 return
 
-            if queued_job.status == queued_job_t.JobStatus.QUEUED:
-                datastore.remove_job_from_queue(command.queued_job_uuid)
+            if queued_job.status != queued_job_t.JobStatus.QUEUED:
                 await command.response_queue.put(
                     CommandCancelJobResponse(
-                        status=CommandCancelJobStatus.CANCELLED_WITH_RESTART
+                        status=CommandCancelJobStatus.JOB_ALREADY_COMPLETED
                     )
                 )
+                return
 
-                logger.log_debug(
-                    "Job cancelled successfully.",
-                    attributes={"job.queued_job_uuid": command.queued_job_uuid},
-                )
-
-                def delayed_exit() -> None:
-                    os._exit(RESTART_EXIT_CODE)
-
-                threading.Timer(interval=5, function=delayed_exit).start()
-
+            datastore.remove_job_from_queue(command.queued_job_uuid)
+            if command.queued_job_uuid in running_job_uuids:
+                cancel_status = CommandCancelJobStatus.CANCELLED_WITH_RESTART
+                threading.Timer(
+                    interval=RESTART_DELAY_SECS,
+                    function=os._exit,
+                    args=[RESTART_EXIT_CODE],
+                ).start()
             else:
-                await command.response_queue.put(
-                    CommandCancelJobResponse(
-                        status=CommandCancelJobStatus.JOB_ALREADY_COMPLETED
-                    )
-                )
+                dispatched_job_uuids.discard(command.queued_job_uuid)
+                cancel_status = CommandCancelJobStatus.CANCELLED_QUEUED
+
+            await command.response_queue.put(
+                CommandCancelJobResponse(status=cancel_status)
+            )
+            logger.log_debug(
+                "Job cancelled successfully.",
+                attributes={
+                    "job.queued_job_uuid": command.queued_job_uuid,
+                    "job.cancel_status": cancel_status,
+                },
+            )
 
         async def _handle_retry_job_command(command: CommandRetryJob) -> None:
             queued_job = datastore.retry_job(command.queued_job_uuid)
@@ -221,17 +252,21 @@ async def start_scheduler(
                 attributes={"job.queued_job_uuids": deleted_uuids},
             )
 
-        for queued_job in queued_jobs:
-            await enqueue_queued_job(queued_job)
+        await _dispatch_undispatched_queued_jobs()
 
         result_task: ResultTask = asyncio.create_task(result_queue.get())
         command_task: CommandTask = asyncio.create_task(command_queue.get())
+        poll_task = asyncio.create_task(asyncio.sleep(QUEUED_JOB_POLL_INTERVAL_SECS))
         shutdown_task: asyncio.Task[object] | None = None
         if shutdown_event is not None:
             shutdown_task = asyncio.create_task(shutdown_event.wait())
 
         while True:
-            waitables: list[asyncio.Task[object]] = [result_task, command_task]
+            waitables: list[asyncio.Task[object]] = [
+                result_task,
+                command_task,
+                poll_task,
+            ]
             if shutdown_task is not None:
                 waitables.append(shutdown_task)
 
@@ -246,6 +281,7 @@ async def start_scheduler(
                         worker_task.cancel()
                     command_task.cancel()
                     result_task.cancel()
+                    poll_task.cancel()
                     process_pool.shutdown(wait=False, cancel_futures=True)
                     return
                 if task == command_task:
@@ -275,4 +311,11 @@ async def start_scheduler(
                                 queued_job_result.queued_job_uuid,
                                 queued_job_t.JobStatus.FAILED,
                             )
+                    dispatched_job_uuids.discard(queued_job_result.queued_job_uuid)
+                    running_job_uuids.discard(queued_job_result.queued_job_uuid)
                     result_task = asyncio.create_task(result_queue.get())
+                elif task == poll_task:
+                    await _dispatch_undispatched_queued_jobs()
+                    poll_task = asyncio.create_task(
+                        asyncio.sleep(QUEUED_JOB_POLL_INTERVAL_SECS)
+                    )