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

uncountable/integration/job.py

@@ -1,45 +1,272 @@
+import functools
+import hmac
+import typing
+from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from uncountable.core.client import Client
-from uncountable.integration.types import JobDefinition
 
-from abc import ABC, abstractmethod
+import simplejson
+
+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 (
+    GenericHttpRequest,
+    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.secret_retrieval.retrieve_secret import retrieve_secret
+from uncountable.integration.telemetry import JobLogger
+from uncountable.types import (
+    base_t,
+    job_definition_t,
+    queued_job_t,
+    webhook_job_t,
+)
+from uncountable.types.job_definition_t import (
+    HttpJobDefinitionBase,
+    JobDefinition,
+    JobResult,
+    ProfileMetadata,
+)
 
 
-@dataclass
-class JobArgumentsBase:
+@dataclass(kw_only=True)
+class JobArguments:
     job_definition: JobDefinition
+    profile_metadata: ProfileMetadata
     client: Client
+    batch_processor: AsyncBatchProcessor
+    logger: JobLogger
+    payload: base_t.JsonValue
+    job_uuid: str
+
+
+# only for compatibility:
+CronJobArguments = JobArguments
+
+
+class Job[PT](ABC):
+    _unc_job_registered: bool = False
+
+    @property
+    @abstractmethod
+    def payload_type(self) -> type[PT]: ...
+
+    @abstractmethod
+    def run_outer(self, args: JobArguments) -> JobResult: ...
+
+    @functools.cached_property
+    def _cached_payload_parser(self) -> CachedParser[PT]:
+        return CachedParser(self.payload_type)
+
+    def get_payload(self, payload: base_t.JsonValue) -> PT:
+        return self._cached_payload_parser.parse_storage(payload)
+
+
+class CronJob(Job):
+    @property
+    def payload_type(self) -> type[None]:
+        return type(None)
+
+    def run_outer(self, args: JobArguments) -> JobResult:
+        assert isinstance(args, CronJobArguments)
+        return self.run(args)
+
+    @abstractmethod
+    def run(self, args: JobArguments) -> JobResult: ...
+
+
+WPT = typing.TypeVar("WPT")
 
 
-@dataclass
-class CronJobArguments(JobArgumentsBase):
-    # can imagine passing additional data such as in the sftp or webhook cases
+@dataclass(kw_only=True)
+class WebhookResponse:
     pass
 
 
-JobArguments = CronJobArguments
+class _RequestValidatorClient(Client):
+    def __init__(self, *, base_url: str, oauth_bearer_token: str):
+        super().__init__(
+            base_url=base_url,
+            auth_details=AuthDetailsOAuth(refresh_token=""),
+            config=None,
+        )
+        self._oauth_bearer_token = oauth_bearer_token
 
+    def _get_oauth_bearer_token(self, *, oauth_details: AuthDetailsOAuth) -> str:
+        return self._oauth_bearer_token
 
-@dataclass
-class JobResult:
-    success: bool
 
+class CustomHttpJob(Job[GenericHttpRequest]):
+    @property
+    def payload_type(self) -> type[GenericHttpRequest]:
+        return GenericHttpRequest
 
-class Job(ABC):
-    _unc_job_registered: bool = False
+    @staticmethod
+    @abstractmethod
+    def validate_request(
+        *,
+        request: GenericHttpRequest,
+        job_definition: HttpJobDefinitionBase,
+        profile_meta: ProfileMetadata,
+    ) -> None:
+        """
+        Validate that the request is valid. If the request is invalid, raise an
+        exception.
+        """
+        ...
 
+    @staticmethod
+    def get_validated_oauth_request_user_id(
+        *, profile_metadata: ProfileMetadata, request: GenericHttpRequest
+    ) -> base_t.ObjectId:
+        token = request.headers.get("Authorization", "").replace("Bearer ", "")
+        if token == "":
+            raise HttpException(
+                message="unauthorized; no bearer token in request", error_code=401
+            )
+        return (
+            _RequestValidatorClient(
+                base_url=profile_metadata.base_url,
+                oauth_bearer_token=token,
+            )
+            .get_current_user_info()
+            .user_id
+        )
+
+    @staticmethod
     @abstractmethod
-    def run(self, args: JobArguments) -> JobResult:
+    def handle_request(
+        *,
+        request: GenericHttpRequest,
+        job_definition: HttpJobDefinitionBase,
+        profile_meta: ProfileMetadata,
+    ) -> GenericHttpResponse:
+        """
+        Handle the request synchronously. Normally this should just enqueue a job
+        and return immediately (see WebhookJob as an example).
+        """
         ...
 
+    def run_outer(self, args: JobArguments) -> JobResult:
+        args.logger.log_warning(
+            message=f"Unexpected call to run_outer for CustomHttpJob: {args.job_definition.id}"
+        )
+        return JobResult(success=False)
 
-class CronJob(Job):
 
+class WebhookJob[WPT](Job[webhook_job_t.WebhookEventPayload]):
+    @property
+    def payload_type(self) -> type[webhook_job_t.WebhookEventPayload]:
+        return webhook_job_t.WebhookEventPayload
+
+    @property
     @abstractmethod
-    def run(self, args: CronJobArguments) -> JobResult:
-        ...
+    def webhook_payload_type(self) -> type[WPT]: ...
+
+    @staticmethod
+    def validate_request(
+        *,
+        request: GenericHttpRequest,
+        job_definition: job_definition_t.HttpJobDefinitionBase,
+        profile_meta: ProfileMetadata,
+    ) -> None:
+        assert isinstance(job_definition, job_definition_t.WebhookJobDefinition)
+        signature_key = retrieve_secret(
+            profile_metadata=profile_meta,
+            secret_retrieval=job_definition.signature_key_secret,
+        )
+        passed_signature = request.headers.get("Uncountable-Webhook-Signature")
+        if passed_signature is None:
+            raise HttpException.no_signature_passed()
+
+        request_body_signature = hmac.new(
+            signature_key.encode("utf-8"), msg=request.body_bytes, digestmod="sha256"
+        ).hexdigest()
+
+        if request_body_signature != passed_signature:
+            raise HttpException.payload_failed_signature()
+
+    @staticmethod
+    def handle_request(
+        *,
+        request: GenericHttpRequest,
+        job_definition: job_definition_t.HttpJobDefinitionBase,
+        profile_meta: ProfileMetadata,  # noqa: ARG004
+    ) -> GenericHttpResponse:
+        try:
+            request_body = simplejson.loads(request.body_text)
+            webhook_payload = typing.cast(base_t.JsonValue, request_body)
+        except (simplejson.JSONDecodeError, ValueError) as e:
+            raise HttpException.body_parse_error() from e
+
+        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
+                    )
+                ),
+                port=get_local_admin_server_port(),
+            )
+        except CommandServerException as e:
+            raise HttpException.unknown_error() from e
 
+        return GenericHttpResponse(
+            response=simplejson.dumps(serialize_for_api(WebhookResponse())),
+            status_code=200,
+        )
 
-def register_job(cls: Job):
+    def run_outer(self, args: JobArguments) -> JobResult:
+        webhook_body = self.get_payload(args.payload)
+        inner_payload = CachedParser(self.webhook_payload_type).parse_api(
+            webhook_body.data
+        )
+        return self.run(args, inner_payload)
+
+    @abstractmethod
+    def run(self, args: JobArguments, payload: WPT) -> JobResult: ...
+
+
+def register_job(cls: type[Job]) -> type[Job]:
     cls._unc_job_registered = True
     return cls
+
+
+class RunsheetWebhookJob(WebhookJob[webhook_job_t.RunsheetWebhookPayload]):
+    @property
+    def webhook_payload_type(self) -> type:
+        return webhook_job_t.RunsheetWebhookPayload
+
+    @abstractmethod
+    def build_runsheet(
+        self,
+        *,
+        args: JobArguments,
+        payload: webhook_job_t.RunsheetWebhookPayload,
+    ) -> FileUpload: ...
+
+    def run(
+        self, args: JobArguments, payload: webhook_job_t.RunsheetWebhookPayload
+    ) -> JobResult:
+        runsheet = self.build_runsheet(args=args, payload=payload)
+
+        files = args.client.upload_files(file_uploads=[runsheet])
+        args.client.complete_async_upload(
+            async_job_id=payload.async_job_id, file_id=files[0].file_id
+        )
+
+        return JobResult(
+            success=True,
+        )

uncountable/integration/queue_runner/command_server/__init__.py

@@ -0,0 +1,28 @@
+from .command_client import check_health, send_job_queue_message
+from .command_server import serve
+from .types import (
+    CommandEnqueueJob,
+    CommandEnqueueJobResponse,
+    CommandQueue,
+    CommandRetryJob,
+    CommandRetryJobResponse,
+    CommandServerBadResponse,
+    CommandServerException,
+    CommandServerTimeout,
+    CommandTask,
+)
+
+__all__: list[str] = [
+    "serve",
+    "check_health",
+    "send_job_queue_message",
+    "CommandEnqueueJob",
+    "CommandEnqueueJobResponse",
+    "CommandRetryJob",
+    "CommandRetryJobResponse",
+    "CommandTask",
+    "CommandQueue",
+    "CommandServerTimeout",
+    "CommandServerException",
+    "CommandServerBadResponse",
+]

uncountable/integration/queue_runner/command_server/command_client.py

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

uncountable/integration/queue_runner/command_server/command_server.py

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

uncountable/integration/queue_runner/command_server/constants.py

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

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

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