refactor: add structure to api routes

Author: ZohebShaikh

src/blueapi/service/main.py

@@ -0,0 +1,85 @@
+from typing import Annotated
+
+from fastapi import Depends, HTTPException, status
+from fastapi.requests import HTTPConnection
+from observability_utils.tracing import get_tracer
+
+from blueapi.config import ApplicationConfig
+from blueapi.service import interface
+from blueapi.service.authentication import Fedid
+from blueapi.service.authorization import OpaUserClient, opa
+from blueapi.service.model import WorkerTask
+from blueapi.service.runner import WorkerDispatcher
+
+TRACER = get_tracer("interface")
+
+RUNNER: WorkerDispatcher | None = None
+
+
+def _runner() -> WorkerDispatcher:
+    """Intended to be used only with FastAPI Depends"""
+    if RUNNER is None:
+        raise ValueError()
+    return RUNNER
+
+
+Runner = Annotated[WorkerDispatcher, Depends(_runner)]
+Opa = Annotated[OpaUserClient | None, Depends(opa)]
+
+
+def setup_runner(
+    config: ApplicationConfig | None = None,
+    runner: WorkerDispatcher | None = None,
+):
+    global RUNNER
+    runner = runner or WorkerDispatcher(config)
+    runner.start()
+
+    RUNNER = runner
+
+
+def teardown_runner():
+    global RUNNER
+    if RUNNER is None:
+        return
+    RUNNER.stop()
+    RUNNER = None
+
+
+async def access_task_permission(
+    opa: Opa,
+    task_id: str,
+    fedid: Fedid,
+    runner: Runner,
+):
+    task = runner.run(interface.get_task_by_id, task_id)
+
+    if (
+        opa
+        and not await opa.admin()
+        and (task and fedid != task.task.metadata.get("user"))
+    ):
+        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)
+
+
+# start_task_permission is used when there is WorkerTask
+async def start_task_permission(
+    task: WorkerTask,
+    opa: Opa,
+    fedid: Fedid,
+    runner: Runner,
+):
+    if not task.task_id:
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
+            detail="No task id provided",
+        )
+    await access_task_permission(opa, task.task_id, fedid, runner)
+
+
+def get_passthrough_headers(request: HTTPConnection) -> dict[str, str]:
+    return {
+        key: value
+        for key, value in request.headers.items()
+        if key.casefold() in ApplicationConfig.PROPAGATED_HEADERS
+    }

src/blueapi/service/routes/__init__.py

@@ -0,0 +1,25 @@
+from fastapi import APIRouter
+from observability_utils.tracing import start_as_current_span
+
+from blueapi.config import Tag
+from blueapi.service import interface
+from blueapi.service.model import DeviceModel, DeviceResponse
+
+from .dependencies import TRACER, Runner
+
+router = APIRouter(tags=[Tag.DEVICE])
+
+
+@router.get("/devices")
+@start_as_current_span(TRACER)
+def get_devices(runner: Runner) -> DeviceResponse:
+    """Retrieve information about all available devices."""
+    devices = runner.run(interface.get_devices)
+    return DeviceResponse(devices=devices)
+
+
+@router.get("/devices/{name}")
+@start_as_current_span(TRACER, "name")
+def get_device_by_name(name: str, runner: Runner) -> DeviceModel:
+    """Retrieve information about a devices by its (unique) name."""
+    return runner.run(interface.get_device, name)

src/blueapi/service/routes/dependencies.py

@@ -0,0 +1,48 @@
+from fastapi import APIRouter, BackgroundTasks
+from observability_utils.tracing import start_as_current_span
+
+from blueapi.config import Tag
+from blueapi.service import interface
+from blueapi.service.model import (
+    EnvironmentResponse,
+    PythonEnvironmentResponse,
+    SourceInfo,
+)
+
+from .dependencies import TRACER, Runner
+
+router = APIRouter(tags=[Tag.ENV])
+
+
+@router.get("/environment")
+@start_as_current_span(TRACER)
+def get_environment(runner: Runner) -> EnvironmentResponse:
+    """Get the current state of the environment, i.e. initialization state."""
+    return runner.state
+
+
+@router.delete("/environment")
+async def delete_environment(
+    background_tasks: BackgroundTasks,
+    runner: Runner,
+) -> EnvironmentResponse:
+    """Delete the current environment, causing internal components to be reloaded."""
+    environment_id = runner.state.environment_id
+    if runner.state.initialized or runner.state.error_message is not None:
+        background_tasks.add_task(runner.reload)
+    return EnvironmentResponse(environment_id=environment_id, initialized=False)
+
+
+@router.get("/python_environment")
+@start_as_current_span(TRACER)
+def get_python_environment(
+    runner: Runner,
+    name: str | None = None,
+    source: SourceInfo | None = None,
+) -> PythonEnvironmentResponse:
+    """
+    Retrieve the Python environment details.
+    This endpoint fetches information about the Python environment,
+    such as the installed packages and scratch packages.
+    """
+    return runner.run(interface.get_python_env, name, source)

src/blueapi/service/routes/devices.py

@@ -0,0 +1,59 @@
+import urllib.parse
+
+from fastapi import APIRouter, HTTPException, Response, status
+from fastapi.responses import RedirectResponse
+from observability_utils.tracing import start_as_current_span
+
+from blueapi.config import ApplicationConfig, OIDCConfig, Tag
+from blueapi.service import interface
+from blueapi.service.model import Health, HealthProbeResponse
+
+from .dependencies import TRACER, Runner
+
+open_router = APIRouter()
+router = APIRouter()
+
+
+@open_router.get("/", include_in_schema=False)
+def root_redirect() -> RedirectResponse:
+    """Redirect to docs url"""
+    return RedirectResponse(
+        status_code=status.HTTP_307_TEMPORARY_REDIRECT,
+        url=ApplicationConfig.DOCS_ENDPOINT,
+    )
+
+
+@open_router.get(
+    "/config/oidc",
+    tags=[Tag.META],
+    responses={
+        status.HTTP_204_NO_CONTENT: {"description": "No Authentication configured"}
+    },
+)
+@start_as_current_span(TRACER)
+def get_oidc_config(runner: Runner) -> OIDCConfig:
+    """Retrieve the OpenID Connect (OIDC) configuration for the server."""
+    config = runner.run(interface.get_oidc_config)
+    if config is None:
+        raise HTTPException(status_code=status.HTTP_204_NO_CONTENT)
+    return config
+
+
+@open_router.get("/healthz", status_code=status.HTTP_200_OK, tags=[Tag.META])
+def health_probe() -> HealthProbeResponse:
+    """If able to serve this, server is live and ready for requests."""
+    return HealthProbeResponse(status=Health.OK)
+
+
+@router.get("/logout", include_in_schema=False)
+def logout(runner: Runner) -> Response:
+    """Redirect to logout url"""
+    config = runner.run(interface.get_oidc_config)
+    if config is None or not config.logout_redirect_endpoint:
+        raise HTTPException(status_code=status.HTTP_205_RESET_CONTENT)
+
+    encoded_url = urllib.parse.quote_plus(config.end_session_endpoint)
+    return RedirectResponse(
+        status_code=status.HTTP_308_PERMANENT_REDIRECT,
+        url=config.logout_redirect_endpoint.rstrip("/") + "?rd=" + encoded_url,
+    )

src/blueapi/service/routes/environment.py

@@ -0,0 +1,25 @@
+from fastapi import APIRouter
+from observability_utils.tracing import start_as_current_span
+
+from blueapi.config import Tag
+from blueapi.service import interface
+from blueapi.service.model import PlanModel, PlanResponse
+
+from .dependencies import TRACER, Runner
+
+router = APIRouter(tags=[Tag.PLAN])
+
+
+@router.get("/plans")
+@start_as_current_span(TRACER)
+def get_plans(runner: Runner) -> PlanResponse:
+    """Retrieve information about all available plans."""
+    plans = runner.run(interface.get_plans)
+    return PlanResponse(plans=plans)
+
+
+@router.get("/plans/{name}")
+@start_as_current_span(TRACER, "name")
+def get_plan_by_name(name: str, runner: Runner) -> PlanModel:
+    """Retrieve information about a plan by its (unique) name."""
+    return runner.run(interface.get_plan, name)

src/blueapi/service/routes/meta.py

@@ -0,0 +1,105 @@
+import logging
+
+from fastapi import APIRouter, WebSocket, WebSocketDisconnect
+from pydantic import ValidationError
+from starlette.status import WS_1007_INVALID_FRAME_PAYLOAD_DATA, WS_1013_TRY_AGAIN_LATER
+
+from blueapi.service import interface, protocol
+from blueapi.service.authentication import Fedid
+from blueapi.service.model import WorkerTask
+from blueapi.service.protocol import (
+    InvalidArgs,
+    PlanNotFound,
+    ServerBusy,
+    Submit,
+    Unauthorized,
+    Update,
+)
+from blueapi.worker.event import WorkerEvent
+from blueapi.worker.task import TaskParameterError
+from blueapi.worker.worker_errors import WorkerBusyError
+
+from .dependencies import Opa, Runner, get_passthrough_headers
+
+LOGGER = logging.getLogger(__name__)
+
+router = APIRouter()
+
+
+@router.websocket("/run_plan")
+async def run_plan(
+    ws: WebSocket,
+    runner: Runner,
+    user: Fedid,
+    opa: Opa,
+):
+    LOGGER.info("Starting WS plan as %s", user)
+    await ws.accept()
+    rq = await ws.receive_text()
+    try:
+        task_request = Submit.model_validate_json(rq)
+    except ValidationError:
+        LOGGER.info("Failed to deserialize request: %r", rq, exc_info=True)
+        await ws.close(
+            code=WS_1007_INVALID_FRAME_PAYLOAD_DATA, reason="Invalid Request"
+        )
+        return
+    LOGGER.info("Plan request: %s", task_request)
+
+    if opa:
+        try:
+            await opa.can_submit_task(task_request.task)
+        except Exception as e:
+            LOGGER.info(
+                "User %s does not have permission to run task", user, exc_info=e
+            )
+            await ws.send_text(Unauthorized().model_dump_json())
+            await ws.close(code=protocol.AUTHZ_ERROR, reason="Unauthorized")
+            return
+
+    try:
+        task_id: str = runner.run(
+            interface.submit_task, task_request.task, {"user": user}
+        )
+        LOGGER.info("Task ID: %s", task_id)
+    except TaskParameterError as tpe:
+        LOGGER.info("Plan args not valid: %s - %s", task_request, tpe)
+        await ws.send_text(InvalidArgs.from_parameter_error(tpe).model_dump_json())
+        await ws.close(code=protocol.INVALID_ARGS, reason="Invalid Args")
+        return
+    except KeyError as ke:
+        LOGGER.info("Plan %r not recognised", ke.args[0])
+        await ws.send_text(PlanNotFound(plan_name=ke.args[0]).model_dump_json())
+        await ws.close(code=protocol.UNKNOWN_PLAN, reason="Unknown Plan")
+        return
+
+    try:
+        with runner.event_pipe() as events:
+            active_task = runner.run(interface.get_active_task)
+            if active_task is not None and not active_task.is_complete:
+                raise WorkerBusyError("Task already running")
+            runner.run(
+                interface.begin_task,
+                task=WorkerTask(task_id=task_id),
+                pass_through_headers=get_passthrough_headers(ws),
+            )
+            async for evt in events:
+                if evt.task_id != task_id:
+                    continue
+                LOGGER.debug("Event: %s", evt)
+                await ws.send_text(Update(data=evt).model_dump_json())
+                if isinstance(evt, WorkerEvent) and evt.is_complete():
+                    LOGGER.debug("End of stream")
+                    break
+    except WorkerBusyError:
+        LOGGER.error("Worker was busy")
+        await ws.send_text(ServerBusy().model_dump_json())
+        await ws.close(code=WS_1013_TRY_AGAIN_LATER, reason="Worker busy")
+    except WebSocketDisconnect:
+        LOGGER.info("Client disconnected")
+        runner.run(
+            interface.cancel_active_task, failure=True, reason="Client disconnected"
+        )
+    else:
+        LOGGER.info("Plan complete")
+        await ws.close()