daq-queuing-service 0.1.13__py3-none-any.whl

daq_queuing_service/__init__.py

@@ -0,0 +1,11 @@
+"""Top level API.
+
+.. data:: __version__
+    :type: str
+
+    Version number as calculated by https://github.com/pypa/setuptools_scm
+"""
+
+from ._version import __version__
+
+__all__ = ["__version__"]

daq_queuing_service/__main__.py

@@ -0,0 +1,34 @@
+"""Interface for ``python -m daq_queuing_service``."""
+
+from argparse import ArgumentParser
+from collections.abc import Sequence
+from pathlib import Path
+
+import uvicorn
+
+from daq_queuing_service.app._config import get_default_config_path
+
+from . import __version__
+
+__all__ = ["main"]
+
+
+def main(args: Sequence[str] | None = None) -> None:
+    """Argument parser for the CLI."""
+    parser = ArgumentParser()
+    parser.add_argument("-v", "--version", action="version", version=__version__)
+    parser.add_argument("-p", "--port", type=int, default=8000)
+    parser.add_argument("--dev", action="store_true", default=False)
+    parser.add_argument("--config", type=Path, default=get_default_config_path())
+
+    parsed_args = parser.parse_args(args)
+
+    from daq_queuing_service.app.app import create_app
+
+    app = create_app(config_path=parsed_args.config, dev=parsed_args.dev)
+
+    uvicorn.run(app, host="0.0.0.0", port=parsed_args.port, workers=1)
+
+
+if __name__ == "__main__":
+    main()

daq_queuing_service/_version.py

@@ -0,0 +1,24 @@
+# file generated by vcs-versioning
+# don't change, don't track in version control
+from __future__ import annotations
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+version: str
+__version__: str
+__version_tuple__: tuple[int | str, ...]
+version_tuple: tuple[int | str, ...]
+commit_id: str | None
+__commit_id__: str | None
+
+__version__ = version = '0.1.13'
+__version_tuple__ = version_tuple = (0, 1, 13)
+
+__commit_id__ = commit_id = None

daq_queuing_service/api/api.py

@@ -0,0 +1,180 @@
+import asyncio
+import json
+from collections.abc import AsyncGenerator
+from typing import Annotated
+
+from blueapi.service.model import TaskRequest
+from fastapi import APIRouter, Depends, Request, Response
+from fastapi.responses import EventSourceResponse
+from pydantic import BaseModel
+
+from daq_queuing_service.app._config import AppConfig
+from daq_queuing_service.app.authentication import User
+from daq_queuing_service.blueapi_interaction.blueapi_call import BlueapiCallResponse
+from daq_queuing_service.broadcaster import Broadcaster
+from daq_queuing_service.plugins.converter import Converter, ValidateError
+from daq_queuing_service.task_queue.queue import (
+    QUEUE_EVENTS,
+    PauseReason,
+    QueueState,
+    TaskQueue,
+    TaskWithPosition,
+)
+from daq_queuing_service.task_queue.task import Experiment, Status, Task
+
+# pyright: reportUnusedFunction=false
+
+
+def get_current_user(request: Request) -> User | None:
+    if hasattr(request.state, "user"):
+        return request.state.user
+
+
+CurrentUser = Annotated[User | None, Depends(get_current_user)]
+
+
+class QueueStateUpdate(BaseModel):
+    paused: bool
+
+
+class TaskCancelRequest(BaseModel):
+    task_ids: list[str]
+
+
+def _filter_by_status(
+    tasks: list[TaskWithPosition], status: Status | None
+) -> list[TaskWithPosition]:
+    if status is None:
+        return tasks
+    return [task for task in tasks if task.status == status]
+
+
+def public_routes(queue: TaskQueue) -> APIRouter:
+    """No authentication is required to access these endpoints."""
+    router = APIRouter()
+
+    @router.get("/")
+    def read_root(request: Request):
+        base_url = str(request.base_url)
+        return (
+            f"Welcome to the daq queuing service. Visit {base_url}docs for Uvicorn API."
+        )
+
+    @router.get("/healthz")
+    async def healthz():
+        return Response()
+
+    @router.get("/queue/state")
+    def get_queue_state() -> QueueState:
+        return queue.state
+
+    return router
+
+
+def protected_routes(
+    queue: TaskQueue,
+    broadcaster: Broadcaster[QUEUE_EVENTS],
+    config: AppConfig,
+    converter: Converter,
+) -> APIRouter:
+    """Authentication is required to access these endpoints (if turned on in config).
+    Additionally, for endpoints that depend on whitelist_check, you must be in the
+    whitelist of authorised fedIDs to access them.
+    """
+    router = APIRouter()
+
+    @router.get("/config")
+    def get_config() -> AppConfig:
+        return config
+
+    @router.patch("/queue/state")
+    async def update_queue_state(payload: QueueStateUpdate) -> QueueState:
+        if payload.paused:
+            return await queue.pause_queue(PauseReason.USER_REQUESTED)
+        else:
+            return await queue.resume_queue()
+
+    @router.get("/queue")
+    async def get_queued_tasks(status: Status | None = None) -> list[TaskWithPosition]:
+        return _filter_by_status(await queue.get_queue(), status)
+
+    @router.post("/queue")
+    async def add_tasks_to_queue(
+        experiments: list[TaskRequest | Experiment],
+        user: CurrentUser,
+        position: int | None = None,
+    ) -> list[str]:
+        try:
+            converter.validate(experiments)
+        except Exception as e:
+            raise ValidateError(*e.args) from e
+
+        tasks = [Task(experiment=experiment, user=user) for experiment in experiments]
+        task_ids = [task.id for task in tasks]
+        await queue.add_tasks(tasks, position)
+        return task_ids
+
+    @router.delete("/queue")
+    async def cancel_all_tasks() -> list[TaskWithPosition]:
+        return await queue.cancel_all_tasks()
+
+    @router.post("/queue/move")
+    async def move_task(task_id: str, new_position: int) -> int:
+        return await queue.move_task(task_id, new_position)
+
+    @router.delete("/queue/tasks")
+    async def cancel_tasks(payload: TaskCancelRequest) -> list[TaskWithPosition]:
+        return await queue.cancel_tasks(payload.task_ids)
+
+    @router.get("/queue/{position}")
+    async def get_task_by_position(position: int) -> TaskWithPosition | None:
+        return await queue.get_task_by_position(position)
+
+    @router.get("/tasks")
+    async def get_all_tasks(status: Status | None = None) -> list[TaskWithPosition]:
+        return _filter_by_status(await queue.get_tasks(), status)
+
+    @router.get("/tasks/{task_id}")
+    async def get_task_by_id(task_id: str) -> TaskWithPosition:
+        return await queue.get_task_by_id(task_id)
+
+    @router.get("/history")
+    async def get_completed_tasks(
+        status: Status | None = None,
+    ) -> list[TaskWithPosition]:
+        return _filter_by_status(await queue.get_history(), status)
+
+    @router.delete("/history")
+    async def clear_history():
+        return await queue.clear_history()
+
+    @router.get("/call_queue")
+    async def get_call_queue() -> list[BlueapiCallResponse]:
+        return await queue.get_call_queue()
+
+    @router.get("/call_history")
+    async def get_call_history() -> list[BlueapiCallResponse]:
+        return await queue.get_call_history()
+
+    @router.get("/events")
+    async def stream_events() -> EventSourceResponse:
+        subscriber = broadcaster.subscribe()
+
+        async def event_generator() -> AsyncGenerator[str, None]:
+            try:
+                while True:
+                    event = await subscriber.get()
+                    event_str = (
+                        f"event: {event['type']}\ndata: {json.dumps(event['data'])}\n\n"
+                    )
+                    yield event_str
+
+            except asyncio.CancelledError:
+                # Client disconnected
+                raise
+            finally:
+                broadcaster.unsubscribe(subscriber)
+
+        return EventSourceResponse(event_generator())
+
+    return router

daq_queuing_service/api/errors.py

@@ -0,0 +1,71 @@
+from collections.abc import Awaitable, Callable
+from typing import TypeVar
+
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse
+
+from daq_queuing_service.log import LOGGER
+from daq_queuing_service.plugins.converter import ConverterError, ValidateError
+from daq_queuing_service.task_queue.queue_utils import (
+    NegativePositionError,
+    QueueError,
+    TaskInProgressError,
+    TaskNotFoundError,
+    TaskNotInQueueError,
+)
+
+# pyright: reportUnusedFunction=false
+
+E = TypeVar("E", bound=Exception)
+
+Handler = Callable[[Request, E], Awaitable[JSONResponse]]
+
+
+def make_exception_handler(
+    status_code: int, error_code: str
+) -> Callable[[Request, Exception], Awaitable[JSONResponse]]:
+    async def handler(request: Request, exception: Exception):
+        LOGGER.exception("Error while handling request: %s", request)
+        return JSONResponse(
+            status_code=status_code,
+            content={"error": error_code, "message": str(exception)},
+        )
+
+    return handler
+
+
+def register_exception_handlers(app: FastAPI):
+    app.add_exception_handler(
+        TaskInProgressError,
+        make_exception_handler(409, "task_in_progress"),
+    )
+
+    app.add_exception_handler(
+        TaskNotFoundError,
+        make_exception_handler(404, "task_not_found"),
+    )
+
+    app.add_exception_handler(
+        TaskNotInQueueError,
+        make_exception_handler(409, "task_not_in_queue"),
+    )
+
+    app.add_exception_handler(
+        NegativePositionError,
+        make_exception_handler(400, "negative_position"),
+    )
+
+    app.add_exception_handler(
+        QueueError,
+        make_exception_handler(409, "queue_error"),
+    )
+
+    app.add_exception_handler(
+        ValidateError,
+        make_exception_handler(422, "validation_error"),
+    )
+
+    app.add_exception_handler(
+        ConverterError,
+        make_exception_handler(422, "converter_error"),
+    )

daq_queuing_service/app/_config.py

@@ -0,0 +1,36 @@
+import os
+from pathlib import Path
+
+import yaml
+from blueapi.config import OIDCConfig, RestConfig, StompConfig
+from pydantic import BaseModel, Field
+
+CONFIG_PATH = "/etc/config/config.yaml"
+TEST_CONFIG_PATH = "tests/test_data/test_config.yaml"
+
+
+class ConverterConfig(BaseModel):
+    path: str
+    name: str
+
+
+class BlueapiConfig(BaseModel):
+    stomp: StompConfig = Field(default_factory=StompConfig)
+    api: RestConfig = Field(default_factory=RestConfig)
+
+
+class AppConfig(BaseModel):
+    blueapi: BlueapiConfig
+    converter: ConverterConfig
+    oidc: OIDCConfig | None = None
+    authorisation_whitelist: list[str] | None = None
+
+
+def get_default_config_path() -> str:
+    return CONFIG_PATH if os.path.isfile(CONFIG_PATH) else TEST_CONFIG_PATH
+
+
+def load_config(config_path: Path) -> AppConfig:
+    with open(config_path) as f:
+        data = yaml.safe_load(f)
+    return AppConfig(**data)

daq_queuing_service/app/app.py

@@ -0,0 +1,109 @@
+import asyncio
+import logging
+from contextlib import asynccontextmanager
+from pathlib import Path
+from typing import NoReturn
+
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.param_functions import Depends
+from fastapi.params import Depends as DependsType
+
+from daq_queuing_service.api.api import protected_routes, public_routes
+from daq_queuing_service.api.errors import register_exception_handlers
+from daq_queuing_service.app.authentication import (
+    build_access_token_check,
+    build_get_current_user,
+)
+from daq_queuing_service.app.authorisation import (
+    build_ensure_current_user_is_in_whitelist,
+)
+from daq_queuing_service.blueapi_interaction.blueapi_adapter import BlueapiClientAdapter
+from daq_queuing_service.blueapi_interaction.get_client import get_blueapi_client
+from daq_queuing_service.broadcaster import Broadcaster
+from daq_queuing_service.plugins.converter import get_converter
+from daq_queuing_service.task_queue.queue import QUEUE_EVENTS, TaskQueue
+from daq_queuing_service.worker.worker import QueueWorker
+
+from ._config import load_config
+
+logging.basicConfig(
+    level=logging.INFO, format="%(asctime)s [%(name)s] %(levelname)s: %(message)s"
+)
+
+
+def create_app(config_path: Path, dev: bool = False) -> FastAPI:
+    @asynccontextmanager
+    async def lifespan(app: FastAPI):
+        worker_task = asyncio.create_task(app.state.worker.run_loop())
+        app.state.worker_task = worker_task
+
+        def log_task_exception(task: asyncio.Task[NoReturn]):
+            try:
+                exc = task.exception()
+                if exc:
+                    logging.error("Worker crashed", exc_info=exc)
+            except asyncio.CancelledError:
+                pass
+
+        worker_task.add_done_callback(log_task_exception)
+        try:
+            yield
+        finally:
+            worker_task.cancel()
+            await asyncio.gather(worker_task, return_exceptions=True)
+
+    config = load_config(config_path)
+
+    broadcaster: Broadcaster[QUEUE_EVENTS] = Broadcaster()
+
+    converter_path = config.converter.path
+    converter_name = config.converter.name
+    converter = get_converter(converter_path, converter_name)
+
+    app = FastAPI(lifespan=lifespan)
+
+    dependencies: list[DependsType] = []
+    whitelist_check = None
+    if config.oidc:
+        validate_token = build_access_token_check(config.oidc)
+        get_current_user = build_get_current_user(validate_token)
+
+        app.swagger_ui_init_oauth = {
+            "clientId": "NOT_SUPPORTED",
+        }
+
+        whitelist_check = build_ensure_current_user_is_in_whitelist(
+            config.authorisation_whitelist, get_current_user
+        )
+
+        dependencies.append(Depends(get_current_user))
+        dependencies.append(Depends(whitelist_check))
+
+    if dev:  # Allows local client/UI through CORS
+        app.add_middleware(
+            CORSMiddleware,
+            allow_origin_regex=r"http://(localhost|127\.0\.0\.1):\d+",
+            allow_credentials=True,
+            allow_methods=["*"],
+            allow_headers=["*"],
+        )
+
+    app.state.queue = TaskQueue(converter, broadcaster)
+
+    blueapi_client = get_blueapi_client(config.blueapi)
+    blueapi_client_adapter = BlueapiClientAdapter(blueapi_client)
+
+    app.state.worker = QueueWorker(
+        queue=app.state.queue,
+        blueapi_client=blueapi_client_adapter,
+    )
+
+    register_exception_handlers(app)
+    app.include_router(public_routes(app.state.queue))
+    app.include_router(
+        protected_routes(app.state.queue, broadcaster, config, converter),
+        dependencies=dependencies,
+    )
+
+    return app

daq_queuing_service/app/authentication.py

@@ -0,0 +1,103 @@
+from collections.abc import Callable
+from typing import Annotated, Any
+
+import jwt
+from blueapi.config import OIDCConfig
+from fastapi import Depends, HTTPException, Request
+from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
+from jwt.exceptions import DecodeError, ExpiredSignatureError
+from pydantic import BaseModel, ValidationError
+from starlette.status import HTTP_401_UNAUTHORIZED
+
+from daq_queuing_service.log import LOGGER
+
+
+class User(BaseModel):
+    fedid: str
+    email: str | None = None
+    name: str | None = None
+
+
+# Some of the following code was copied from blueapi
+# See https://github.com/DiamondLightSource/blueapi/blob/2108ee0c89b4399d961106f7f23082a58d48a564/src/blueapi/service/authentication.py#L281-L340
+
+bearer_scheme = HTTPBearer(auto_error=False)
+
+
+def unchecked_bearer_token(
+    credentials: Annotated[HTTPAuthorizationCredentials | None, Depends(bearer_scheme)],
+) -> str | None:
+    if credentials is None:
+        return None
+    return credentials.credentials
+
+
+UncheckedBearerToken = Annotated[str | None, Depends(unchecked_bearer_token)]
+
+
+def build_access_token_check(
+    config: OIDCConfig,
+) -> Callable[[UncheckedBearerToken], dict[str, Any]]:
+    """
+    Create a function to validate the bearer token of requests
+
+    The returned function should be used via fastAPI's 'Depends' mechanism to
+    ensure users are authenticated
+    """
+    LOGGER.info(f"JWKS URI: {config.jwks_uri}")
+    jwkclient = jwt.PyJWKClient(config.jwks_uri)
+
+    def validate_bearer_token(token: UncheckedBearerToken):
+        """Check that a bearer token is valid and inject into request state"""
+        if not token:
+            raise HTTPException(
+                status_code=HTTP_401_UNAUTHORIZED,
+                detail="Not authenticated",
+                headers={"WWW-Authenticate": "Bearer"},
+            )
+
+        try:
+            signing_key = jwkclient.get_signing_key_from_jwt(token)
+        except DecodeError as e:
+            raise HTTPException(
+                status_code=HTTP_401_UNAUTHORIZED,
+                detail="Cannot decode token",
+            ) from e
+        except ExpiredSignatureError as e:
+            raise HTTPException(
+                status_code=HTTP_401_UNAUTHORIZED,
+                detail="Token expired",
+            ) from e
+
+        decoded: dict[str, Any] = jwt.decode(
+            token,
+            signing_key.key,
+            algorithms=config.id_token_signing_alg_values_supported,
+            verify=True,
+            audience=config.client_audience,
+            issuer=config.issuer,
+        )
+        LOGGER.debug(f"Decoded valid token: {decoded}")
+        return decoded
+
+    return validate_bearer_token
+
+
+def build_get_current_user(
+    validate_token: Callable[..., dict[str, Any]],
+) -> Callable[[Request, dict[str, Any]], User]:
+    def get_current_user(
+        request: Request,
+        decoded: Annotated[dict[str, Any], Depends(validate_token)],
+    ) -> User:
+        try:
+            user = User.model_validate(decoded)
+        except ValidationError as e:
+            raise HTTPException(
+                status_code=HTTP_401_UNAUTHORIZED,
+                detail="Invalid token claims",
+            ) from e
+        request.state.user = user
+        return user
+
+    return get_current_user

daq_queuing_service/app/authorisation.py

@@ -0,0 +1,31 @@
+from collections.abc import Callable
+from typing import Annotated
+
+from fastapi import Depends, HTTPException
+from starlette.status import HTTP_403_FORBIDDEN
+
+from daq_queuing_service.app.authentication import User
+from daq_queuing_service.worker.worker import LOGGER
+
+
+def build_ensure_current_user_is_in_whitelist(
+    whitelist: list[str] | None, get_current_user: Callable[..., User]
+) -> Callable[[User], User]:
+    def ensure_current_user_is_in_whitelist(
+        current_user: Annotated[User, Depends(get_current_user)],
+    ) -> User:
+        LOGGER.debug(f"Got user: {current_user}")
+        if whitelist is None:
+            LOGGER.debug("No user whitelist. All authenticated users are authorised.")
+            return current_user
+        elif current_user.fedid in whitelist:
+            LOGGER.debug(
+                f"FedID {current_user.fedid} found in whitelist, user authorised."
+            )
+            return current_user
+        raise HTTPException(
+            status_code=HTTP_403_FORBIDDEN,
+            detail="Not authorised. You are not in the whitelist of authorised FedIDs",
+        )
+
+    return ensure_current_user_is_in_whitelist

daq_queuing_service/blueapi_interaction/blueapi_adapter.py

@@ -0,0 +1,70 @@
+import asyncio
+from dataclasses import dataclass
+from typing import Generic, TypeVar
+
+from blueapi.client import BlueapiClient
+from blueapi.client.event_bus import OnAnyEvent
+from blueapi.client.rest import (
+    BlueskyRemoteControlError,
+    InvalidParametersError,
+    ServiceUnavailableError,
+    UnknownPlanError,
+)
+from blueapi.service.model import TaskRequest
+from blueapi.worker import TaskStatus, WorkerState
+
+from daq_queuing_service.log import LOGGER
+
+T = TypeVar("T")
+E = TypeVar("E", bound=Exception)
+
+
+@dataclass(frozen=True)
+class BlueapiResult(Generic[T, E]):
+    value: T | None = None
+    error: E | None = None
+
+    def __post_init__(self):
+        if (self.value is None) == (self.error is None):
+            raise ValueError("Exactly one of value or error must be set")
+
+
+class BlueapiClientAdapter:
+    def __init__(
+        self,
+        client: BlueapiClient,
+    ):
+        self._client = client
+
+    async def get_state(self) -> BlueapiResult[WorkerState, ServiceUnavailableError]:
+        try:
+            state = await asyncio.to_thread(self._client.get_state)
+            return BlueapiResult(value=state)
+        except ServiceUnavailableError as e:
+            LOGGER.error(f"Lost connection to blueapi: {e}")
+            return BlueapiResult(error=e)
+
+    async def run_task(
+        self,
+        task_request: TaskRequest,
+        on_event: OnAnyEvent | None = None,
+    ) -> BlueapiResult[
+        TaskStatus,
+        BlueskyRemoteControlError
+        | InvalidParametersError
+        | UnknownPlanError
+        | ServiceUnavailableError,
+    ]:
+        try:
+            task_status = await asyncio.to_thread(
+                self._client.run_task, task_request, on_event=on_event
+            )
+            return BlueapiResult(value=task_status)
+        except (
+            BlueskyRemoteControlError,
+            InvalidParametersError,
+            UnknownPlanError,
+            ServiceUnavailableError,
+        ) as e:
+            LOGGER.exception(e)
+            return BlueapiResult(error=e)