UncountablePythonSDK 0.0.69__py3-none-any.whl → 0.0.71__py3-none-any.whl

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

@@ -0,0 +1,38 @@
+# ruff: noqa
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import message as _message
+from typing import ClassVar as _ClassVar, Optional as _Optional
+
+DESCRIPTOR: _descriptor.FileDescriptor
+
+class EnqueueJobRequest(_message.Message):
+    __slots__ = ("job_ref_name", "serialized_payload")
+    JOB_REF_NAME_FIELD_NUMBER: _ClassVar[int]
+    SERIALIZED_PAYLOAD_FIELD_NUMBER: _ClassVar[int]
+    job_ref_name: str
+    serialized_payload: str
+    def __init__(
+        self,
+        job_ref_name: _Optional[str] = ...,
+        serialized_payload: _Optional[str] = ...,
+    ) -> None: ...
+
+class EnqueueJobResult(_message.Message):
+    __slots__ = ("successfully_queued", "queued_job_uuid")
+    SUCCESSFULLY_QUEUED_FIELD_NUMBER: _ClassVar[int]
+    QUEUED_JOB_UUID_FIELD_NUMBER: _ClassVar[int]
+    successfully_queued: bool
+    queued_job_uuid: str
+    def __init__(
+        self, successfully_queued: bool = ..., queued_job_uuid: _Optional[str] = ...
+    ) -> None: ...
+
+class CheckHealthRequest(_message.Message):
+    __slots__ = ()
+    def __init__(self) -> None: ...
+
+class CheckHealthResult(_message.Message):
+    __slots__ = ("success",)
+    SUCCESS_FIELD_NUMBER: _ClassVar[int]
+    success: bool
+    def __init__(self, success: bool = ...) -> None: ...

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

@@ -0,0 +1,129 @@
+# mypy: disable-error-code="no-untyped-def"
+# ruff: noqa
+# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
+"""Client and server classes corresponding to protobuf-defined services."""
+
+import grpc
+
+from uncountable.integration.queue_runner.command_server.protocol import (
+    command_server_pb2 as uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2,
+)
+
+
+class CommandServerStub(object):
+    """Missing associated documentation comment in .proto file."""
+
+    def __init__(self, channel):
+        """Constructor.
+
+        Args:
+            channel: A grpc.Channel.
+        """
+        self.EnqueueJob = channel.unary_unary(
+            "/CommandServer/EnqueueJob",
+            request_serializer=uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.EnqueueJobRequest.SerializeToString,
+            response_deserializer=uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.EnqueueJobResult.FromString,
+        )
+        self.CheckHealth = channel.unary_unary(
+            "/CommandServer/CheckHealth",
+            request_serializer=uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.CheckHealthRequest.SerializeToString,
+            response_deserializer=uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.CheckHealthResult.FromString,
+        )
+
+
+class CommandServerServicer(object):
+    """Missing associated documentation comment in .proto file."""
+
+    def EnqueueJob(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 CheckHealth(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 = {
+        "EnqueueJob": grpc.unary_unary_rpc_method_handler(
+            servicer.EnqueueJob,
+            request_deserializer=uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.EnqueueJobRequest.FromString,
+            response_serializer=uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.EnqueueJobResult.SerializeToString,
+        ),
+        "CheckHealth": grpc.unary_unary_rpc_method_handler(
+            servicer.CheckHealth,
+            request_deserializer=uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.CheckHealthRequest.FromString,
+            response_serializer=uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.CheckHealthResult.SerializeToString,
+        ),
+    }
+    generic_handler = grpc.method_handlers_generic_handler(
+        "CommandServer", rpc_method_handlers
+    )
+    server.add_generic_rpc_handlers((generic_handler,))
+
+
+# This class is part of an EXPERIMENTAL API.
+class CommandServer(object):
+    """Missing associated documentation comment in .proto file."""
+
+    @staticmethod
+    def EnqueueJob(
+        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/EnqueueJob",
+            uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.EnqueueJobRequest.SerializeToString,
+            uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.EnqueueJobResult.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+        )
+
+    @staticmethod
+    def CheckHealth(
+        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/CheckHealth",
+            uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.CheckHealthRequest.SerializeToString,
+            uncountable_dot_integration_dot_queue__runner_dot_command__server_dot_protocol_dot_command__server__pb2.CheckHealthResult.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+        )

uncountable/integration/queue_runner/command_server/types.py

@@ -0,0 +1,52 @@
+import asyncio
+import typing
+from dataclasses import dataclass
+from enum import StrEnum
+
+from uncountable.types import queued_job_t
+
+
+class CommandType(StrEnum):
+    ENQUEUE_JOB = "enqueue_job"
+
+
+RT = typing.TypeVar("RT")
+
+
+@dataclass(kw_only=True)
+class CommandBase(typing.Generic[RT]):
+    type: CommandType
+    response_queue: asyncio.Queue[RT]
+
+
+@dataclass(kw_only=True)
+class CommandEnqueueJobResponse:
+    queued_job_uuid: str
+
+
+@dataclass(kw_only=True)
+class CommandEnqueueJob(CommandBase[CommandEnqueueJobResponse]):
+    type: CommandType = CommandType.ENQUEUE_JOB
+    job_ref_name: str
+    payload: queued_job_t.QueuedJobPayload
+    response_queue: asyncio.Queue[CommandEnqueueJobResponse]
+
+
+_Command = CommandEnqueueJob
+
+
+CommandQueue = asyncio.Queue[_Command]
+
+CommandTask = asyncio.Task[_Command]
+
+
+class CommandServerException(Exception):
+    pass
+
+
+class CommandServerTimeout(CommandServerException):
+    pass
+
+
+class CommandServerBadResponse(CommandServerException):
+    pass

uncountable/integration/queue_runner/datastore/__init__.py

@@ -0,0 +1,3 @@
+from .datastore_sqlite import DatastoreSqlite
+
+__all__: list[str] = ["DatastoreSqlite"]

uncountable/integration/queue_runner/datastore/datastore_sqlite.py

@@ -0,0 +1,93 @@
+import uuid
+from datetime import datetime, timezone
+
+from sqlalchemy import delete, insert, select, update
+from sqlalchemy.engine import Engine
+
+from pkgs.argument_parser import CachedParser
+from pkgs.serialization_util import serialize_for_storage
+from uncountable.integration.db.session import DBSessionMaker
+from uncountable.integration.queue_runner.datastore.interface import Datastore
+from uncountable.integration.queue_runner.datastore.model import Base, QueuedJob
+from uncountable.types import queued_job_t
+
+queued_job_payload_parser = CachedParser(queued_job_t.QueuedJobPayload)
+
+
+class DatastoreSqlite(Datastore):
+    def __init__(self, session_maker: DBSessionMaker) -> None:
+        self.session_maker = session_maker
+        super().__init__()
+
+    @classmethod
+    def setup(cls, engine: Engine) -> None:
+        Base.metadata.create_all(engine)
+
+    def add_job_to_queue(
+        self, job_payload: queued_job_t.QueuedJobPayload, job_ref_name: str
+    ) -> queued_job_t.QueuedJob:
+        with self.session_maker() as session:
+            serialized_payload = serialize_for_storage(job_payload)
+            queued_job_uuid = str(uuid.uuid4())
+            num_attempts = 0
+            submitted_at = datetime.now(timezone.utc)
+            insert_stmt = insert(QueuedJob).values({
+                QueuedJob.id.key: queued_job_uuid,
+                QueuedJob.job_ref_name.key: job_ref_name,
+                QueuedJob.payload.key: serialized_payload,
+                QueuedJob.num_attempts: num_attempts,
+                QueuedJob.submitted_at: submitted_at,
+            })
+            session.execute(insert_stmt)
+            return queued_job_t.QueuedJob(
+                queued_job_uuid=queued_job_uuid,
+                job_ref_name=job_ref_name,
+                payload=job_payload,
+                submitted_at=submitted_at,
+                num_attempts=num_attempts,
+            )
+
+    def increment_num_attempts(self, queued_job_uuid: str) -> int:
+        with self.session_maker() as session:
+            update_stmt = (
+                update(QueuedJob)
+                .values({QueuedJob.num_attempts.key: QueuedJob.num_attempts + 1})
+                .filter(QueuedJob.id == queued_job_uuid)
+            )
+            session.execute(update_stmt)
+            session.flush()
+            # IMPROVE: python3.12's sqlite does not support the RETURNING clause
+            select_stmt = select(QueuedJob.num_attempts).filter(
+                QueuedJob.id == queued_job_uuid
+            )
+            return int(session.execute(select_stmt).one().num_attempts)
+
+    def remove_job_from_queue(self, queued_job_uuid: str) -> None:
+        with self.session_maker() as session:
+            delete_stmt = delete(QueuedJob).filter(QueuedJob.id == queued_job_uuid)
+            session.execute(delete_stmt)
+
+    def load_job_queue(self) -> list[queued_job_t.QueuedJob]:
+        with self.session_maker() as session:
+            select_stmt = select(
+                QueuedJob.id,
+                QueuedJob.payload,
+                QueuedJob.num_attempts,
+                QueuedJob.job_ref_name,
+                QueuedJob.submitted_at,
+            ).order_by(QueuedJob.submitted_at)
+
+            queued_jobs: list[queued_job_t.QueuedJob] = []
+            for row in session.execute(select_stmt):
+                parsed_payload = queued_job_payload_parser.parse_storage(row.payload)
+                queued_jobs.append(
+                    queued_job_t.QueuedJob(
+                        queued_job_uuid=row.id,
+                        job_ref_name=row.job_ref_name,
+                        num_attempts=row.num_attempts,
+                        submitted_at=row.submitted_at,
+                        payload=parsed_payload,
+                    )
+                )
+
+            return queued_jobs

uncountable/integration/queue_runner/datastore/interface.py

@@ -0,0 +1,19 @@
+from abc import ABC, abstractmethod
+
+from uncountable.types import queued_job_t
+
+
+class Datastore(ABC):
+    @abstractmethod
+    def add_job_to_queue(
+        self, job_payload: queued_job_t.QueuedJobPayload, job_ref_name: str
+    ) -> queued_job_t.QueuedJob: ...
+
+    @abstractmethod
+    def remove_job_from_queue(self, queued_job_uuid: str) -> None: ...
+
+    @abstractmethod
+    def increment_num_attempts(self, queued_job_uuid: str) -> int: ...
+
+    @abstractmethod
+    def load_job_queue(self) -> list[queued_job_t.QueuedJob]: ...

uncountable/integration/queue_runner/datastore/model.py

@@ -0,0 +1,17 @@
+from sqlalchemy import JSON, BigInteger, Column, DateTime, Text
+from sqlalchemy.orm import declarative_base
+from sqlalchemy.sql import func
+
+Base = declarative_base()
+
+
+class QueuedJob(Base):
+    __tablename__ = "queued_jobs"
+
+    id = Column(Text, primary_key=True)
+    job_ref_name = Column(Text, nullable=False, index=True)
+    submitted_at = Column(
+        DateTime(timezone=True), server_default=func.current_timestamp(), nullable=False
+    )
+    payload = Column(JSON, nullable=False)
+    num_attempts = Column(BigInteger, nullable=False, default=0, server_default="0")

uncountable/integration/queue_runner/job_scheduler.py

@@ -0,0 +1,128 @@
+import asyncio
+import typing
+from concurrent.futures import ProcessPoolExecutor
+from dataclasses import dataclass
+
+from opentelemetry.trace import get_current_span
+
+from uncountable.integration.db.connect import IntegrationDBService, create_db_engine
+from uncountable.integration.db.session import get_session_maker
+from uncountable.integration.queue_runner.command_server import (
+    CommandEnqueueJobResponse,
+    CommandQueue,
+    CommandTask,
+)
+from uncountable.integration.queue_runner.datastore import DatastoreSqlite
+from uncountable.integration.queue_runner.datastore.interface import Datastore
+from uncountable.integration.queue_runner.worker import Worker
+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
+
+_MAX_JOB_WORKERS = 5
+
+
+@dataclass(kw_only=True, frozen=True)
+class JobListenerKey:
+    profile_name: str
+    subqueue_name: str = "default"
+
+
+def _get_job_worker_key(
+    job_definition: job_definition_t.JobDefinition, profile_name: str
+) -> JobListenerKey:
+    return JobListenerKey(profile_name=profile_name)
+
+
+def on_worker_crash(
+    worker_key: JobListenerKey,
+) -> typing.Callable[[asyncio.Task], None]:
+    def hook(task: asyncio.Task) -> None:
+        raise Exception(
+            f"worker {worker_key.profile_name}_{worker_key.subqueue_name} crashed unexpectedly"
+        )
+
+    return hook
+
+
+def _start_workers(
+    process_pool: ProcessPoolExecutor, result_queue: ResultQueue, datastore: Datastore
+) -> dict[str, Worker]:
+    profiles = load_profiles()
+    job_queue_worker_lookup: dict[JobListenerKey, Worker] = {}
+    job_worker_lookup: dict[str, Worker] = {}
+    job_definition_lookup: dict[str, job_definition_t.JobDefinition] = {}
+    for profile in profiles:
+        for job_definition in profile.definition.jobs:
+            job_definition_lookup[job_definition.id] = job_definition
+            job_worker_key = _get_job_worker_key(job_definition, profile.name)
+            if job_worker_key not in job_queue_worker_lookup:
+                worker = Worker(
+                    process_pool=process_pool,
+                    listen_queue=asyncio.Queue(),
+                    result_queue=result_queue,
+                    datastore=datastore,
+                )
+                task = asyncio.create_task(worker.run_worker_loop())
+                task.add_done_callback(on_worker_crash(job_worker_key))
+                job_queue_worker_lookup[job_worker_key] = worker
+            job_worker_lookup[job_definition.id] = job_queue_worker_lookup[
+                job_worker_key
+            ]
+    return job_worker_lookup
+
+
+async def start_scheduler(command_queue: CommandQueue) -> None:
+    logger = Logger(get_current_span())
+    result_queue: ResultQueue = asyncio.Queue()
+    engine = create_db_engine(IntegrationDBService.RUNNER)
+    session_maker = get_session_maker(engine)
+
+    datastore = DatastoreSqlite(session_maker)
+    datastore.setup(engine)
+
+    with ProcessPoolExecutor(max_workers=_MAX_JOB_WORKERS) as process_pool:
+        job_worker_lookup = _start_workers(
+            process_pool, result_queue, datastore=datastore
+        )
+
+        queued_jobs = datastore.load_job_queue()
+
+        async def enqueue_queued_job(queued_job: queued_job_t.QueuedJob) -> None:
+            try:
+                worker = job_worker_lookup[queued_job.job_ref_name]
+            except KeyError as e:
+                logger.log_exception(e)
+                datastore.remove_job_from_queue(queued_job.queued_job_uuid)
+                return
+            await worker.listen_queue.put(queued_job)
+
+        for queued_job in queued_jobs:
+            await enqueue_queued_job(queued_job)
+
+        result_task: ResultTask = asyncio.create_task(result_queue.get())
+        command_task: CommandTask = asyncio.create_task(command_queue.get())
+        while True:
+            finished, _ = await asyncio.wait(
+                [result_task, command_task], return_when=asyncio.FIRST_COMPLETED
+            )
+
+            for task in finished:
+                if task == command_task:
+                    command = command_task.result()
+                    queued_job = datastore.add_job_to_queue(
+                        job_payload=command.payload, job_ref_name=command.job_ref_name
+                    )
+                    await command.response_queue.put(
+                        CommandEnqueueJobResponse(
+                            queued_job_uuid=queued_job.queued_job_uuid
+                        )
+                    )
+                    await enqueue_queued_job(queued_job)
+                    command_task = asyncio.create_task(command_queue.get())
+                elif task == result_task:
+                    queued_job_result = result_task.result()
+                    datastore.remove_job_from_queue(queued_job_result.queued_job_uuid)
+                    result_task = asyncio.create_task(result_queue.get())

uncountable/integration/queue_runner/queue_runner.py

@@ -0,0 +1,26 @@
+import asyncio
+
+from uncountable.integration.queue_runner.command_server import serve
+from uncountable.integration.queue_runner.command_server.types import CommandQueue
+from uncountable.integration.queue_runner.job_scheduler import start_scheduler
+
+
+async def queue_runner_loop() -> None:
+    command_queue: CommandQueue = asyncio.Queue()
+
+    command_server = asyncio.create_task(serve(command_queue))
+
+    scheduler = asyncio.create_task(start_scheduler(command_queue))
+
+    await scheduler
+    await command_server
+
+
+def start_queue_runner() -> None:
+    loop = asyncio.new_event_loop()
+    loop.run_until_complete(queue_runner_loop())
+    loop.close()
+
+
+if __name__ == "__main__":
+    start_queue_runner()

uncountable/integration/queue_runner/types.py

@@ -0,0 +1,7 @@
+from asyncio import Queue, Task
+
+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]

uncountable/integration/queue_runner/worker.py

@@ -0,0 +1,109 @@
+import asyncio
+from concurrent.futures import ProcessPoolExecutor
+from dataclasses import dataclass
+
+from uncountable.core.async_batch import AsyncBatchProcessor
+from uncountable.integration.construct_client import construct_uncountable_client
+from uncountable.integration.executors.executors import execute_job
+from uncountable.integration.job import JobArguments
+from uncountable.integration.queue_runner.datastore.interface import Datastore
+from uncountable.integration.queue_runner.types import ListenQueue, ResultQueue
+from uncountable.integration.scan_profiles import load_profiles
+from uncountable.integration.telemetry import JobLogger, get_otel_tracer
+from uncountable.types import base_t, job_definition_t, queued_job_t
+
+
+class Worker:
+    def __init__(
+        self,
+        *,
+        process_pool: ProcessPoolExecutor,
+        listen_queue: ListenQueue,
+        result_queue: ResultQueue,
+        datastore: Datastore,
+    ) -> None:
+        self.process_pool = process_pool
+        self.listen_queue = listen_queue
+        self.result_queue = result_queue
+        self.datastore = datastore
+
+    async def run_worker_loop(self) -> None:
+        while True:
+            queued_job = await self.listen_queue.get()
+            self.datastore.increment_num_attempts(queued_job.queued_job_uuid)
+            loop = asyncio.get_event_loop()
+            result = await loop.run_in_executor(
+                self.process_pool, run_queued_job, queued_job
+            )
+            assert isinstance(result, job_definition_t.JobResult)
+            await self.result_queue.put(
+                queued_job_t.QueuedJobResult(
+                    job_result=result, queued_job_uuid=queued_job.queued_job_uuid
+                )
+            )
+
+
+@dataclass(kw_only=True)
+class RegisteredJobDetails:
+    profile_metadata: job_definition_t.ProfileMetadata
+    job_definition: job_definition_t.JobDefinition
+
+
+def get_registered_job_details(job_ref_name: str) -> RegisteredJobDetails:
+    profiles = load_profiles()
+    for profile in profiles:
+        for job_definition in profile.definition.jobs:
+            if job_definition.id == job_ref_name:
+                return RegisteredJobDetails(
+                    profile_metadata=job_definition_t.ProfileMetadata(
+                        name=profile.name,
+                        base_url=profile.definition.base_url,
+                        auth_retrieval=profile.definition.auth_retrieval,
+                        client_options=profile.definition.client_options,
+                    ),
+                    job_definition=job_definition,
+                )
+    raise Exception(f"profile not found for job {job_ref_name}")
+
+
+def _resolve_queued_job_payload(queued_job: queued_job_t.QueuedJob) -> base_t.JsonValue:
+    match queued_job.payload.invocation_context:
+        case queued_job_t.InvocationContextCron():
+            return None
+        case queued_job_t.InvocationContextManual():
+            return None
+        case queued_job_t.InvocationContextWebhook():
+            return queued_job.payload.invocation_context.webhook_payload
+
+
+def run_queued_job(
+    queued_job: queued_job_t.QueuedJob,
+) -> job_definition_t.JobResult:
+    with get_otel_tracer().start_as_current_span(name="run_queued_job") as span:
+        job_details = get_registered_job_details(queued_job.job_ref_name)
+        job_logger = JobLogger(
+            base_span=span,
+            profile_metadata=job_details.profile_metadata,
+            job_definition=job_details.job_definition,
+        )
+        client = construct_uncountable_client(
+            profile_meta=job_details.profile_metadata, job_logger=job_logger
+        )
+        batch_processor = AsyncBatchProcessor(client=client)
+
+        payload = _resolve_queued_job_payload(queued_job)
+
+        args = JobArguments(
+            job_definition=job_details.job_definition,
+            client=client,
+            batch_processor=batch_processor,
+            profile_metadata=job_details.profile_metadata,
+            logger=job_logger,
+            payload=payload,
+        )
+
+        return execute_job(
+            args=args,
+            profile_metadata=job_details.profile_metadata,
+            job_definition=job_details.job_definition,
+        )

uncountable/integration/scan_profiles.py

@@ -1,3 +1,4 @@
+import functools
 import os
 from dataclasses import dataclass
 from importlib import resources
@@ -14,6 +15,7 @@ class ProfileDetails:
     definition: job_definition_t.ProfileDefinition
 
 
+@functools.cache
 def load_profiles() -> list[ProfileDetails]:
     profiles_module = os.environ["UNC_PROFILES_MODULE"]
     profiles = [