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

uncountable/integration/server.py

@@ -1,28 +1,38 @@
 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
 from apscheduler.schedulers.background import BackgroundScheduler
 from apscheduler.schedulers.base import BaseScheduler
 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
-from uncountable.types.client_config_t import ClientConfigOptions
+from uncountable.types import base_t, job_definition_t
 from uncountable.types.job_definition_t import (
-    AuthRetrieval,
     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
@@ -36,69 +46,83 @@ class IntegrationServer:
             jobstores={"default": SQLAlchemyJobStore(engine=engine)},
             executors={"default": ThreadPoolExecutor(_MAX_APSCHEDULER_CONCURRENT_JOBS)},
         )
-        self._server_logger = Logger()
+        self._server_logger = Logger(get_current_span())
 
-    def register_profile(
-        self,
-        *,
-        profile_name: str,
-        base_url: str,
-        auth_retrieval: AuthRetrieval,
-        jobs: list[JobDefinition],
-        client_options: Optional[ClientConfigOptions],
-    ) -> None:
-        for job_defn in jobs:
-            profile_metadata = ProfileMetadata(
-                name=profile_name,
-                auth_retrieval=auth_retrieval,
-                base_url=base_url,
-                client_options=client_options,
-            )
-            match job_defn:
-                case CronJobDefinition():
-                    # Add to ap scheduler
-                    job_kwargs = asdict(
-                        CronJobArgs(
-                            definition=job_defn, profile_metadata=profile_metadata
-                        )
-                    )
-                    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,
-                        )
-                        existing_job.reschedule(
-                            CronTrigger.from_crontab(job_defn.cron_spec)
+    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
+                            )
                         )
-                        if not job_defn.enabled:
-                            existing_job.pause()
+                        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:
-                            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,
-                            **job_opts,
-                        )
-                case _:
-                    assert_never(job_defn.trigger)
+                            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()
@@ -111,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,22 +1,34 @@
 import functools
+import json
 import os
-import sys
 import time
+import traceback
+import typing
 from contextlib import contextmanager
 from enum import StrEnum
-from typing import Generator, TextIO, assert_never, cast
+from typing import Generator, assert_never, cast
 
-from opentelemetry import trace
+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 LogRecord
+from opentelemetry.sdk._logs import Logger as OTELLogger
+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
 from opentelemetry.sdk.trace.export import (
     SimpleSpanProcessor,
 )
-from opentelemetry.trace import Tracer
+from opentelemetry.trace import DEFAULT_TRACE_OPTIONS, Span, Tracer
 
-from uncountable.core.version import get_version
+from uncountable.core.environment import (
+    get_otel_enabled,
+    get_server_env,
+    get_version,
+)
 from uncountable.types import base_t, job_definition_t
 
 
@@ -24,8 +36,13 @@ 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_tracer() -> Tracer:
+def get_otel_resource() -> Resource:
     attributes: dict[str, base_t.JsonValue] = {
         "service.name": "integration-server",
         "sdk.version": get_version(),
@@ -33,16 +50,34 @@ def get_tracer() -> Tracer:
     unc_version = os.environ.get("UNC_VERSION")
     if unc_version is not None:
         attributes["service.version"] = unc_version
-    unc_env = os.environ.get("UNC_INTEGRATION_ENV")
+    unc_env = get_server_env()
     if unc_env is not None:
         attributes["deployment.environment"] = unc_env
     resource = Resource.create(attributes=_cast_attributes(attributes))
-    provider = TracerProvider(resource=resource)
-    provider.add_span_processor(SimpleSpanProcessor(OTLPSpanExporter()))
+    return resource
+
+
+@functools.cache
+def get_otel_tracer() -> Tracer:
+    provider = TracerProvider(resource=get_otel_resource())
+    if get_otel_enabled():
+        provider.add_span_processor(SimpleSpanProcessor(OTLPSpanExporter()))
     trace.set_tracer_provider(provider)
     return provider.get_tracer("integration.telemetry")
 
 
+@functools.cache
+def get_otel_logger() -> OTELLogger:
+    provider = LoggerProvider(resource=get_otel_resource())
+    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)
+    return provider.get_logger("integration.telemetry")
+
+
 class LogSeverity(StrEnum):
     INFO = "Info"
     WARN = "Warn"
@@ -50,50 +85,129 @@ class LogSeverity(StrEnum):
 
 
 class Logger:
-    current_span_id: int | None = None
-    current_trace_id: int | None = None
+    current_span: Span
 
-    def _patch_attributes(self, attributes: Attributes | None) -> Attributes:
-        return attributes or {}
+    def __init__(self, base_span: Span) -> None:
+        self.current_span = base_span
+
+    @property
+    def current_span_id(self) -> int:
+        return self.current_span.get_span_context().span_id
+
+    @property
+    def current_trace_id(self) -> int | None:
+        return self.current_span.get_span_context().trace_id
+
+    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
     ) -> None:
+        otel_logger = get_otel_logger()
         log_record = LogRecord(
             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,
+            severity_number=_logs.SeverityNumber.UNSPECIFIED,
+            resource=get_otel_resource(),
         )
-        log_file: TextIO = sys.stderr if severity == LogSeverity.ERROR else sys.stdout
-        log_file.write(log_record.to_json())
-        log_file.flush()
+        otel_logger.emit(log_record)
 
     def log_info(self, message: str, *, attributes: Attributes | None = None) -> None:
-        self._emit_log(message=message, severity=LogSeverity.INFO, attributes=attributes)
+        self._emit_log(
+            message=message, severity=LogSeverity.INFO, attributes=attributes
+        )
 
-    def log_warning(self, message: str, *, attributes: Attributes | None = None) -> None:
-        self._emit_log(message=message, severity=LogSeverity.WARN, attributes=attributes)
+    def log_warning(
+        self, message: str, *, attributes: Attributes | None = None
+    ) -> None:
+        self._emit_log(
+            message=message, severity=LogSeverity.WARN, attributes=attributes
+        )
 
     def log_error(self, message: str, *, attributes: Attributes | None = None) -> None:
-        self._emit_log(message=message, severity=LogSeverity.ERROR, attributes=attributes)
+        self._emit_log(
+            message=message, severity=LogSeverity.ERROR, attributes=attributes
+        )
+
+    def log_exception(
+        self,
+        exception: BaseException,
+        *,
+        message: str | None = None,
+        attributes: Attributes | None = None,
+    ) -> None:
+        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
+        )
+        self.log_error(
+            message=f"error: {message}\nexception: {exception}{traceback_str}",
+            attributes=patched_attributes,
+        )
+
+    @contextmanager
+    def push_scope(
+        self, scope_name: str, *, attributes: Attributes | None = None
+    ) -> Generator[typing.Self, None, None]:
+        with get_otel_tracer().start_as_current_span(
+            scope_name, attributes=self._patch_attributes(attributes)
+        ):
+            yield self
 
 
 class JobLogger(Logger):
     def __init__(
         self,
         *,
+        base_span: Span,
         profile_metadata: job_definition_t.ProfileMetadata,
         job_definition: job_definition_t.JobDefinition,
     ) -> None:
         self.profile_metadata = profile_metadata
         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
@@ -105,6 +219,8 @@ class JobLogger(Logger):
                 patched_attributes["job.definition.cron_spec"] = (
                     self.job_definition.cron_spec
                 )
+            case job_definition_t.HttpJobDefinitionBase():
+                pass
             case _:
                 assert_never(self.job_definition)
         patched_attributes["job.definition.executor.type"] = (
@@ -123,13 +239,13 @@ class JobLogger(Logger):
                 assert_never(self.job_definition.executor)
         return _cast_attributes(patched_attributes)
 
-    @contextmanager
-    def push_scope(
-        self, scope_name: str, *, attributes: Attributes | None = None
-    ) -> Generator["JobLogger", None, None]:
-        with get_tracer().start_as_current_span(
-            scope_name, attributes=self._patch_attributes(attributes)
-        ) as span:
-            self.current_span_id = span.get_span_context().span_id
-            self.current_trace_id = span.get_span_context().trace_id
-            yield self
+
+@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

@@ -0,0 +1,97 @@
+import base64
+
+import flask
+from flask.typing import ResponseReturnValue
+from opentelemetry.trace import get_current_span
+from uncountable.core.environment import (
+    get_http_server_port,
+    get_server_env,
+)
+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.telemetry import Logger
+from uncountable.types import job_definition_t
+
+app = flask.Flask(__name__)
+
+
+def register_route(
+    *,
+    server_logger: Logger,
+    profile_meta: job_definition_t.ProfileMetadata,
+    job: job_definition_t.HttpJobDefinitionBase,
+) -> None:
+    route = f"/{profile_meta.name}/{job.id}"
+
+    def handle_request() -> ResponseReturnValue:
+        with server_logger.push_scope(route):
+            try:
+                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
+                )
+                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),
+                )
+                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
+                )
+
+                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 HttpException.unknown_error().make_error_response()
+
+    app.add_url_rule(
+        route,
+        endpoint=f"handle_request_{job.id}",
+        view_func=handle_request,
+        methods=["POST"],
+    )
+
+    server_logger.log_info(f"job {job.id} webhook registered at: {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.HttpJobDefinitionBase):
+                register_route(
+                    server_logger=server_logger, profile_meta=profile_metadata, job=job
+                )
+
+
+main()
+
+
+if __name__ == "__main__":
+    app.run(
+        host="0.0.0.0",
+        port=get_http_server_port(),
+        debug=get_server_env() == "playground",
+        exclude_patterns=[],
+    )