refactor(robot-server): add experimental protocol router (#7856)

Co-authored-by: Max Marrone <[email protected]>
Co-authored-by: Sanniti Pimpley <[email protected]>

robot-server/robot_server/protocols/__init__.py

@@ -0,0 +1,18 @@
+"""Protocol file upload and management."""
+
+from opentrons.file_runner import ProtocolFileType
+
+from .router import protocols_router
+from .dependencies import get_protocol_store
+from .protocol_store import ProtocolStore, ProtocolResource
+
+__all__ = [
+    # main protocols router
+    "protocols_router",
+    # protocol state management
+    "get_protocol_store",
+    "ProtocolStore",
+    "ProtocolResource",
+    # convenience re-exports from opentrons
+    "ProtocolFileType",
+]

robot-server/robot_server/protocols/dependencies.py

@@ -0,0 +1,23 @@
+"""Protocol router dependency wire-up."""
+from fastapi import Request
+from logging import getLogger
+from pathlib import Path
+from tempfile import gettempdir
+from .protocol_store import ProtocolStore
+
+log = getLogger(__name__)
+
+PROTOCOL_STORE_KEY = "protocol_store"
+PROTOCOL_STORE_DIRECTORY = Path(gettempdir()) / "opentrons-protocols"
+
+
+def get_protocol_store(request: Request) -> ProtocolStore:
+    """Get a singleton ProtocolStore to keep track of created protocols."""
+    protocol_store = getattr(request.app.state, PROTOCOL_STORE_KEY, None)
+
+    if protocol_store is None:
+        log.info(f"Storing protocols in {PROTOCOL_STORE_DIRECTORY}")
+        protocol_store = ProtocolStore(directory=PROTOCOL_STORE_DIRECTORY)
+        setattr(request.app.state, PROTOCOL_STORE_KEY, protocol_store)
+
+    return protocol_store

robot-server/robot_server/protocols/protocol_models.py

@@ -0,0 +1,17 @@
+"""Protocol file models."""
+from __future__ import annotations
+from datetime import datetime
+from pydantic import Field
+from opentrons.file_runner import ProtocolFileType
+from robot_server.service.json_api import ResourceModel
+
+
+class Protocol(ResourceModel):
+    """A model representing an uploaded protocol resource."""
+
+    id: str = Field(..., description="A unique identifier for this protocol.")
+    createdAt: datetime = Field(..., description="When this protocol was uploaded.")
+    protocolType: ProtocolFileType = Field(
+        ...,
+        description="The type of protocol file (JSON or Python).",
+    )

robot-server/robot_server/protocols/protocol_store.py

@@ -0,0 +1,130 @@
+"""Methods for saving and retrieving protocol files."""
+from dataclasses import dataclass
+from datetime import datetime
+from fastapi import UploadFile
+from logging import getLogger
+from pathlib import Path
+from typing import Dict, List, Sequence
+
+from opentrons.file_runner import ProtocolFileType
+
+log = getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class ProtocolResource:
+    """An entry in the session store, used to construct response models."""
+
+    protocol_id: str
+    protocol_type: ProtocolFileType
+    created_at: datetime
+    files: List[Path]
+
+
+class ProtocolNotFoundError(KeyError):
+    """Error raised when a protocol ID was not found in the store."""
+
+    def __init__(self, protocol_id: str) -> None:
+        """Initialize the error message from the missing ID."""
+        super().__init__(f"Protocol {protocol_id} was not found.")
+
+
+class ProtocolFileInvalidError(ValueError):
+    """Error raised when a given source file is invalid.
+
+    May be caused by an empty filename.
+    """
+
+    pass
+
+
+class ProtocolStore:
+    """Methods for storing and retrieving protocol files."""
+
+    def __init__(self, directory: Path) -> None:
+        """Initialize the ProtocolStore.
+
+        Arguments:
+            directory: Directory in which to place created files.
+        """
+        self._directory = directory
+        self._protocols_by_id: Dict[str, ProtocolResource] = {}
+
+    async def create(
+        self,
+        protocol_id: str,
+        created_at: datetime,
+        files: Sequence[UploadFile],
+    ) -> ProtocolResource:
+        """Add a protocol to the store."""
+        protocol_dir = self._get_protocol_dir(protocol_id)
+        # TODO(mc, 2021-06-02): check for protocol collision
+        protocol_dir.mkdir(parents=True)
+        saved_files = []
+
+        for index, upload_file in enumerate(files):
+            if upload_file.filename == "":
+                raise ProtocolFileInvalidError(f"File {index} is missing a filename")
+
+            contents = await upload_file.read()
+            file_path = protocol_dir / upload_file.filename
+
+            if isinstance(contents, str):
+                file_path.write_text(contents, "utf-8")
+            else:
+                file_path.write_bytes(contents)
+
+            saved_files.append(file_path)
+
+        entry = ProtocolResource(
+            protocol_id=protocol_id,
+            protocol_type=self._get_protocol_type(saved_files),
+            created_at=created_at,
+            files=saved_files,
+        )
+
+        self._protocols_by_id[protocol_id] = entry
+
+        return entry
+
+    def get(self, protocol_id: str) -> ProtocolResource:
+        """Get a single protocol by ID."""
+        try:
+            return self._protocols_by_id[protocol_id]
+        except KeyError as e:
+            raise ProtocolNotFoundError(protocol_id) from e
+
+    def get_all(self) -> List[ProtocolResource]:
+        """Get all protocols currently saved in this store."""
+        return list(self._protocols_by_id.values())
+
+    def remove(self, protocol_id: str) -> ProtocolResource:
+        """Remove a protocol from the store."""
+        try:
+            entry = self._protocols_by_id.pop(protocol_id)
+        except KeyError as e:
+            raise ProtocolNotFoundError(protocol_id) from e
+
+        try:
+            for file_path in entry.files:
+                file_path.unlink()
+            self._get_protocol_dir(protocol_id).rmdir()
+        except Exception as e:
+            log.warn(f"Unable to delete all files for protocol {protocol_id}: {str(e)}")
+
+        return entry
+
+    def _get_protocol_dir(self, protocol_id: str) -> Path:
+        return self._directory / protocol_id
+
+    # TODO(mc, 2021-06-01): add python support, add multi-file support, and
+    # honestly, probably ditch all of this logic in favor of whatever
+    # `ProtocolAnalyzer` situation we come up with
+    @staticmethod
+    def _get_protocol_type(files: List[Path]) -> ProtocolFileType:
+        file_path = files[0]
+
+        if file_path.suffix == ".json":
+            return ProtocolFileType.JSON
+        else:
+            raise NotImplementedError()

robot-server/robot_server/protocols/response_builder.py

@@ -0,0 +1,23 @@
+"""Protocol response model factory."""
+from .protocol_store import ProtocolResource
+from .protocol_models import Protocol
+
+
+class ResponseBuilder:
+    """Interface to construct protocol resource models from data."""
+
+    @staticmethod
+    def build(protocol_entry: ProtocolResource) -> Protocol:
+        """Build a protocol resource model.
+
+        Arguments:
+            entry: Protocol data from the ProtocolStore.
+
+        Returns:
+            Protocol model representing the resource.
+        """
+        return Protocol(
+            id=protocol_entry.protocol_id,
+            protocolType=protocol_entry.protocol_type,
+            createdAt=protocol_entry.created_at,
+        )

robot-server/robot_server/protocols/router.py

@@ -0,0 +1,165 @@
+"""Router for /protocols endpoints."""
+from datetime import datetime
+from fastapi import APIRouter, Depends, File, UploadFile, status
+from typing import List
+from typing_extensions import Literal
+
+from robot_server.errors import ErrorDetails, ErrorResponse
+from robot_server.service.dependencies import get_unique_id, get_current_time
+from robot_server.service.json_api import (
+    ResponseModel,
+    MultiResponseModel,
+    EmptyResponseModel,
+)
+
+from .dependencies import get_protocol_store
+from .protocol_models import Protocol
+from .protocol_store import (
+    ProtocolStore,
+    ProtocolNotFoundError,
+    ProtocolFileInvalidError,
+)
+from .response_builder import ResponseBuilder
+
+
+class ProtocolNotFound(ErrorDetails):
+    """An error returned when a given protocol cannot be found."""
+
+    id: Literal["ProtocolNotFound"] = "ProtocolNotFound"
+    title: str = "Protocol Not Found"
+
+
+class ProtocolFileInvalid(ErrorDetails):
+    """An error returned when an uploaded protocol file is invalid."""
+
+    id: Literal["ProtocolFileInvalid"] = "ProtocolFileInvalid"
+    title: str = "Protocol File Invalid"
+
+
+protocols_router = APIRouter()
+
+
+@protocols_router.post(
+    path="/protocols",
+    summary="Upload a protocol",
+    status_code=status.HTTP_201_CREATED,
+    response_model=ResponseModel[Protocol],
+    responses={
+        status.HTTP_400_BAD_REQUEST: {"model": ErrorResponse[ProtocolFileInvalid]},
+    },
+)
+async def create_protocol(
+    files: List[UploadFile] = File(...),
+    response_builder: ResponseBuilder = Depends(ResponseBuilder),
+    protocol_store: ProtocolStore = Depends(get_protocol_store),
+    protocol_id: str = Depends(get_unique_id),
+    created_at: datetime = Depends(get_current_time),
+) -> ResponseModel[Protocol]:
+    """Create a new protocol by uploading its files.
+
+    Arguments:
+        files: List of uploaded files, from form-data.
+        response_builder: Interface to construct response models.
+        protocol_store: In-memory database of protocol resources.
+        protocol_id: Unique identifier to attach to the new resource.
+        created_at: Timestamp to attach to the new resource.
+    """
+    if len(files) > 1:
+        raise NotImplementedError("Multi-file protocols not yet supported.")
+    elif files[0].filename.endswith(".py"):
+        raise NotImplementedError("Python protocols not yet supported")
+
+    try:
+        protocol_entry = await protocol_store.create(
+            protocol_id=protocol_id,
+            created_at=created_at,
+            files=files,
+        )
+    except ProtocolFileInvalidError as e:
+        raise ProtocolFileInvalid(detail=str(e)).as_error(status.HTTP_400_BAD_REQUEST)
+
+    data = response_builder.build(protocol_entry)
+
+    return ResponseModel(data=data)
+
+
+@protocols_router.get(
+    path="/protocols",
+    summary="Get uploaded protocols",
+    status_code=status.HTTP_200_OK,
+    response_model=MultiResponseModel[Protocol],
+)
+async def get_protocols(
+    response_builder: ResponseBuilder = Depends(ResponseBuilder),
+    protocol_store: ProtocolStore = Depends(get_protocol_store),
+) -> MultiResponseModel[Protocol]:
+    """Get a list of all currently uploaded protocols.
+
+    Arguments:
+        response_builder: Interface to construct response models.
+        protocol_store: In-memory database of protocol resources.
+    """
+    protocol_entries = protocol_store.get_all()
+    data = [response_builder.build(e) for e in protocol_entries]
+
+    return MultiResponseModel(data=data)
+
+
+@protocols_router.get(
+    path="/protocols/{protocolId}",
+    summary="Get an uploaded protocol",
+    status_code=status.HTTP_200_OK,
+    response_model=ResponseModel[Protocol],
+    responses={
+        status.HTTP_404_NOT_FOUND: {"model": ErrorResponse[ProtocolNotFound]},
+    },
+)
+async def get_protocol_by_id(
+    protocolId: str,
+    response_builder: ResponseBuilder = Depends(ResponseBuilder),
+    protocol_store: ProtocolStore = Depends(get_protocol_store),
+) -> ResponseModel[Protocol]:
+    """Get an uploaded protocol by ID.
+
+    Arguments:
+        protocolId: Protocol identifier to fetch, pulled from URL.
+        response_builder: Interface to construct response models.
+        protocol_store: In-memory database of protocol resources.
+    """
+    try:
+        protocol_entry = protocol_store.get(protocol_id=protocolId)
+
+    except ProtocolNotFoundError as e:
+        raise ProtocolNotFound(detail=str(e)).as_error(status.HTTP_404_NOT_FOUND)
+
+    data = response_builder.build(protocol_entry)
+
+    return ResponseModel(data=data)
+
+
+@protocols_router.delete(
+    path="/protocols/{protocolId}",
+    summary="Delete an uploaded protocol",
+    status_code=status.HTTP_200_OK,
+    response_model=EmptyResponseModel,
+    responses={
+        status.HTTP_404_NOT_FOUND: {"model": ErrorResponse[ProtocolNotFound]},
+    },
+)
+async def delete_protocol_by_id(
+    protocolId: str,
+    protocol_store: ProtocolStore = Depends(get_protocol_store),
+) -> EmptyResponseModel:
+    """Delete an uploaded protocol by ID.
+
+    Arguments:
+        protocolId: Protocol identifier to delete, pulled from URL.
+        protocol_store: In-memory database of protocol resources.
+    """
+    try:
+        protocol_store.remove(protocol_id=protocolId)
+
+    except ProtocolNotFoundError as e:
+        raise ProtocolNotFound(detail=str(e)).as_error(status.HTTP_404_NOT_FOUND)
+
+    return EmptyResponseModel()