UncountablePythonSDK 0.0.8__py3-none-any.whl → 0.0.92__py3-none-any.whl

uncountable/integration/queue_runner/job_scheduler.py

@@ -0,0 +1,163 @@
+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 (
+    CommandEnqueueJob,
+    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.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)
+
+        async def _enqueue_or_deduplicate_job(
+            job_ref_name: str,
+            payload: queued_job_t.QueuedJobPayload,
+        ) -> str:
+            if isinstance(
+                payload.invocation_context,
+                (
+                    queued_job_t.InvocationContextCron,
+                    queued_job_t.InvocationContextManual,
+                ),
+            ):
+                existing_queued_jobs = datastore.load_job_queue()
+                duplicate_job = next(
+                    (
+                        job
+                        for job in existing_queued_jobs
+                        if job.job_ref_name == job_ref_name
+                    ),
+                    None,
+                )
+                if duplicate_job is not None:
+                    return duplicate_job.queued_job_uuid
+            queued_job = datastore.add_job_to_queue(
+                job_payload=payload,
+                job_ref_name=job_ref_name,
+            )
+            await enqueue_queued_job(queued_job)
+            return queued_job.queued_job_uuid
+
+        async def _handle_enqueue_job_command(command: CommandEnqueueJob) -> None:
+            queued_job_uuid = await _enqueue_or_deduplicate_job(
+                job_ref_name=command.job_ref_name,
+                payload=command.payload,
+            )
+            await command.response_queue.put(
+                CommandEnqueueJobResponse(queued_job_uuid=queued_job_uuid)
+            )
+
+        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()
+                    match command:
+                        case CommandEnqueueJob():
+                            await _handle_enqueue_job_command(command=command)
+                        case _:
+                            typing.assert_never(command)
+                    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,119 @@
+import asyncio
+from concurrent.futures import ProcessPoolExecutor
+from dataclasses import dataclass
+
+from opentelemetry.trace import get_current_span
+
+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, Logger, 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:
+        logger = Logger(get_current_span())
+        while True:
+            try:
+                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
+                    )
+                )
+            except BaseException as e:
+                logger.log_exception(e)
+                raise e
+
+
+@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_metadata in profiles:
+        for job_definition in profile_metadata.jobs:
+            if job_definition.id == job_ref_name:
+                return RegisteredJobDetails(
+                    profile_metadata=profile_metadata,
+                    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,
+        )
+        try:
+            client = construct_uncountable_client(
+                profile_meta=job_details.profile_metadata, 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,
+                job_uuid=queued_job.queued_job_uuid,
+            )
+        except Exception as e:
+            job_logger.log_exception(e)
+            return job_definition_t.JobResult(success=False)
+        except BaseException as e:
+            job_logger.log_exception(e)
+            raise e

uncountable/integration/scan_profiles.py

@@ -0,0 +1,67 @@
+import functools
+from importlib import resources
+
+from pkgs.argument_parser import CachedParser
+from uncountable.core import environment
+from uncountable.types import integration_server_t, job_definition_t
+
+profile_parser = CachedParser(job_definition_t.ProfileDefinition)
+
+_DEFAULT_PROFILE_ENV = integration_server_t.IntegrationEnvironment.PROD
+_IGNORED_PROFILE_FOLDERS = ["__pycache__"]
+
+
+@functools.cache
+def load_profiles() -> list[job_definition_t.ProfileMetadata]:
+    profiles_module = environment.get_profiles_module()
+    integration_envs = environment.get_integration_envs()
+    profiles = [
+        entry
+        for entry in resources.files(profiles_module).iterdir()
+        if entry.is_dir() and entry.name not in _IGNORED_PROFILE_FOLDERS
+    ]
+    profile_details: list[job_definition_t.ProfileMetadata] = []
+    seen_job_ids: set[str] = set()
+    for profile_file in profiles:
+        profile_name = profile_file.name
+        try:
+            definition = profile_parser.parse_yaml_resource(
+                package=".".join([profiles_module, profile_name]),
+                resource="profile.yaml",
+            )
+            for job in definition.jobs:
+                if job.id in seen_job_ids:
+                    raise Exception(f"multiple jobs with id {job.id}")
+                seen_job_ids.add(job.id)
+
+            if definition.environments is not None:
+                for integration_env in integration_envs:
+                    environment_config = definition.environments.get(integration_env)
+                    if environment_config is not None:
+                        profile_details.append(
+                            job_definition_t.ProfileMetadata(
+                                name=profile_name,
+                                jobs=definition.jobs,
+                                base_url=environment_config.base_url,
+                                auth_retrieval=environment_config.auth_retrieval,
+                                client_options=environment_config.client_options,
+                            )
+                        )
+            elif _DEFAULT_PROFILE_ENV in integration_envs:
+                assert (
+                    definition.base_url is not None
+                    and definition.auth_retrieval is not None
+                ), f"define environments in profile.yaml for {profile_name}"
+                profile_details.append(
+                    job_definition_t.ProfileMetadata(
+                        name=profile_name,
+                        jobs=definition.jobs,
+                        base_url=definition.base_url,
+                        auth_retrieval=definition.auth_retrieval,
+                        client_options=definition.client_options,
+                    )
+                )
+        except FileNotFoundError as e:
+            print(f"WARN: profile.yaml not found for {profile_name}", e)
+            continue
+    return profile_details

uncountable/integration/scheduler.py

@@ -0,0 +1,150 @@
+import multiprocessing
+import subprocess
+import sys
+import time
+from dataclasses import dataclass
+from datetime import datetime, timezone
+
+from opentelemetry.trace import get_current_span
+
+from uncountable.core.environment import get_local_admin_server_port
+from uncountable.integration.entrypoint import main as cron_target
+from uncountable.integration.queue_runner.command_server import (
+    CommandServerTimeout,
+    check_health,
+)
+from uncountable.integration.queue_runner.queue_runner import start_queue_runner
+from uncountable.integration.telemetry import Logger
+
+SHUTDOWN_TIMEOUT_SECS = 30
+
+
+@dataclass(kw_only=True)
+class ProcessInfo:
+    name: str
+    process: multiprocessing.Process | subprocess.Popen[bytes]
+
+    @property
+    def is_alive(self) -> bool:
+        match self.process:
+            case multiprocessing.Process():
+                return self.process.is_alive()
+            case subprocess.Popen():
+                return self.process.poll() is None
+
+    @property
+    def pid(self) -> int | None:
+        return self.process.pid
+
+    @property
+    def exitcode(self) -> int | None:
+        match self.process:
+            case multiprocessing.Process():
+                return self.process.exitcode
+            case subprocess.Popen():
+                return self.process.poll()
+
+
+def handle_shutdown(logger: Logger, processes: list[ProcessInfo]) -> None:
+    logger.log_info("received shutdown command, shutting down sub-processes")
+    for proc_info in processes:
+        if proc_info.is_alive:
+            proc_info.process.terminate()
+
+    shutdown_start = time.time()
+    still_living_processes = processes
+    while (
+        time.time() - shutdown_start < SHUTDOWN_TIMEOUT_SECS
+        and len(still_living_processes) > 0
+    ):
+        current_loop_processes = [*still_living_processes]
+        logger.log_info(
+            "waiting for sub-processes to shut down",
+            attributes={
+                "still_living_processes": [
+                    proc_info.name for proc_info in still_living_processes
+                ]
+            },
+        )
+        still_living_processes = []
+        for proc_info in current_loop_processes:
+            if not proc_info.is_alive:
+                logger.log_info(f"{proc_info.name} shut down successfully")
+            else:
+                still_living_processes.append(proc_info)
+        time.sleep(1)
+
+    for proc_info in still_living_processes:
+        logger.log_warning(
+            f"{proc_info.name} failed to shut down after {SHUTDOWN_TIMEOUT_SECS} seconds, forcefully terminating"
+        )
+        proc_info.process.kill()
+
+
+def check_process_alive(logger: Logger, processes: list[ProcessInfo]) -> None:
+    for proc_info in processes:
+        if not proc_info.is_alive:
+            logger.log_error(
+                f"process {proc_info.name} shut down unexpectedly! shutting down scheduler; exit code is {proc_info.exitcode}"
+            )
+            handle_shutdown(logger, processes)
+            sys.exit(1)
+
+
+def _wait_queue_runner_online() -> None:
+    MAX_QUEUE_RUNNER_HEALTH_CHECKS = 10
+    QUEUE_RUNNER_HEALTH_CHECK_DELAY_SECS = 1
+
+    num_attempts = 0
+    before = datetime.now(timezone.utc)
+    while num_attempts < MAX_QUEUE_RUNNER_HEALTH_CHECKS:
+        try:
+            if check_health(port=get_local_admin_server_port()):
+                return
+        except CommandServerTimeout:
+            pass
+        num_attempts += 1
+        time.sleep(QUEUE_RUNNER_HEALTH_CHECK_DELAY_SECS)
+    after = datetime.now(timezone.utc)
+    duration_secs = (after - before).seconds
+    raise Exception(f"queue runner failed to come online after {duration_secs} seconds")
+
+
+def main() -> None:
+    logger = Logger(get_current_span())
+    processes: list[ProcessInfo] = []
+
+    def add_process(process: ProcessInfo) -> None:
+        processes.append(process)
+        logger.log_info(f"started process {process.name}")
+
+    runner_process = multiprocessing.Process(target=start_queue_runner)
+    runner_process.start()
+    add_process(ProcessInfo(name="queue runner", process=runner_process))
+
+    try:
+        _wait_queue_runner_online()
+    except Exception as e:
+        logger.log_exception(e)
+        handle_shutdown(logger, processes=processes)
+        return
+
+    cron_process = multiprocessing.Process(target=cron_target)
+    cron_process.start()
+    add_process(ProcessInfo(name="cron server", process=cron_process))
+
+    uwsgi_process = subprocess.Popen([
+        "uwsgi",
+        "--die-on-term",
+    ])
+    add_process(ProcessInfo(name="uwsgi", process=uwsgi_process))
+
+    try:
+        while True:
+            check_process_alive(logger, processes=processes)
+            time.sleep(1)
+    except KeyboardInterrupt:
+        handle_shutdown(logger, processes=processes)
+
+
+main()

uncountable/integration/secret_retrieval/__init__.py

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

uncountable/integration/secret_retrieval/retrieve_secret.py

@@ -0,0 +1,93 @@
+import base64
+import functools
+import json
+import os
+
+import boto3
+
+from pkgs.argument_parser import CachedParser
+from uncountable.types import overrides_t
+from uncountable.types.job_definition_t import ProfileMetadata
+from uncountable.types.secret_retrieval_t import (
+    SecretRetrieval,
+    SecretRetrievalAWS,
+    SecretRetrievalEnv,
+)
+
+
+class SecretRetrievalError(Exception):
+    def __init__(
+        self, secret_retrieval: SecretRetrieval, message: str | None = None
+    ) -> None:
+        self.secret_retrieval = secret_retrieval
+        self.message = message
+
+    def __str__(self) -> str:
+        append_message = ""
+        if self.message is not None:
+            append_message = f": {self.message}"
+        return f"{self.secret_retrieval.type} secret retrieval failed{append_message}"
+
+
+@functools.cache
+def _get_aws_secret(*, secret_name: str, region_name: str, sub_key: str | None) -> str:
+    client = boto3.client("secretsmanager", region_name=region_name)
+    response = client.get_secret_value(SecretId=secret_name)
+
+    if "SecretString" in response:
+        secret = response["SecretString"]
+    else:
+        secret = base64.b64decode(response["SecretBinary"])
+
+    value = json.loads(secret)
+
+    if sub_key is not None:
+        assert isinstance(value, dict)
+        return str(value[sub_key])
+    else:
+        return str(value)
+
+
+@functools.cache
+def _load_secret_overrides(profile_name: str) -> dict[SecretRetrieval, str]:
+    overrides_parser = CachedParser(overrides_t.Overrides)
+    profiles_module = os.environ["UNC_PROFILES_MODULE"]
+    try:
+        overrides = overrides_parser.parse_yaml_resource(
+            package=".".join([profiles_module, profile_name]),
+            resource="local_overrides.yaml",
+        )
+        return {
+            override.secret_retrieval: override.value for override in overrides.secrets
+        }
+    except FileNotFoundError:
+        return {}
+
+
+def retrieve_secret(
+    secret_retrieval: SecretRetrieval, profile_metadata: ProfileMetadata
+) -> str:
+    value_from_override = _load_secret_overrides(profile_metadata.name).get(
+        secret_retrieval
+    )
+    if value_from_override is not None:
+        return value_from_override
+
+    match secret_retrieval:
+        case SecretRetrievalEnv():
+            env_name = f"UNC_{profile_metadata.name.upper()}_{secret_retrieval.env_key.upper()}"
+            secret = os.environ.get(env_name)
+            if secret is None:
+                raise SecretRetrievalError(
+                    secret_retrieval, f"environment variable {env_name} missing"
+                )
+            return secret
+        case SecretRetrievalAWS():
+            try:
+                return _get_aws_secret(
+                    secret_name=secret_retrieval.secret_name,
+                    region_name=secret_retrieval.region,
+                    sub_key=secret_retrieval.sub_key,
+                )
+            except Exception as e:
+                raise SecretRetrievalError(secret_retrieval) from e