UncountablePythonSDK 0.0.82__py3-none-any.whl → 0.0.132__py3-none-any.whl

uncountable/integration/scheduler.py

@@ -1,9 +1,11 @@
+import datetime
 import multiprocessing
 import subprocess
 import sys
 import time
 from dataclasses import dataclass
-from datetime import datetime, timezone
+from datetime import UTC
+from enum import StrEnum
 
 from opentelemetry.trace import get_current_span
 
@@ -18,11 +20,19 @@ 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: str
-    process: multiprocessing.Process | subprocess.Popen[bytes]
+    name: ProcessName
+    process: AnyProcess
 
     @property
     def is_alive(self) -> bool:
@@ -45,14 +55,14 @@ class ProcessInfo:
                 return self.process.poll()
 
 
-def handle_shutdown(logger: Logger, processes: list[ProcessInfo]) -> None:
+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:
+    for proc_info in processes.values():
         if proc_info.is_alive:
             proc_info.process.terminate()
 
     shutdown_start = time.time()
-    still_living_processes = processes
+    still_living_processes = list(processes.values())
     while (
         time.time() - shutdown_start < SHUTDOWN_TIMEOUT_SECS
         and len(still_living_processes) > 0
@@ -81,46 +91,84 @@ def handle_shutdown(logger: Logger, processes: list[ProcessInfo]) -> None:
         proc_info.process.kill()
 
 
-def check_process_alive(logger: Logger, processes: list[ProcessInfo]) -> None:
-    for proc_info in processes:
+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:
-            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)
+            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
+    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:
+    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.now(timezone.utc)
+        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: list[ProcessInfo] = []
+    processes: dict[ProcessName, ProcessInfo] = {}
+
+    multiprocessing.set_start_method("forkserver")
 
     def add_process(process: ProcessInfo) -> None:
-        processes.append(process)
+        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="queue runner", process=runner_process))
+    add_process(ProcessInfo(name=ProcessName.QUEUE_RUNNER, process=runner_process))
 
     try:
         _wait_queue_runner_online()
@@ -131,13 +179,13 @@ def main() -> None:
 
     cron_process = multiprocessing.Process(target=cron_target)
     cron_process.start()
-    add_process(ProcessInfo(name="cron server", process=cron_process))
+    add_process(ProcessInfo(name=ProcessName.CRON_SERVER, process=cron_process))
 
     uwsgi_process = subprocess.Popen([
         "uwsgi",
         "--die-on-term",
     ])
-    add_process(ProcessInfo(name="uwsgi", process=uwsgi_process))
+    add_process(ProcessInfo(name=ProcessName.UWSGI, process=uwsgi_process))
 
     try:
         while True:
@@ -147,4 +195,5 @@ def main() -> None:
         handle_shutdown(logger, processes=processes)
 
 
-main()
+if __name__ == "__main__":
+    main()

uncountable/integration/secret_retrieval/retrieve_secret.py

@@ -54,7 +54,7 @@ def _load_secret_overrides(profile_name: str) -> dict[SecretRetrieval, str]:
     profiles_module = os.environ["UNC_PROFILES_MODULE"]
     try:
         overrides = overrides_parser.parse_yaml_resource(
-            package=".".join([profiles_module, profile_name]),
+            package=f"{profiles_module}.{profile_name}",
             resource="local_overrides.yaml",
         )
         return {

uncountable/integration/server.py

@@ -1,7 +1,7 @@
 import signal
 from dataclasses import asdict
 from types import TracebackType
-from typing import Optional, assert_never
+from typing import assert_never
 
 from apscheduler.executors.pool import ThreadPoolExecutor
 from apscheduler.jobstores.sqlalchemy import SQLAlchemyJobStore
@@ -11,16 +11,28 @@ from apscheduler.triggers.cron import CronTrigger
 from opentelemetry.trace import get_current_span
 from sqlalchemy.engine.base import Engine
 
+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,
-    WebhookJobDefinition,
+    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
@@ -36,11 +48,27 @@ class IntegrationServer:
         )
         self._server_logger = Logger(get_current_span())
 
+    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 = []
+        valid_job_ids: set[str] = set()
         for profile_metadata in profiles:
             for job_defn in profile_metadata.jobs:
-                valid_job_ids.append(job_defn.id)
+                valid_job_ids.add(job_defn.id)
                 match job_defn:
                     case CronJobDefinition():
                         # Add to ap scheduler
@@ -86,14 +114,15 @@ class IntegrationServer:
                                 misfire_grace_time=None,
                                 **job_opts,
                             )
-                    case WebhookJobDefinition():
+                    case HttpJobDefinitionBase():
                         pass
                     case _:
                         assert_never(job_defn)
-        all_jobs = self._scheduler.get_jobs()
-        for job in all_jobs:
-            if job.id not in valid_job_ids:
-                self._scheduler.remove_job(job.id)
+        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()
@@ -106,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()

uncountable/integration/telemetry.py

@@ -1,4 +1,5 @@
 import functools
+import json
 import os
 import time
 import traceback
@@ -11,7 +12,10 @@ from opentelemetry import _logs, trace
 from opentelemetry.exporter.otlp.proto.http._log_exporter import OTLPLogExporter
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
 from opentelemetry.sdk._logs import Logger as OTELLogger
-from opentelemetry.sdk._logs import LoggerProvider, LogRecord
+from opentelemetry.sdk._logs import (
+    LoggerProvider,
+    LogRecord,
+)
 from opentelemetry.sdk._logs.export import BatchLogRecordProcessor, ConsoleLogExporter
 from opentelemetry.sdk.resources import Attributes, Resource
 from opentelemetry.sdk.trace import TracerProvider
@@ -32,6 +36,11 @@ def _cast_attributes(attributes: dict[str, base_t.JsonValue]) -> Attributes:
     return cast(Attributes, attributes)
 
 
+def one_line_formatter(record: LogRecord) -> str:
+    json_data = record.to_json()
+    return json.dumps(json.loads(json_data), separators=(",", ":")) + "\n"
+
+
 @functools.cache
 def get_otel_resource() -> Resource:
     attributes: dict[str, base_t.JsonValue] = {
@@ -60,7 +69,9 @@ def get_otel_tracer() -> Tracer:
 @functools.cache
 def get_otel_logger() -> OTELLogger:
     provider = LoggerProvider(resource=get_otel_resource())
-    provider.add_log_record_processor(BatchLogRecordProcessor(ConsoleLogExporter()))
+    provider.add_log_record_processor(
+        BatchLogRecordProcessor(ConsoleLogExporter(formatter=one_line_formatter))
+    )
     if get_otel_enabled():
         provider.add_log_record_processor(BatchLogRecordProcessor(OTLPLogExporter()))
     _logs.set_logger_provider(provider)
@@ -87,8 +98,27 @@ class Logger:
     def current_trace_id(self) -> int | None:
         return self.current_span.get_span_context().trace_id
 
-    def _patch_attributes(self, attributes: Attributes | None) -> Attributes:
-        return attributes or {}
+    def _patch_attributes(
+        self,
+        attributes: Attributes | None,
+        *,
+        message: str | None = None,
+        severity: LogSeverity | None = None,
+    ) -> Attributes:
+        patched_attributes = {**(attributes if attributes is not None else {})}
+        if message is not None:
+            patched_attributes["message"] = message
+        elif "body" in patched_attributes:
+            patched_attributes["message"] = patched_attributes["body"]
+
+        if severity is not None:
+            patched_attributes["status"] = severity.lower()
+        elif "severity_text" in patched_attributes and isinstance(
+            patched_attributes["severity_text"], str
+        ):
+            patched_attributes["status"] = patched_attributes["severity_text"].lower()
+
+        return patched_attributes
 
     def _emit_log(
         self, message: str, *, severity: LogSeverity, attributes: Attributes | None
@@ -98,7 +128,9 @@ class Logger:
             body=message,
             severity_text=severity,
             timestamp=time.time_ns(),
-            attributes=self._patch_attributes(attributes),
+            attributes=self._patch_attributes(
+                message=message, severity=severity, attributes=attributes
+            ),
             span_id=self.current_span_id,
             trace_id=self.current_trace_id,
             trace_flags=DEFAULT_TRACE_OPTIONS,
@@ -131,8 +163,10 @@ class Logger:
         message: str | None = None,
         attributes: Attributes | None = None,
     ) -> None:
-        traceback_str = "".join(traceback.format_tb(exception.__traceback__))
-        patched_attributes = self._patch_attributes(attributes)
+        traceback_str = "".join(traceback.format_exception(exception))
+        patched_attributes = self._patch_attributes(
+            message=message, severity=LogSeverity.ERROR, attributes=attributes
+        )
         self.current_span.record_exception(
             exception=exception, attributes=patched_attributes
         )
@@ -163,9 +197,17 @@ class JobLogger(Logger):
         self.job_definition = job_definition
         super().__init__(base_span)
 
-    def _patch_attributes(self, attributes: Attributes | None) -> Attributes:
+    def _patch_attributes(
+        self,
+        attributes: Attributes | None,
+        *,
+        message: str | None = None,
+        severity: LogSeverity | None = None,
+    ) -> Attributes:
         patched_attributes: dict[str, base_t.JsonValue] = {
-            **(attributes if attributes is not None else {})
+            **super()._patch_attributes(
+                attributes=attributes, message=message, severity=severity
+            )
         }
         patched_attributes["profile.name"] = self.profile_metadata.name
         patched_attributes["profile.base_url"] = self.profile_metadata.base_url
@@ -177,7 +219,7 @@ class JobLogger(Logger):
                 patched_attributes["job.definition.cron_spec"] = (
                     self.job_definition.cron_spec
                 )
-            case job_definition_t.WebhookJobDefinition():
+            case job_definition_t.HttpJobDefinitionBase():
                 pass
             case _:
                 assert_never(self.job_definition)
@@ -196,3 +238,14 @@ class JobLogger(Logger):
             case _:
                 assert_never(self.job_definition.executor)
         return _cast_attributes(patched_attributes)
+
+
+@contextmanager
+def push_scope_optional(
+    logger: Logger | None, scope_name: str, *, attributes: Attributes | None = None
+) -> Generator[None, None, None]:
+    if logger is None:
+        yield
+    else:
+        with logger.push_scope(scope_name, attributes=attributes):
+            yield

uncountable/integration/webhook_server/entrypoint.py

@@ -1,146 +1,71 @@
-import hmac
-import typing
-from dataclasses import dataclass
+import base64
 
 import flask
-import simplejson
 from flask.typing import ResponseReturnValue
-from flask.wrappers import Response
 from opentelemetry.trace import get_current_span
 from uncountable.core.environment import (
-    get_local_admin_server_port,
+    get_http_server_port,
     get_server_env,
-    get_webhook_server_port,
-)
-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.executors.script_executor import resolve_script_executor
+from uncountable.integration.http_server import GenericHttpRequest, HttpException
+from uncountable.integration.job import CustomHttpJob, WebhookJob
 from uncountable.integration.scan_profiles import load_profiles
-from uncountable.integration.secret_retrieval.retrieve_secret import retrieve_secret
 from uncountable.integration.telemetry import Logger
-from uncountable.types import base_t, job_definition_t, queued_job_t, webhook_job_t
-
-from pkgs.argument_parser import CachedParser
+from uncountable.types import job_definition_t
 
 app = flask.Flask(__name__)
 
 
-@dataclass(kw_only=True)
-class WebhookResponse:
-    pass
-
-
-webhook_payload_parser = CachedParser(webhook_job_t.WebhookEventBody)
-
-
-class WebhookException(BaseException):
-    error_code: int
-    message: str
-
-    def __init__(self, *, error_code: int, message: str) -> None:
-        self.error_code = error_code
-        self.message = message
-
-    @staticmethod
-    def payload_failed_signature() -> "WebhookException":
-        return WebhookException(
-            error_code=401, message="webhook payload did not match signature"
-        )
-
-    @staticmethod
-    def no_signature_passed() -> "WebhookException":
-        return WebhookException(error_code=400, message="missing signature")
-
-    @staticmethod
-    def body_parse_error() -> "WebhookException":
-        return WebhookException(error_code=400, message="body parse error")
-
-    @staticmethod
-    def unknown_error() -> "WebhookException":
-        return WebhookException(error_code=500, message="internal server error")
-
-    def __str__(self) -> str:
-        return f"[{self.error_code}]: {self.message}"
-
-    def make_error_response(self) -> Response:
-        return Response(
-            status=self.error_code, response={"error": {"message": str(self)}}
-        )
-
-
-def _parse_webhook_payload(
-    *, raw_request_body: bytes, signature_key: str, passed_signature: str
-) -> base_t.JsonValue:
-    request_body_signature = hmac.new(
-        signature_key.encode("utf-8"), msg=raw_request_body, digestmod="sha256"
-    ).hexdigest()
-
-    if request_body_signature != passed_signature:
-        raise WebhookException.payload_failed_signature()
-
-    try:
-        request_body = simplejson.loads(raw_request_body.decode())
-        return typing.cast(base_t.JsonValue, request_body)
-    except (simplejson.JSONDecodeError, ValueError) as e:
-        raise WebhookException.body_parse_error() from e
-
-
 def register_route(
     *,
     server_logger: Logger,
     profile_meta: job_definition_t.ProfileMetadata,
-    job: job_definition_t.WebhookJobDefinition,
+    job: job_definition_t.HttpJobDefinitionBase,
 ) -> None:
     route = f"/{profile_meta.name}/{job.id}"
 
-    def handle_webhook() -> ResponseReturnValue:
+    def handle_request() -> ResponseReturnValue:
         with server_logger.push_scope(route):
             try:
-                signature_key = retrieve_secret(
-                    profile_metadata=profile_meta,
-                    secret_retrieval=job.signature_key_secret,
+                if not isinstance(job.executor, job_definition_t.JobExecutorScript):
+                    raise HttpException.configuration_error(
+                        message="[internal] http job must use a script executor"
+                    )
+                job_instance = resolve_script_executor(
+                    executor=job.executor, profile_metadata=profile_meta
                 )
-
-                passed_signature = flask.request.headers.get(
-                    "Uncountable-Webhook-Signature"
+                if not isinstance(job_instance, (CustomHttpJob, WebhookJob)):
+                    raise HttpException.configuration_error(
+                        message="[internal] http job must descend from CustomHttpJob"
+                    )
+                http_request = GenericHttpRequest(
+                    body_base64=base64.b64encode(flask.request.get_data()).decode(),
+                    headers=dict(flask.request.headers),
                 )
-                if passed_signature is None:
-                    raise WebhookException.no_signature_passed()
-
-                webhook_payload = _parse_webhook_payload(
-                    raw_request_body=flask.request.data,
-                    signature_key=signature_key,
-                    passed_signature=passed_signature,
+                job_instance.validate_request(
+                    request=http_request, job_definition=job, profile_meta=profile_meta
+                )
+                http_response = job_instance.handle_request(
+                    request=http_request, job_definition=job, profile_meta=profile_meta
                 )
 
-                try:
-                    send_job_queue_message(
-                        job_ref_name=job.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 WebhookException.unknown_error() from e
-
-                return flask.jsonify(WebhookResponse())
-            except WebhookException as e:
+                return flask.make_response(
+                    http_response.response,
+                    http_response.status_code,
+                    http_response.headers,
+                )
+            except HttpException as e:
                 server_logger.log_exception(e)
                 return e.make_error_response()
             except Exception as e:
                 server_logger.log_exception(e)
-                return WebhookException.unknown_error().make_error_response()
+                return HttpException.unknown_error().make_error_response()
 
     app.add_url_rule(
         route,
-        endpoint=f"handle_webhook_{job.id}",
-        view_func=handle_webhook,
+        endpoint=f"handle_request_{job.id}",
+        view_func=handle_request,
         methods=["POST"],
     )
 
@@ -148,11 +73,13 @@ def register_route(
 
 
 def main() -> None:
+    app.add_url_rule("/health", "health", lambda: ("OK", 200))
+
     profiles = load_profiles()
     for profile_metadata in profiles:
         server_logger = Logger(get_current_span())
         for job in profile_metadata.jobs:
-            if isinstance(job, job_definition_t.WebhookJobDefinition):
+            if isinstance(job, job_definition_t.HttpJobDefinitionBase):
                 register_route(
                     server_logger=server_logger, profile_meta=profile_metadata, job=job
                 )
@@ -164,7 +91,7 @@ main()
 if __name__ == "__main__":
     app.run(
         host="0.0.0.0",
-        port=get_webhook_server_port(),
+        port=get_http_server_port(),
         debug=get_server_env() == "playground",
         exclude_patterns=[],
     )