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

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=f"{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,199 @@
+import datetime
+import multiprocessing
+import subprocess
+import sys
+import time
+from dataclasses import dataclass
+from datetime import UTC
+from enum import StrEnum
+
+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
+
+AnyProcess = multiprocessing.Process | subprocess.Popen[bytes]
+
+
+class ProcessName(StrEnum):
+    QUEUE_RUNNER = "queue_runner"
+    CRON_SERVER = "cron_server"
+    UWSGI = "uwsgi"
+
+
+@dataclass(kw_only=True)
+class ProcessInfo:
+    name: ProcessName
+    process: AnyProcess
+
+    @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: dict[ProcessName, ProcessInfo]) -> None:
+    logger.log_info("received shutdown command, shutting down sub-processes")
+    for proc_info in processes.values():
+        if proc_info.is_alive:
+            proc_info.process.terminate()
+
+    shutdown_start = time.time()
+    still_living_processes = list(processes.values())
+    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 restart_process(
+    logger: Logger, proc_info: ProcessInfo, processes: dict[ProcessName, ProcessInfo]
+) -> None:
+    logger.log_error(
+        f"process {proc_info.name} shut down unexpectedly - exit code {proc_info.exitcode}. Restarting..."
+    )
+
+    match proc_info.name:
+        case ProcessName.QUEUE_RUNNER:
+            queue_proc = multiprocessing.Process(target=start_queue_runner)
+            queue_proc.start()
+            new_info = ProcessInfo(name=ProcessName.QUEUE_RUNNER, process=queue_proc)
+            processes[ProcessName.QUEUE_RUNNER] = new_info
+            try:
+                _wait_queue_runner_online()
+                logger.log_info("queue runner restarted successfully")
+            except Exception as e:
+                logger.log_exception(e)
+                logger.log_error(
+                    "queue runner failed to restart, shutting down scheduler"
+                )
+                handle_shutdown(logger, processes)
+                sys.exit(1)
+
+        case ProcessName.CRON_SERVER:
+            cron_proc = multiprocessing.Process(target=cron_target)
+            cron_proc.start()
+            new_info = ProcessInfo(name=ProcessName.CRON_SERVER, process=cron_proc)
+            processes[ProcessName.CRON_SERVER] = new_info
+            logger.log_info("cron server restarted successfully")
+
+        case ProcessName.UWSGI:
+            uwsgi_proc: AnyProcess = subprocess.Popen(["uwsgi", "--die-on-term"])
+            new_info = ProcessInfo(name=ProcessName.UWSGI, process=uwsgi_proc)
+            processes[ProcessName.UWSGI] = new_info
+            logger.log_info("uwsgi restarted successfully")
+
+
+def check_process_alive(
+    logger: Logger, processes: dict[ProcessName, ProcessInfo]
+) -> None:
+    for proc_info in processes.values():
+        if not proc_info.is_alive:
+            restart_process(logger, proc_info, processes)
+
+
+def _wait_queue_runner_online() -> None:
+    MAX_QUEUE_RUNNER_HEALTH_CHECKS = 10
+    QUEUE_RUNNER_HEALTH_CHECK_DELAY_SECS = 1
+
+    num_attempts = 0
+    before = datetime.datetime.now(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.datetime.now(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: dict[ProcessName, ProcessInfo] = {}
+
+    multiprocessing.set_start_method("forkserver")
+
+    def add_process(process: ProcessInfo) -> None:
+        processes[process.name] = process
+        logger.log_info(f"started process {process.name}")
+
+    runner_process = multiprocessing.Process(target=start_queue_runner)
+    runner_process.start()
+    add_process(ProcessInfo(name=ProcessName.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=ProcessName.CRON_SERVER, process=cron_process))
+
+    uwsgi_process = subprocess.Popen([
+        "uwsgi",
+        "--die-on-term",
+    ])
+    add_process(ProcessInfo(name=ProcessName.UWSGI, process=uwsgi_process))
+
+    try:
+        while True:
+            check_process_alive(logger, processes=processes)
+            time.sleep(1)
+    except KeyboardInterrupt:
+        handle_shutdown(logger, processes=processes)
+
+
+if __name__ == "__main__":
+    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=f"{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

uncountable/integration/server.py

@@ -1,29 +1,43 @@
 import signal
-from types import TracebackType
 from dataclasses import asdict
-from typing import Optional, assert_never
+from types import TracebackType
+from typing import assert_never
+
+from apscheduler.executors.pool import ThreadPoolExecutor
+from apscheduler.jobstores.sqlalchemy import SQLAlchemyJobStore
 from apscheduler.schedulers.background import BackgroundScheduler
 from apscheduler.schedulers.base import BaseScheduler
-from apscheduler.jobstores.sqlalchemy import SQLAlchemyJobStore
-from apscheduler.executors.pool import ThreadPoolExecutor
-from uncountable.integration.cron import CronJobArgs, cron_job_executor
 from apscheduler.triggers.cron import CronTrigger
+from opentelemetry.trace import get_current_span
 from sqlalchemy.engine.base import Engine
 
-from uncountable.integration.types import (
-    AuthRetrieval,
+from uncountable.core.environment import get_local_admin_server_port
+from uncountable.integration.cron import CronJobArgs, cron_job_executor
+from uncountable.integration.queue_runner.command_server.command_client import (
+    send_vaccuum_queued_jobs_message,
+)
+from uncountable.integration.telemetry import Logger
+from uncountable.types import base_t, job_definition_t
+from uncountable.types.job_definition_t import (
     CronJobDefinition,
-    JobDefinition,
-    ProfileMetadata,
+    HttpJobDefinitionBase,
 )
 
-
 _MAX_APSCHEDULER_CONCURRENT_JOBS = 1
 
+VACCUUM_QUEUED_JOBS_JOB_ID = "vacuum_queued_jobs"
+
+STATIC_JOB_IDS = {VACCUUM_QUEUED_JOBS_JOB_ID}
+
+
+def vaccuum_queued_jobs() -> None:
+    send_vaccuum_queued_jobs_message(port=get_local_admin_server_port())
+
 
 class IntegrationServer:
     _scheduler: BaseScheduler
     _engine: Engine
+    _server_logger: Logger
 
     def __init__(self, engine: Engine) -> None:
         self._engine = engine
@@ -32,49 +46,83 @@ class IntegrationServer:
             jobstores={"default": SQLAlchemyJobStore(engine=engine)},
             executors={"default": ThreadPoolExecutor(_MAX_APSCHEDULER_CONCURRENT_JOBS)},
         )
+        self._server_logger = Logger(get_current_span())
 
-    def register_profile(
-        self,
-        *,
-        profile_name: str,
-        base_url: str,
-        auth_retrieval: AuthRetrieval,
-        jobs: list[JobDefinition],
-    ) -> None:
-        for job_defn in jobs:
-            profile_metadata = ProfileMetadata(
-                name=profile_name, auth_retrieval=auth_retrieval, base_url=base_url
-            )
-            match job_defn:
-                case CronJobDefinition():
-                    # Add to ap scheduler
-                    job_kwargs = asdict(
-                        CronJobArgs(
-                            definition=job_defn, profile_metadata=profile_metadata
-                        )
-                    )
-                    existing_job = self._scheduler.get_job(job_defn.id)
-                    if existing_job is not None:
-                        existing_job.modify(
-                            name=job_defn.name,
-                            kwargs=job_kwargs,
-                        )
-                        existing_job.reschedule(
-                            CronTrigger.from_crontab(job_defn.cron_spec)
-                        )
-                    else:
-                        self._scheduler.add_job(
-                            cron_job_executor,
-                            # IMPROVE: reconsider these defaults
-                            max_instances=1,
-                            coalesce=True,
-                            trigger=CronTrigger.from_crontab(job_defn.cron_spec),
-                            name=job_defn.name,
-                            id=job_defn.id,
-                            kwargs=job_kwargs,
+    def _register_static_jobs(self) -> None:
+        all_job_ids = {job.id for job in self._scheduler.get_jobs()}
+        if VACCUUM_QUEUED_JOBS_JOB_ID in all_job_ids:
+            self._scheduler.remove_job(VACCUUM_QUEUED_JOBS_JOB_ID)
+
+        self._scheduler.add_job(
+            vaccuum_queued_jobs,
+            max_instances=1,
+            coalesce=True,
+            trigger=CronTrigger.from_crontab("5 4 * * 4"),
+            name="Vaccuum queued jobs",
+            id=VACCUUM_QUEUED_JOBS_JOB_ID,
+            kwargs={},
+            misfire_grace_time=None,
+        )
+
+    def register_jobs(self, profiles: list[job_definition_t.ProfileMetadata]) -> None:
+        valid_job_ids: set[str] = set()
+        for profile_metadata in profiles:
+            for job_defn in profile_metadata.jobs:
+                valid_job_ids.add(job_defn.id)
+                match job_defn:
+                    case CronJobDefinition():
+                        # Add to ap scheduler
+                        job_kwargs = asdict(
+                            CronJobArgs(
+                                definition=job_defn, profile_metadata=profile_metadata
+                            )
                         )
-                case _:
-                    assert_never(job_defn.trigger)
+                        try:
+                            existing_job = self._scheduler.get_job(job_defn.id)
+                        except ValueError as e:
+                            self._server_logger.log_warning(
+                                f"could not reconstitute job {job_defn.id}: {e}"
+                            )
+                            self._scheduler.remove_job(job_defn.id)
+                            existing_job = None
+                        if existing_job is not None:
+                            existing_job.modify(
+                                name=job_defn.name,
+                                kwargs=job_kwargs,
+                                misfire_grace_time=None,
+                            )
+                            existing_job.reschedule(
+                                CronTrigger.from_crontab(job_defn.cron_spec)
+                            )
+                            if not job_defn.enabled:
+                                existing_job.pause()
+                            else:
+                                existing_job.resume()
+                        else:
+                            job_opts: dict[str, base_t.JsonValue] = {}
+                            if not job_defn.enabled:
+                                job_opts["next_run_time"] = None
+                            self._scheduler.add_job(
+                                cron_job_executor,
+                                # IMPROVE: reconsider these defaults
+                                max_instances=1,
+                                coalesce=True,
+                                trigger=CronTrigger.from_crontab(job_defn.cron_spec),
+                                name=job_defn.name,
+                                id=job_defn.id,
+                                kwargs=job_kwargs,
+                                misfire_grace_time=None,
+                                **job_opts,
+                            )
+                    case HttpJobDefinitionBase():
+                        pass
+                    case _:
+                        assert_never(job_defn)
+        all_job_ids = {job.id for job in self._scheduler.get_jobs()}
+        invalid_job_ids = all_job_ids.difference(valid_job_ids.union(STATIC_JOB_IDS))
+
+        for job_id in invalid_job_ids:
+            self._scheduler.remove_job(job_id)
 
     def serve_forever(self) -> None:
         signal.pause()
@@ -87,12 +135,13 @@ class IntegrationServer:
 
     def __enter__(self) -> "IntegrationServer":
         self._start_apscheduler()
+        self._register_static_jobs()
         return self
 
     def __exit__(
         self,
-        exc_type: Optional[type[BaseException]],
-        exc_val: Optional[BaseException],
-        exc_tb: Optional[TracebackType],
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
     ) -> None:
         self._stop_apscheduler()