Scheduler rate limit functionality

boefjes/boefjes/katalogus/models.py

@@ -2,7 +2,7 @@
 from enum import Enum
 from typing import List, Literal, NewType, Optional, Set, Union
 
-from pydantic import AnyHttpUrl, BaseModel, Field
+from pydantic import AnyHttpUrl, BaseModel, Field, constr
 
 RESERVED_LOCAL_ID = "LOCAL"
 
@@ -18,6 +18,11 @@ class Organisation(BaseModel):
     name: str
 
 
+class RateLimit(BaseModel):
+    interval: str
+    identifier: str = constr(min_length=1)
+
+
 class Plugin(BaseModel):
     id: str
     repository_id: str = RESERVED_LOCAL_ID
@@ -41,6 +46,7 @@ class Boefje(Plugin):
     produces: List[str] = Field(default_factory=list)
     options: Optional[List[str]]
     runnable_hash: Optional[str]
+    rate_limit: Optional[RateLimit]
 
 
 class Normalizer(Plugin):

mula/pyproject.toml

@@ -9,6 +9,7 @@ license = "EUPL"
 python = "^3.8"
 alembic = "^1.8.1"
 celery = "^5.2.7"
+limits = "^3.5.0"
 fastapi = "^0.101.0"
 mmh3 = "^4.0.0"
 pika = "^1.2.0"
@@ -28,6 +29,7 @@ opentelemetry-instrumentation-fastapi = "^0.40b0"
 opentelemetry-instrumentation-psycopg2 = "^0.40b0"
 opentelemetry-instrumentation-requests = "^0.40b0"
 certifi = "^2023.7.22"
+jinja2 = "^3.1.2"
 
 [tool.poetry.group.dev.dependencies]
 factory_boy = "^3.2.1"

mula/scheduler/alembic/versions/0008_add_delayed_status_for_tasks.py

@@ -0,0 +1,26 @@
+"""Add DELAYED status for tasks
+
+Revision ID: 0008
+Revises: 0007_add_cancelled_status_for_tasks
+Create Date: 2023-07-24 12:22:42.329666
+
+"""
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "0008"
+down_revision = "0007"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.execute("ALTER TYPE taskstatus ADD VALUE 'DELAYED'")
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.execute("ALTER TYPE taskstatus DROP VALUE 'DELAYED'")
+    # ### end Alembic commands ###

mula/scheduler/context/context.py

@@ -1,8 +1,10 @@
 import json
 import logging.config
+import threading
 from pathlib import Path
 from types import SimpleNamespace
 
+from limits import storage, strategies
 from prometheus_client import CollectorRegistry, Gauge, Info
 
 import scheduler
@@ -79,6 +81,10 @@ def __init__(self) -> None:
         self.task_store: stores.TaskStorer = sqlalchemy.TaskStore(datastore)
         self.pq_store: stores.PriorityQueueStorer = sqlalchemy.PriorityQueueStore(datastore)
 
+        # Rate limiter
+        self.rate_limiter = strategies.MovingWindowRateLimiter(storage=storage.MemoryStorage())
+        self.rate_limiter_lock = threading.Lock()
+
         # Metrics collector registry
         self.metrics_registry: CollectorRegistry = CollectorRegistry()
 

mula/scheduler/models/__init__.py

@@ -1,5 +1,5 @@
 from .base import Base
-from .boefje import Boefje, BoefjeMeta
+from .boefje import Boefje, BoefjeMeta, RateLimit
 from .events import NormalizerMetaReceivedEvent, RawData, RawDataReceivedEvent
 from .filter import Filter
 from .health import ServiceHealth

mula/scheduler/models/boefje.py

@@ -4,11 +4,17 @@
 from pydantic import BaseModel, Field
 
 
+class RateLimit(BaseModel):
+    identifier: str
+    interval: str
+
+
 class Boefje(BaseModel):
     """Boefje representation."""
 
     id: str
-    version: Optional[str] = Field(default=None)
+    version: Optional[str]
+    rate_limit: Optional[RateLimit]
 
 
 class BoefjeMeta(BaseModel):

mula/scheduler/models/tasks.py

@@ -24,6 +24,7 @@ class TaskStatus(str, enum.Enum):
     """Status of a task."""
 
     PENDING = "pending"
+    DELAYED = "delayed"
     QUEUED = "queued"
     DISPATCHED = "dispatched"
     RUNNING = "running"

mula/scheduler/repositories/sqlalchemy/task_store.py

@@ -1,6 +1,8 @@
 import datetime
 from typing import List, Optional, Tuple
 
+from sqlalchemy import asc, desc
+
 from scheduler import models
 
 from ..stores import TaskStorer  # noqa: TID252
@@ -27,6 +29,7 @@ def get_tasks(
         status: Optional[str] = None,
         min_created_at: Optional[datetime.datetime] = None,
         max_created_at: Optional[datetime.datetime] = None,
+        order_by: Optional[str] = None,
         filters: Optional[List[models.Filter]] = None,
     ) -> Tuple[List[models.Task], int]:
         with self.datastore.session.begin() as session:
@@ -49,6 +52,12 @@ def get_tasks(
             if max_created_at is not None:
                 query = query.filter(models.TaskORM.created_at <= max_created_at)
 
+            if order_by is not None:
+                if order_by.startswith("-"):
+                    query = query.order_by(desc(order_by[1:]))
+                else:
+                    query = query.order_by(asc(order_by))
+
             if filters is not None:
                 for f in filters:
                     query.filter(models.TaskORM.p_item[f.get_field()].astext == f.value)
@@ -120,6 +129,11 @@ def update_task(self, task: models.Task) -> None:
         with self.datastore.session.begin() as session:
             (session.query(models.TaskORM).filter(models.TaskORM.id == task.id).update(task.dict()))
 
+    @retry()
+    def update_task_status(self, task_id: str, status: models.TaskStatus) -> None:
+        with self.datastore.session.begin() as session:
+            session.query(models.TaskORM).filter(models.TaskORM.id == task_id).update({"status": status.name})
+
     @retry()
     def cancel_tasks(self, scheduler_id: str, task_ids: List[str]) -> None:
         with self.datastore.session.begin() as session:

mula/scheduler/repositories/stores.py

@@ -23,6 +23,7 @@ def get_tasks(
         status: Optional[str] = None,
         min_created_at: Optional[datetime.datetime] = None,
         max_created_at: Optional[datetime.datetime] = None,
+        order_by: Optional[str] = None,
         filters: Optional[List[models.Filter]] = None,
     ):
         raise NotImplementedError
@@ -47,6 +48,9 @@ def create_task(self, task: models.Task) -> Optional[models.Task]:
     def update_task(self, task: models.Task) -> Optional[models.Task]:
         raise NotImplementedError
 
+    def update_task_status(self, task_id: str, status: models.TaskStatus) -> None:
+        raise NotImplementedError
+
     def cancel_tasks(self, scheduler_id: str, task_ids: List[str]) -> None:
         raise NotImplementedError