diff --git a/boefjes/boefjes/__main__.py b/boefjes/boefjes/__main__.py index 0805fcf5158..e79fb806b77 100644 --- a/boefjes/boefjes/__main__.py +++ b/boefjes/boefjes/__main__.py @@ -1,11 +1,13 @@ import json import logging.config +import typing import click import structlog from boefjes.app import get_runtime_manager from boefjes.config import settings +from boefjes.runtime_interfaces import WorkerManager with settings.log_cfg.open() as f: logging.config.dictConfig(json.load(f)) @@ -34,9 +36,9 @@ @click.command() -@click.argument("worker_type", type=click.Choice(["boefje", "normalizer"])) +@click.argument("worker_type", type=click.Choice(typing.get_args(WorkerManager.WorkerType))) @click.option("--log-level", type=click.Choice(["DEBUG", "INFO", "WARNING", "ERROR"]), help="Log level", default="INFO") -def cli(worker_type: str, log_level: str) -> None: +def cli(worker_type: WorkerManager.WorkerType, log_level: str) -> None: logger.setLevel(log_level) logger.info("Starting runtime for %s", worker_type) diff --git a/boefjes/boefjes/app.py b/boefjes/boefjes/app.py index 86685453362..6c8a7058d50 100644 --- a/boefjes/boefjes/app.py +++ b/boefjes/boefjes/app.py @@ -265,7 +265,7 @@ def _start_working( logger.exception("Could not patch scheduler task to %s", status.value) -def get_runtime_manager(settings: Settings, queue: str, log_level: str) -> WorkerManager: +def get_runtime_manager(settings: Settings, queue: WorkerManager.WorkerType, log_level: str) -> WorkerManager: local_repository = get_local_repository() session = sessionmaker(bind=get_engine())() diff --git a/boefjes/boefjes/config.py b/boefjes/boefjes/config.py index 8299dad166f..d1e5bdc3c0a 100644 --- a/boefjes/boefjes/config.py +++ b/boefjes/boefjes/config.py @@ -8,6 +8,7 @@ from pydantic_settings.sources import EnvSettingsSource from boefjes.models import EncryptionMiddleware +from boefjes.runtime_interfaces import WorkerManager BASE_DIR: Path = Path(__file__).parent.resolve() @@ -132,7 +133,7 @@ class Settings(BaseSettings): examples=[[], ["ipv4", "wifi-pineapple"]], ) - runner_type: Literal["boefje", "normalizer"] = Field(None, examples=["boefje", "normalizer"]) + runner_type: WorkerManager.WorkerType | None = Field(None, examples=["boefje", "normalizer"]) logging_format: Literal["text", "json"] = Field("text", description="Logging format") diff --git a/boefjes/boefjes/runtime_interfaces.py b/boefjes/boefjes/runtime_interfaces.py index dab3c8bfdfc..be31610a370 100644 --- a/boefjes/boefjes/runtime_interfaces.py +++ b/boefjes/boefjes/runtime_interfaces.py @@ -1,3 +1,5 @@ +from typing import Literal + from boefjes.job_models import BoefjeMeta, NormalizerMeta, NormalizerResults @@ -17,7 +19,9 @@ def run(self, normalizer_meta: NormalizerMeta, raw: bytes) -> NormalizerResults: class WorkerManager: - def run(self, queue: str) -> None: + WorkerType = Literal["boefje", "normalizer"] + + def run(self, queue: WorkerType) -> None: raise NotImplementedError() diff --git a/boefjes/entrypoint.sh b/boefjes/entrypoint.sh index 4e15b8de5ac..968d2f4052a 100755 --- a/boefjes/entrypoint.sh +++ b/boefjes/entrypoint.sh @@ -17,7 +17,7 @@ if [ "$DATABASE_MIGRATION" = "1" ] || [[ $DATABASE_MIGRATION == "true" ]]; then fi if [ "$1" = "katalogus" ]; then - exec env BOEFJES_RUNNER_TYPE="$1" python -m uvicorn --host 0.0.0.0 boefjes.katalogus.root:app + exec python -m uvicorn --host 0.0.0.0 boefjes.katalogus.root:app fi exec "$@"