Flexible scheduling

boefjes/boefjes/api.py

@@ -122,7 +122,7 @@ async def boefje_output(
     bytes_client.save_boefje_meta(boefje_meta)
 
     if boefje_output.files:
-        mime_types = _default_mime_types(task.p_item.data.boefje)
+        mime_types = _default_mime_types(task.data.boefje)
         for file in boefje_output.files:
             raw = base64.b64decode(file.content)
             # when supported, also save file.name to Bytes
@@ -149,13 +149,13 @@ def get_task(task_id, scheduler_client):
 
 
 def create_boefje_meta(task, local_repository):
-    boefje = task.p_item.data.boefje
+    boefje = task.data.boefje
     boefje_resource = local_repository.by_id(boefje.id)
     env_keys = boefje_resource.environment_keys
-    environment = get_environment_settings(task.p_item.data, env_keys) if env_keys else {}
+    environment = get_environment_settings(task.data, env_keys) if env_keys else {}
 
-    organization = task.p_item.data.organization
-    input_ooi = task.p_item.data.input_ooi
+    organization = task.data.organization
+    input_ooi = task.data.input_ooi
     arguments = {"oci_arguments": boefje_resource.oci_arguments}
 
     if input_ooi:

boefjes/boefjes/app.py

@@ -9,12 +9,7 @@
 from httpx import HTTPError
 from pydantic import ValidationError
 
-from boefjes.clients.scheduler_client import (
-    QueuePrioritizedItem,
-    SchedulerAPIClient,
-    SchedulerClientInterface,
-    TaskStatus,
-)
+from boefjes.clients.scheduler_client import SchedulerAPIClient, SchedulerClientInterface, Task, TaskStatus
 from boefjes.config import Settings
 from boefjes.job_handler import BoefjeHandler, NormalizerHandler, bytes_api_client
 from boefjes.local import LocalBoefjeJobRunner, LocalNormalizerJobRunner
@@ -192,7 +187,7 @@ def exit(self, queue_type: WorkerManager.Queue, signum: int | None = None):
                 logger.info("Received %s, exiting", signal.Signals(signum).name)
 
             if not self.task_queue.empty():
-                items: list[QueuePrioritizedItem] = [self.task_queue.get() for _ in range(self.task_queue.qsize())]
+                items: list[Task] = [self.task_queue.get() for _ in range(self.task_queue.qsize())]
 
                 for p_item in items:
                     try:

boefjes/boefjes/clients/scheduler_client.py

@@ -13,18 +13,6 @@ class Queue(BaseModel):
     size: int
 
 
-class QueuePrioritizedItem(BaseModel):
-    """Representation of a queue.PrioritizedItem on the priority queue. Used
-    for unmarshalling of priority queue prioritized items to a JSON
-    representation.
-    """
-
-    id: uuid.UUID
-    priority: int
-    hash: str | None = None
-    data: BoefjeMeta | NormalizerMeta
-
-
 class TaskStatus(Enum):
     """Status of a task."""
 
@@ -34,14 +22,18 @@ class TaskStatus(Enum):
     RUNNING = "running"
     COMPLETED = "completed"
     FAILED = "failed"
+    CANCELLED = "cancelled"
 
 
 class Task(BaseModel):
     id: uuid.UUID
     scheduler_id: str
-    type: str
-    p_item: QueuePrioritizedItem
+    schedule_id: str | None
+    priority: int
     status: TaskStatus
+    type: str
+    hash: str | None = None
+    data: BoefjeMeta | NormalizerMeta
     created_at: datetime.datetime
     modified_at: datetime.datetime
 
@@ -50,7 +42,7 @@ class SchedulerClientInterface:
     def get_queues(self) -> list[Queue]:
         raise NotImplementedError()
 
-    def pop_item(self, queue: str) -> QueuePrioritizedItem | None:
+    def pop_item(self, queue: str) -> Task | None:
         raise NotImplementedError()
 
     def patch_task(self, task_id: uuid.UUID, status: TaskStatus) -> None:
@@ -59,7 +51,7 @@ def patch_task(self, task_id: uuid.UUID, status: TaskStatus) -> None:
     def get_task(self, task_id: uuid.UUID) -> Task:
         raise NotImplementedError()
 
-    def push_item(self, queue_id: str, p_item: QueuePrioritizedItem) -> None:
+    def push_item(self, queue_id: str, p_item: Task) -> None:
         raise NotImplementedError()
 
 
@@ -77,13 +69,13 @@ def get_queues(self) -> list[Queue]:
 
         return TypeAdapter(list[Queue]).validate_json(response.content)
 
-    def pop_item(self, queue: str) -> QueuePrioritizedItem | None:
+    def pop_item(self, queue: str) -> Task | None:
         response = self._session.post(f"/queues/{queue}/pop")
         self._verify_response(response)
 
-        return TypeAdapter(QueuePrioritizedItem | None).validate_json(response.content)
+        return TypeAdapter(Task | None).validate_json(response.content)
 
-    def push_item(self, queue_id: str, p_item: QueuePrioritizedItem) -> None:
+    def push_item(self, queue_id: str, p_item: Task) -> None:
         response = self._session.post(f"/queues/{queue_id}/push", content=p_item.json())
         self._verify_response(response)
 

boefjes/tests/conftest.py

@@ -14,7 +14,7 @@
 
 from boefjes.app import SchedulerWorkerManager
 from boefjes.clients.bytes_client import BytesAPIClient
-from boefjes.clients.scheduler_client import Queue, QueuePrioritizedItem, SchedulerClientInterface, Task, TaskStatus
+from boefjes.clients.scheduler_client import Queue, SchedulerClientInterface, Task, TaskStatus
 from boefjes.config import Settings, settings
 from boefjes.job_handler import bytes_api_client
 from boefjes.job_models import BoefjeMeta, NormalizerMeta
@@ -52,25 +52,25 @@ def __init__(
 
         self._iterations = 0
         self._tasks: dict[str, Task] = multiprocessing.Manager().dict()
-        self._popped_items: dict[str, QueuePrioritizedItem] = multiprocessing.Manager().dict()
-        self._pushed_items: dict[str, tuple[str, QueuePrioritizedItem]] = multiprocessing.Manager().dict()
+        self._popped_items: dict[str, Task] = multiprocessing.Manager().dict()
+        self._pushed_items: dict[str, tuple[str, Task]] = multiprocessing.Manager().dict()
 
     def get_queues(self) -> list[Queue]:
         time.sleep(self.sleep_time)
         return TypeAdapter(list[Queue]).validate_json(self.queue_response)
 
-    def pop_item(self, queue: str) -> QueuePrioritizedItem | None:
+    def pop_item(self, queue: str) -> Task | None:
         time.sleep(self.sleep_time)
 
         try:
             if WorkerManager.Queue.BOEFJES.value in queue:
-                p_item = TypeAdapter(QueuePrioritizedItem).validate_json(self.boefje_responses.pop(0))
+                p_item = TypeAdapter(Task).validate_json(self.boefje_responses.pop(0))
                 self._popped_items[str(p_item.id)] = p_item
                 self._tasks[str(p_item.id)] = self._task_from_id(p_item.id)
                 return p_item
 
             if WorkerManager.Queue.NORMALIZERS.value in queue:
-                p_item = TypeAdapter(QueuePrioritizedItem).validate_json(self.normalizer_responses.pop(0))
+                p_item = TypeAdapter(Task).validate_json(self.normalizer_responses.pop(0))
                 self._popped_items[str(p_item.id)] = p_item
                 return p_item
         except IndexError:
@@ -92,17 +92,9 @@ def get_task(self, task_id: UUID) -> Task:
         return self._task_from_id(task_id) if str(task_id) not in self._tasks else self._tasks[str(task_id)]
 
     def _task_from_id(self, task_id: UUID):
-        return Task(
-            id=task_id,
-            scheduler_id="test",
-            type="test",
-            p_item=self._popped_items[str(task_id)],
-            status=TaskStatus.DISPATCHED,
-            created_at=datetime.now(timezone.utc),
-            modified_at=datetime.now(timezone.utc),
-        )
+        return self._popped_items[str(task_id)]
 
-    def push_item(self, queue_id: str, p_item: QueuePrioritizedItem) -> None:
+    def push_item(self, queue_id: str, p_item: Task) -> None:
         self._pushed_items[str(p_item.id)] = (queue_id, p_item)
 
 
@@ -131,10 +123,12 @@ def item_handler(tmp_path: Path):
 @pytest.fixture
 def manager(item_handler: MockHandler, tmp_path: Path) -> SchedulerWorkerManager:
     scheduler_client = MockSchedulerClient(
-        get_dummy_data("scheduler/queues_response.json"),
-        2 * [get_dummy_data("scheduler/pop_response_boefje.json")] + [get_dummy_data("scheduler/should_crash.json")],
-        [get_dummy_data("scheduler/pop_response_normalizer.json")],
-        tmp_path / "patch_task_log",
+        queue_response=get_dummy_data("scheduler/queues_response.json"),
+        boefje_responses=(
+            2 * [get_dummy_data("scheduler/pop_response_boefje.json")] + [get_dummy_data("scheduler/should_crash.json")]
+        ),
+        normalizer_responses=[get_dummy_data("scheduler/pop_response_normalizer.json")],
+        log_path=tmp_path / "patch_task_log",
     )
 
     return SchedulerWorkerManager(item_handler, scheduler_client, Settings(pool_size=1, poll_interval=0.01), "DEBUG")

boefjes/tests/examples/scheduler/pop_response_boefje.json

@@ -1,6 +1,10 @@
 {
   "id": "70da7d4f-f41f-4940-901b-d98a92e9014b",
   "priority": 1,
+  "scheduler_id": "boefje-_dev",
+  "schedule_id": null,
+  "status": "dispatched",
+  "type": "boefje",
   "hash": "70da7d4f-f41f-4940-901b-d98a92e9014b",
   "data": {
     "id": "70da7d4f-f41f-4940-901b-d98a92e9014b",
@@ -15,5 +19,7 @@
     "runnable_hash": null,
     "environment": null,
     "ended_at": null
-  }
+  },
+  "created_at": "2021-06-29T14:00:00",
+  "modified_at": "2021-06-29T14:00:00"
 }

boefjes/tests/examples/scheduler/pop_response_boefje_no_ooi.json

@@ -1,6 +1,10 @@
 {
   "id": "70da7d4f-f41f-4940-901b-d98a92e9014b",
   "priority": 1,
+  "scheduler_id": "boefje-_dev",
+  "schedule_id": null,
+  "status": "dispatched",
+  "type": "boefje",
   "hash": "70da7d4f-f41f-4940-901b-d98a92e9014b",
   "data": {
     "id": "70da7d4f-f41f-4940-901b-d98a92e9014b",
@@ -15,5 +19,7 @@
     "runnable_hash": null,
     "environment": null,
     "ended_at": null
-  }
+  },
+  "created_at": "2021-06-29T14:00:00",
+  "modified_at": "2021-06-29T14:00:00"
 }

boefjes/tests/examples/scheduler/pop_response_normalizer.json

@@ -2,6 +2,9 @@
   "id": "60da7d4ff41f4940901bd98a92e9014b",
   "priority": 1,
   "scheduler_id": "normalizer-_dev",
+  "schedule_id": null,
+  "status": "dispatched",
+  "type": "normalizer",
   "hash": "7e698c377cfd85015c0d7086b76b76b4",
   "data": {
     "id": "60da7d4ff41f4940901bd98a92e9014b",
@@ -48,5 +51,7 @@
     "normalizer": {
       "id": "kat_dns_normalize"
     }
-  }
+  },
+  "created_at": "2021-06-29T14:00:00",
+  "modified_at": "2021-06-29T14:00:00"
 }

boefjes/tests/examples/scheduler/should_crash.json

@@ -2,6 +2,9 @@
   "id": "9071c9fd-2b9f-440f-a524-ef1ca4824fd4",
   "priority": 1,
   "scheduler_id": "boefje-_dev",
+  "schedule_id": null,
+  "status": "dispatched",
+  "type": "boefje",
   "hash": "7e698c377cfd85015c0d7086b76b76b4",
   "data": {
     "id": "9071c9fd-2b9f-440f-a524-ef1ca4824fd4",
@@ -34,5 +37,7 @@
     "input_ooi": "Hostname|internet|test.test",
     "organization": "_dev",
     "dispatches": []
-  }
+  },
+  "created_at": "2021-06-29T14:00:00",
+  "modified_at": "2021-06-29T14:00:00"
 }

boefjes/tests/test_api.py

@@ -8,10 +8,10 @@
 
 def _mocked_scheduler_client(tmp_path: Path):
     return MockSchedulerClient(
-        get_dummy_data("scheduler/queues_response.json"),
-        [get_dummy_data("scheduler/pop_response_boefje_no_ooi.json")],
-        [],
-        tmp_path / "patch_task_log",
+        queue_response=get_dummy_data("scheduler/queues_response.json"),
+        boefje_responses=[get_dummy_data("scheduler/pop_response_boefje_no_ooi.json")],
+        normalizer_responses=[],
+        log_path=tmp_path / "patch_task_log",
     )
 
 

mula/.ci/docker-compose.yml

@@ -28,7 +28,9 @@ services:
       - ci_rabbitmq
 
   ci_postgres:
-    image: docker.io/library/postgres:15
+    image: postgres:15
+    command:
+      ["postgres", "-c", "log_statement=all", "-c", "log_destination=stderr"]
     healthcheck:
       test: ["CMD", "gosu", "postgres", "pg_isready"]
       interval: 3s