Merge pull request #81 from ScilifelabDataCentre/celery-pause

SQ787: Prevent adding new jobs to the queue prior to doing migrations/upgrades

Author: RMCrean

docs/development/celery_task_implementation.md

@@ -185,6 +185,7 @@ class TaskHistoryResult(BaseModel):
   ```
 
 - The dependency injection in `project_and_user_and_role` and the helper function `has_required_role` are used to connect to the database and check that the user has permission to submit tasks to this project. See the example below for how to use it in an endpoint function.
+- Tasks should also call the `check_queue_closed_for_new_tasks` fn to validate the queue is not down (i.e. due to planned maintenance). This function raises a `QueueClosedError` if the queue is closed, which is handled in the API layer to return a 400 error with a message to the user.
 - To enqueue the task function defined in `tasks.py` in the job system, use `result = <TASK_FUNCTION>.apply_async()`. This returns some initial Celery task metadata, including the Celery task UUID.
   - The established pattern in DivBase is - for clarity - to call `.apply_async()` with keyword arguments and not arguments. Specifically, the Pydantic kwargs model can be populated, and then dumped in to Python dict in `.apply_async()` since Celery cannot de/serialise Pydantic mdoels. This ensures that kwarg types are validated before enqueuing the task.
 - Call the `create_task_history_entry()` CRUD function to record the Celery task UUID in `TaskHistoryDB` along with user and project ID. This function will return the DivBase task ID, which is the autoincrementing id from the postgreSQL table. This is an integer and much easier for the users to handle than long UUIDs.
@@ -211,6 +212,9 @@ async def create_bcftools_jobs(
     if not has_required_role(role, ProjectRoles.EDIT):
         raise AuthorizationError("You don't have permission to query this project.")
 
+    # validation queue is not closed, will raise custom error to user if so.
+    await check_queue_closed_for_new_tasks(db=db, is_admin=current_user.is_admin)
+
     # Pack the required task arguments in the corresponding Pydantic model (see Section 2.2.)
     # This ensures that the kwargs are type validated
     task_kwargs = BcftoolsQueryKwargs(

packages/divbase-api/src/divbase_api/admin_panel.py

@@ -40,6 +40,7 @@
 from divbase_api.models.announcements import AnnouncementDB, AnnouncementLevel, AnnouncementTarget
 from divbase_api.models.project_versions import ProjectVersionDB
 from divbase_api.models.projects import ProjectDB, ProjectMembershipDB, ProjectRoles
+from divbase_api.models.queue_status import QueueStatus
 from divbase_api.models.revoked_tokens import RevokedTokenDB, TokenRevokeReason
 from divbase_api.models.task_history import CeleryTaskMeta, TaskHistoryDB, TaskStartedAtDB
 from divbase_api.models.users import UserDB
@@ -628,6 +629,101 @@ class AnnouncementView(ModelView):
     exclude_fields_from_detail = []
 
 
+class QueueStatusView(ModelView):
+    """
+    Custom admin panel View for the QueueStatus model.
+
+    This is a singleton table (only 1 row allowed).
+    Deletion and creation are disabled - only editing the existing row is allowed.
+    """
+
+    page_size_options = PAGINATION_DEFAULTS
+    fields = [
+        IntegerField("id", disabled=True),
+        BooleanField(
+            "is_closed",
+            help_text=(
+                "Is the queue closed for new tasks? "
+                "If the scheduled start time is in the past or not set, the queue is closed. "
+                "If scheduled start time is in the future, queue will be closed at that time."
+            ),
+        ),
+        DateTimeField(
+            "scheduled_start",
+            label="Scheduled closure start time (UTC time)",
+            help_text=(
+                "(THIS IS UTC TIME) Optional: When should the queue closure set above take effect? "
+                "Leave empty for queue to be closed immediatialy "
+                "If the queue is open and you set this field, it will do nothing."
+            ),
+        ),
+        TextAreaField(
+            "reason_for_users",
+            help_text=(
+                "Message shown to users when the queue is closed (max 500 characters). "
+                "You could write: "
+                "The queuing system is currently closed for new tasks due to planned upcoming maintenance. Please try again later."
+            ),
+        ),
+        DateTimeField(
+            "created_at",
+            label="Created At (UTC time)",
+            help_text="Timestamp when the entry was created. Value determined by system.",
+            disabled=True,
+        ),
+        DateTimeField(
+            "updated_at",
+            label="Updated At (UTC time)",
+            help_text="Timestamp when the entry was last updated. Value determined by system.",
+            disabled=True,
+        ),
+    ]
+
+    exclude_fields_from_list = []
+    exclude_fields_from_edit = ["id", "created_at", "updated_at"]
+    exclude_fields_from_detail = []
+
+    def can_delete(self, request: Request) -> bool:
+        """Disable deletion - this is a singleton table."""
+        return False
+
+    def can_create(self, request: Request) -> bool:
+        """Disable creation - this is a singleton table with only 1 row."""
+        # 1st row created by the alembic migration
+        return False
+
+    async def edit(self, request: Request, pk: Any, data: dict) -> Any:
+        """
+        Override the default edit method to ensure that the `scheduled_start` field is set to None
+        if `is_closed` is set to False.
+        """
+        if not data.get("is_closed"):
+            data["scheduled_start"] = None
+
+        return await super().edit(request=request, pk=pk, data=data)
+
+    async def validate(self, request: Request, data: dict[str, Any]) -> None:
+        errors: dict[str, str] = {}
+
+        if data.get("scheduled_start") and not data.get("is_closed"):
+            errors["scheduled_start"] = "Cannot have a scheduled start time without the queue being closed."
+
+        reason = data.get("reason_for_users") or ""
+        if data.get("is_closed") and not reason.strip():
+            errors["reason_for_users"] = "Reason for users is required when the queue is closed."
+
+        if len(reason) > 500:
+            errors["reason_for_users"] = "Message must be 500 characters or less."
+
+        if errors:
+            raise FormValidationError(errors=errors)
+
+        return await super().validate(request=request, data=data)
+
+    # NOTE, no serialize_field_value/_format_cet_datetime override on purpose here
+    # as does not work well when modifying the timestamp in edit view, easier to just display as UTC
+
+
 class DivBaseAuthProvider(AuthProvider):
     """
     This class enables starlette-admin to make use of DivBase's pre-existing auth system.
@@ -711,4 +807,5 @@ def register_admin_panel(app: FastAPI, engine: AsyncEngine) -> None:
     admin.add_view(
         AnnouncementView(AnnouncementDB, icon="fas fa-bullhorn", label="Announcements", identity="announcement")
     )
+    admin.add_view(QueueStatusView(QueueStatus, icon="fas fa-power-off", label="Queue Status", identity="queue-status"))
     admin.mount_to(app)

packages/divbase-api/src/divbase_api/crud/queue_status.py

@@ -0,0 +1,39 @@
+"""
+Get the current status of the queuing system as defined in the QueueStatus model.
+
+This is used to determine if we should reject new tasks from being created when the queue is closed for/ahead of planned maintenance.
+
+CREATE, EDIT and DELETE all handled by admin panel.
+"""
+
+from datetime import datetime
+
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from divbase_api.exceptions import QueueClosedError
+from divbase_api.models.queue_status import QueueStatus
+
+
+async def check_queue_closed_for_new_tasks(db: AsyncSession, is_admin: bool) -> None:
+    """
+    Checks if the queue is closed for new tasks based on the QueueStatus model.
+
+    Will raise a QueueClosedError if so.
+    """
+    if is_admin:
+        # admins can always submit new tasks to allow for testing and validation during maintenance periods
+        return
+
+    result = await db.execute(select(QueueStatus).filter_by(id=1))
+    status = result.scalar_one_or_none()
+
+    if not status or not status.is_closed:
+        return
+
+    if not status.scheduled_start:  # effective immediately if no scheduled start time
+        raise QueueClosedError(message=status.reason_for_users)
+
+    # if the queue is scheduled to be closed in the future, we can still allow new tasks to be created until that time
+    if status.scheduled_start < datetime.now(tz=status.scheduled_start.tzinfo):
+        raise QueueClosedError(message=status.reason_for_users)

packages/divbase-api/src/divbase_api/exception_handlers.py

@@ -30,6 +30,7 @@
     ProjectVersionCreationError,
     ProjectVersionNotFoundError,
     ProjectVersionSoftDeletedError,
+    QueueClosedError,
     TaskNotFoundInBackendError,
     TooManyObjectsInRequestError,
     UserRegistrationError,
@@ -378,6 +379,18 @@ async def object_does_not_exist_error_handler(request: Request, exc: ObjectDoesN
         )
 
 
+async def queue_closed_error_handler(request: Request, exc: QueueClosedError):
+    logger.debug(f"Queue closed error for {request.method} {request.url.path}: {exc.message}", exc_info=False)
+    if is_api_request(request):
+        return JSONResponse(
+            status_code=exc.status_code,
+            content={"detail": exc.message, "type": "queue_closed_error"},
+            headers=exc.headers,
+        )
+    else:
+        return await render_error_page(request, exc.message, status_code=exc.status_code)
+
+
 def register_exception_handlers(app: FastAPI) -> None:
     """
     Register all exception handlers with FastAPI app.
@@ -402,6 +415,7 @@ def register_exception_handlers(app: FastAPI) -> None:
     app.add_exception_handler(TaskNotFoundInBackendError, task_not_found_in_backend_error_handler)  # type: ignore
     app.add_exception_handler(DownloadedFileChecksumMismatchError, downloaded_file_checksum_mismatch_error_handler)  # type: ignore
     app.add_exception_handler(ObjectDoesNotExistError, object_does_not_exist_error_handler)  # type: ignore
+    app.add_exception_handler(QueueClosedError, queue_closed_error_handler)  # type: ignore
 
     # These cover more generic/unexpected HTTP errors - the exceptions above take precedence
     app.add_exception_handler(HTTPException, generic_http_exception_handler)  # type: ignore

packages/divbase-api/src/divbase_api/exceptions.py

@@ -162,3 +162,10 @@ def __init__(self, filename: str | None = None, project_name: str | None = None)
                 "Please check that you have spelled the file name correctly and that the file has been uploaded to the project."
             )
         super().__init__(message=message, status_code=status.HTTP_404_NOT_FOUND)
+
+
+class QueueClosedError(DivBaseAPIException):
+    """Raised when the queue is closed for new tasks (e.g. before a maintenance window)."""
+
+    def __init__(self, message: str):
+        super().__init__(message=message, status_code=status.HTTP_400_BAD_REQUEST)

packages/divbase-api/src/divbase_api/middleware.py

@@ -84,7 +84,7 @@ def register_middleware(app: FastAPI) -> None:
     """
     Register all middleware to the FastAPI app.
 
-    NOTE: When modyfying/extending this function remember that Middleware order is important.
+    NOTE: When modifying/extending this function remember that Middleware order is important.
 
     So for the request-response cycle:
     request -> last added middleware -> ... -> first added middleware -> route handler

packages/divbase-api/src/divbase_api/migrations/versions/2026-03-19_add_queue_status_db_model.py

@@ -0,0 +1,53 @@
+"""add_queue_status_db_model
+
+Revision ID: 755a240de0ab
+Revises: 03964a15e9ed
+Create Date: 2026-03-19 11:57:16.792455
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "755a240de0ab"
+down_revision: Union[str, Sequence[str], None] = "03964a15e9ed"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    """Upgrade schema."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    # NOTE: The 'id' is set to autoincrement, which is not actually needed
+    # This is because we are inheriting it from basedb model and it does no harm to have autoincrement=true.
+    op.create_table(
+        "queue_status",
+        sa.Column("is_closed", sa.Boolean(), nullable=False),
+        sa.Column("scheduled_start", sa.DateTime(timezone=True), nullable=True),
+        sa.Column("reason_for_users", sa.String(length=500), nullable=False),
+        sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
+        sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
+        sa.Column("updated_at", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
+        sa.CheckConstraint("id = 1", name="queue_status_singleton_row"),
+        sa.PrimaryKeyConstraint("id"),
+    )
+    op.create_index(op.f("ix_queue_status_id"), "queue_status", ["id"], unique=False)
+    # ### end Alembic commands ###
+    ### human added - create the initial singleton row
+    conn = op.get_bind()
+    conn.execute(
+        sa.text(
+            "INSERT INTO queue_status (id, is_closed, reason_for_users) VALUES (1, false, 'The queuing system is currently closed for new tasks due to planned upcoming maintenance. Please try again later.')"
+        )
+    )
+
+
+def downgrade() -> None:
+    """Downgrade schema."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_index(op.f("ix_queue_status_id"), table_name="queue_status")
+    op.drop_table("queue_status")
+    # ### end Alembic commands ###

packages/divbase-api/src/divbase_api/models/__init__.py

@@ -8,10 +8,11 @@
 from divbase_api.models.base import Base, BaseDBModel
 from divbase_api.models.project_versions import ProjectVersionDB
 from divbase_api.models.projects import ProjectDB, ProjectMembershipDB, ProjectRoles
+from divbase_api.models.queue_status import QueueStatus
 from divbase_api.models.revoked_tokens import RevokedTokenDB, TokenRevokeReason
 from divbase_api.models.task_history import TaskHistoryDB, TaskStartedAtDB
 from divbase_api.models.users import UserDB
-from divbase_api.models.vcf_dimensions import SkippedVCFDB, VCFMetadataDB
+from divbase_api.models.vcf_dimensions import SkippedVCFDB, VCFMetadataDB, VCFMetadataSamplesDB, VCFMetadataScaffoldsDB
 
 __all__ = [
     "Base",
@@ -22,13 +23,14 @@
     "ProjectMembershipDB",
     "VCFMetadataDB",
     "SkippedVCFDB",
+    "VCFMetadataSamplesDB",
+    "VCFMetadataScaffoldsDB",
     "TaskHistoryDB",
     "TaskStartedAtDB",
     "TokenRevokeReason",
     "RevokedTokenDB",
     "ProjectVersionDB",
     "AnnouncementDB",
     "AnnouncementTarget",
-    "VCFMetadataSamplesDB",
-    "VCFMetadataScaffoldsDB",
+    "QueueStatus",
 ]

packages/divbase-api/src/divbase_api/models/queue_status.py

@@ -0,0 +1,28 @@
+"""
+QueueStatus db model is used to define whether the DivBase queuing system is closed to new tasks.
+
+The use case for this is prior to a scheduled upgrade, we may want to empty the queue:
+1. (Day before upgrade): Go to admin panel, set is_closed to True (optionally schedule when the queue will be closed).
+2. (Day of upgrade): Validate queue is empty, perform upgrade,
+3. Once happy, go to admin panel, set is_closed to False to allow new tasks to be created again.
+
+NOTE: Admins can submit jobs even when the queue is closed to allow for validation and testing during the maintenance period
+"""
+
+from datetime import datetime
+
+from sqlalchemy import CheckConstraint, DateTime, String
+from sqlalchemy.orm import Mapped, mapped_column
+
+from divbase_api.models.base import BaseDBModel
+
+
+class QueueStatus(BaseDBModel):
+    __tablename__ = "queue_status"
+
+    is_closed: Mapped[bool] = mapped_column(default=False)
+    scheduled_start: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
+    reason_for_users: Mapped[str] = mapped_column(String(500))
+
+    # this arg makes the table a singleton, so only 1 row (with id = 1) in the table
+    __table_args__ = (CheckConstraint("id = 1", name="queue_status_singleton_row"),)

packages/divbase-api/src/divbase_api/routes/queries.py

@@ -12,6 +12,7 @@
 
 from divbase_api.api_config import settings
 from divbase_api.crud.projects import has_required_role
+from divbase_api.crud.queue_status import check_queue_closed_for_new_tasks
 from divbase_api.crud.task_history import create_task_history_entry, update_task_history_entry_with_celery_task_id
 from divbase_api.db import get_db
 from divbase_api.deps import get_project_member
@@ -65,6 +66,7 @@ async def sample_metadata_query(
 
     if not has_required_role(role, ProjectRoles.EDIT):
         raise AuthorizationError("You don't have permission to query this project.")
+    await check_queue_closed_for_new_tasks(db=db, is_admin=current_user.is_admin)
 
     task_kwargs = SampleMetadataQueryKwargs(
         tsv_filter=sample_metadata_query_request.tsv_filter,
@@ -128,6 +130,7 @@ async def create_bcftools_jobs(
 
     if not has_required_role(role, ProjectRoles.EDIT):
         raise AuthorizationError("You don't have permission to query this project.")
+    await check_queue_closed_for_new_tasks(db=db, is_admin=current_user.is_admin)
 
     job_id = await create_task_history_entry(
         user_id=current_user.id,