feat(robot-server): Allow adding multiple labware offsets in a single…

… request (#17436)
Opentrons · Feb 6, 2025 · ba07963 · ba07963
1 parent aae2bfa
commit ba07963
Show file tree

Hide file tree

Showing 9 changed files with 345 additions and 82 deletions.
diff --git a/api-client/src/runs/createLabwareOffset.ts b/api-client/src/runs/createLabwareOffset.ts
@@ -2,17 +2,25 @@ import { POST, request } from '../request'
 
 import type { ResponsePromise } from '../request'
 import type { HostConfig } from '../types'
-import type { LegacyLabwareOffsetCreateData, Run } from './types'
+import type { LabwareOffset, LegacyLabwareOffsetCreateData } from './types'
 
 export function createLabwareOffset(
   config: HostConfig,
   runId: string,
   data: LegacyLabwareOffsetCreateData
-): ResponsePromise<Run> {
-  return request<Run, { data: LegacyLabwareOffsetCreateData }>(
-    POST,
-    `/runs/${runId}/labware_offsets`,
-    { data },
-    config
-  )
+): ResponsePromise<LabwareOffset>
+export function createLabwareOffset(
+  config: HostConfig,
+  runId: string,
+  data: LegacyLabwareOffsetCreateData[]
+): ResponsePromise<LabwareOffset[]>
+export function createLabwareOffset(
+  config: HostConfig,
+  runId: string,
+  data: LegacyLabwareOffsetCreateData | LegacyLabwareOffsetCreateData[]
+): ResponsePromise<LabwareOffset | LabwareOffset[]> {
+  return request<
+    LabwareOffset | LabwareOffset[],
+    { data: LegacyLabwareOffsetCreateData | LegacyLabwareOffsetCreateData[] }
+  >(POST, `/runs/${runId}/labware_offsets`, { data }, config)
 }
diff --git a/react-api-client/src/runs/useCreateLabwareOffsetMutation.ts b/react-api-client/src/runs/useCreateLabwareOffsetMutation.ts
@@ -3,8 +3,8 @@ import { createLabwareOffset } from '@opentrons/api-client'
 import { useHost } from '../api'
 import type {
   HostConfig,
-  Run,
   LegacyLabwareOffsetCreateData,
+  LabwareOffset,
 } from '@opentrons/api-client'
 import type { UseMutationResult, UseMutateAsyncFunction } from 'react-query'
 
@@ -14,12 +14,12 @@ interface CreateLabwareOffsetParams {
 }
 
 export type UseCreateLabwareOffsetMutationResult = UseMutationResult<
-  Run,
+  LabwareOffset,
   unknown,
   CreateLabwareOffsetParams
 > & {
   createLabwareOffset: UseMutateAsyncFunction<
-    Run,
+    LabwareOffset,
     unknown,
     CreateLabwareOffsetParams
   >
@@ -29,19 +29,22 @@ export function useCreateLabwareOffsetMutation(): UseCreateLabwareOffsetMutation
   const host = useHost()
   const queryClient = useQueryClient()
 
-  const mutation = useMutation<Run, unknown, CreateLabwareOffsetParams>(
-    ({ runId, data }) =>
-      createLabwareOffset(host as HostConfig, runId, data)
-        .then(response => {
-          queryClient.invalidateQueries([host, 'runs']).catch((e: Error) => {
-            console.error(`error invalidating runs query: ${e.message}`)
-          })
-          return response.data
-        })
-        .catch((e: Error) => {
-          console.error(`error creating labware offsets: ${e.message}`)
-          throw e
+  const mutation = useMutation<
+    LabwareOffset,
+    unknown,
+    CreateLabwareOffsetParams
+  >(({ runId, data }) =>
+    createLabwareOffset(host as HostConfig, runId, data)
+      .then(response => {
+        queryClient.invalidateQueries([host, 'runs']).catch((e: Error) => {
+          console.error(`error invalidating runs query: ${e.message}`)
         })
+        return response.data
+      })
+      .catch((e: Error) => {
+        console.error(`error creating labware offsets: ${e.message}`)
+        throw e
+      })
   )
 
   return {

diff --git a/robot-server/robot_server/labware_offsets/router.py b/robot-server/robot_server/labware_offsets/router.py
@@ -12,7 +12,10 @@
 from opentrons.protocol_engine import ModuleModel
 
 from robot_server.labware_offsets.models import LabwareOffsetNotFound
-from robot_server.service.dependencies import get_current_time, get_unique_id
+from robot_server.service.dependencies import (
+    UniqueIDFactory,
+    get_current_time,
+)
 from robot_server.service.json_api.request import RequestModel
 from robot_server.service.json_api.response import (
     MultiBodyMeta,
@@ -42,42 +45,67 @@
 @PydanticResponse.wrap_route(
     router.post,
     path="/labwareOffsets",
-    summary="Store a labware offset",
+    summary="Store labware offsets",
     description=textwrap.dedent(
         """\
-        Store a labware offset for later retrieval through `GET /labwareOffsets`.
+        Store labware offsets for later retrieval through `GET /labwareOffsets`.
 
         On its own, this does not affect robot motion.
-        To do that, you must add the offset to a run, through the `/runs` endpoints.
+        To do that, you must add the offsets to a run, through the `/runs` endpoints.
+
+        The response body's `data` will either be a single offset or a list of offsets,
+        depending on whether you provided a single offset or a list in the request body's `data`.
         """
     ),
     status_code=201,
     include_in_schema=False,  # todo(mm, 2025-01-08): Include for v8.4.0.
 )
-async def post_labware_offset(  # noqa: D103
+async def post_labware_offsets(  # noqa: D103
     store: Annotated[LabwareOffsetStore, fastapi.Depends(get_labware_offset_store)],
-    new_offset_id: Annotated[str, fastapi.Depends(get_unique_id)],
+    new_offset_id_factory: Annotated[UniqueIDFactory, fastapi.Depends(UniqueIDFactory)],
     new_offset_created_at: Annotated[datetime, fastapi.Depends(get_current_time)],
-    request_body: Annotated[RequestModel[StoredLabwareOffsetCreate], fastapi.Body()],
-) -> PydanticResponse[SimpleBody[StoredLabwareOffset]]:
-    new_offset = IncomingStoredLabwareOffset(
-        id=new_offset_id,
-        createdAt=new_offset_created_at,
-        definitionUri=request_body.data.definitionUri,
-        locationSequence=request_body.data.locationSequence,
-        vector=request_body.data.vector,
+    request_body: Annotated[
+        RequestModel[StoredLabwareOffsetCreate | list[StoredLabwareOffsetCreate]],
+        fastapi.Body(),
+    ],
+) -> PydanticResponse[SimpleBody[StoredLabwareOffset | list[StoredLabwareOffset]]]:
+    new_offsets = [
+        IncomingStoredLabwareOffset(
+            id=new_offset_id_factory.get(),
+            createdAt=new_offset_created_at,
+            definitionUri=request_body_element.definitionUri,
+            locationSequence=request_body_element.locationSequence,
+            vector=request_body_element.vector,
+        )
+        for request_body_element in (
+            request_body.data
+            if isinstance(request_body.data, list)
+            else [request_body.data]
+        )
+    ]
+
+    for new_offset in new_offsets:
+        store.add(new_offset)
+
+    stored_offsets = [
+        StoredLabwareOffset.model_construct(
+            id=incoming.id,
+            createdAt=incoming.createdAt,
+            definitionUri=incoming.definitionUri,
+            locationSequence=incoming.locationSequence,
+            vector=incoming.vector,
+        )
+        for incoming in new_offsets
+    ]
+
+    # Return a list if the client POSTed a list, or an object if the client POSTed an object.
+    # For some reason, mypy needs to be given the type annotation explicitly.
+    response_data: StoredLabwareOffset | list[StoredLabwareOffset] = (
+        stored_offsets if isinstance(request_body.data, list) else stored_offsets[0]
     )
-    store.add(new_offset)
+
     return await PydanticResponse.create(
-        content=SimpleBody.model_construct(
-            data=StoredLabwareOffset(
-                id=new_offset_id,
-                createdAt=new_offset_created_at,
-                definitionUri=request_body.data.definitionUri,
-                locationSequence=request_body.data.locationSequence,
-                vector=request_body.data.vector,
-            )
-        ),
+        content=SimpleBody.model_construct(data=response_data),
         status_code=201,
     )
 

diff --git a/robot-server/robot_server/maintenance_runs/router/labware_router.py b/robot-server/robot_server/maintenance_runs/router/labware_router.py
@@ -36,33 +36,57 @@
         "There is no matching `GET /maintenance_runs/{runId}/labware_offsets` endpoint."
         " To read the list of labware offsets currently on the run,"
         " see the run's `labwareOffsets` field."
+        "\n\n"
+        "The response body's `data` will either be a single offset or a list of offsets,"
+        " depending on whether you provided a single offset or a list in the request body's `data`."
     ),
     status_code=status.HTTP_201_CREATED,
     responses={
-        status.HTTP_201_CREATED: {"model": SimpleBody[LabwareOffset]},
+        status.HTTP_201_CREATED: {
+            "model": SimpleBody[LabwareOffset | list[LabwareOffset]]
+        },
         status.HTTP_404_NOT_FOUND: {"model": ErrorBody[RunNotFound]},
         status.HTTP_409_CONFLICT: {"model": ErrorBody[RunNotIdle]},
     },
 )
 async def add_labware_offset(
-    request_body: RequestModel[LabwareOffsetCreate | LegacyLabwareOffsetCreate],
+    request_body: RequestModel[
+        LabwareOffsetCreate
+        | LegacyLabwareOffsetCreate
+        | list[LabwareOffsetCreate | LegacyLabwareOffsetCreate]
+    ],
     run_orchestrator_store: Annotated[
         MaintenanceRunOrchestratorStore, Depends(get_maintenance_run_orchestrator_store)
     ],
     run: Annotated[MaintenanceRun, Depends(get_run_data_from_url)],
-) -> PydanticResponse[SimpleBody[LabwareOffset]]:
-    """Add a labware offset to a maintenance run.
+) -> PydanticResponse[SimpleBody[LabwareOffset | list[LabwareOffset]]]:
+    """Add labware offsets to a maintenance run.
 
     Args:
         request_body: New labware offset request data from request body.
         run_orchestrator_store: Engine storage interface.
         run: Run response data by ID from URL; ensures 404 if run not found.
     """
-    added_offset = run_orchestrator_store.add_labware_offset(request_body.data)
-    log.info(f'Added labware offset "{added_offset.id}"' f' to run "{run.id}".')
+    offsets_to_add = (
+        request_body.data
+        if isinstance(request_body.data, list)
+        else [request_body.data]
+    )
+
+    added_offsets: list[LabwareOffset] = []
+    for offset_to_add in offsets_to_add:
+        added_offset = run_orchestrator_store.add_labware_offset(offset_to_add)
+        added_offsets.append(added_offset)
+        log.info(f'Added labware offset "{added_offset.id}" to run "{run.id}".')
+
+    # Return a list if the client POSTed a list, or an object if the client POSTed an object.
+    # For some reason, mypy needs to be given the type annotation explicitly.
+    response_data: LabwareOffset | list[LabwareOffset] = (
+        added_offsets if isinstance(request_body.data, list) else added_offsets[0]
+    )
 
     return await PydanticResponse.create(
-        content=SimpleBody.model_construct(data=added_offset),
+        content=SimpleBody.model_construct(data=response_data),
         status_code=status.HTTP_201_CREATED,
     )
 

diff --git a/robot-server/robot_server/runs/router/labware_router.py b/robot-server/robot_server/runs/router/labware_router.py
@@ -35,29 +35,38 @@
 @PydanticResponse.wrap_route(
     labware_router.post,
     path="/runs/{runId}/labware_offsets",
-    summary="Add a labware offset to a run",
+    summary="Add labware offsets to a run",
     description=(
-        "Add a labware offset to an existing run, returning the created offset."
+        "Add labware offsets to an existing run, returning the created offsets."
         "\n\n"
         "There is no matching `GET /runs/{runId}/labware_offsets` endpoint."
         " To read the list of labware offsets currently on the run,"
         " see the run's `labwareOffsets` field."
+        "\n\n"
+        "The response body's `data` will either be a single offset or a list of offsets,"
+        " depending on whether you provided a single offset or a list in the request body's `data`."
     ),
     status_code=status.HTTP_201_CREATED,
     responses={
-        status.HTTP_201_CREATED: {"model": SimpleBody[LabwareOffset]},
+        status.HTTP_201_CREATED: {
+            "model": SimpleBody[LabwareOffset | list[LabwareOffset]]
+        },
         status.HTTP_404_NOT_FOUND: {"model": ErrorBody[RunNotFound]},
         status.HTTP_409_CONFLICT: {"model": ErrorBody[Union[RunStopped, RunNotIdle]]},
     },
 )
 async def add_labware_offset(
-    request_body: RequestModel[LegacyLabwareOffsetCreate | LabwareOffsetCreate],
+    request_body: RequestModel[
+        LegacyLabwareOffsetCreate
+        | LabwareOffsetCreate
+        | list[LegacyLabwareOffsetCreate | LabwareOffsetCreate]
+    ],
     run_orchestrator_store: Annotated[
         RunOrchestratorStore, Depends(get_run_orchestrator_store)
     ],
     run: Annotated[Run, Depends(get_run_data_from_url)],
-) -> PydanticResponse[SimpleBody[LabwareOffset]]:
-    """Add a labware offset to a run.
+) -> PydanticResponse[SimpleBody[LabwareOffset | list[LabwareOffset]]]:
+    """Add labware offsets to a run.
 
     Args:
         request_body: New labware offset request data from request body.
@@ -69,11 +78,26 @@ async def add_labware_offset(
             status.HTTP_409_CONFLICT
         )
 
-    added_offset = run_orchestrator_store.add_labware_offset(request_body.data)
-    log.info(f'Added labware offset "{added_offset.id}"' f' to run "{run.id}".')
+    offsets_to_add = (
+        request_body.data
+        if isinstance(request_body.data, list)
+        else [request_body.data]
+    )
+
+    added_offsets: list[LabwareOffset] = []
+    for offset_to_add in offsets_to_add:
+        added_offset = run_orchestrator_store.add_labware_offset(offset_to_add)
+        added_offsets.append(added_offset)
+        log.info(f'Added labware offset "{added_offset.id}" to run "{run.id}".')
+
+    # Return a list if the client POSTed a list, or an object if the client POSTed an object.
+    # For some reason, mypy needs to be given the type annotation explicitly.
+    response_data: LabwareOffset | list[LabwareOffset] = (
+        added_offsets if isinstance(request_body.data, list) else added_offsets[0]
+    )
 
     return await PydanticResponse.create(
-        content=SimpleBody.model_construct(data=added_offset),
+        content=SimpleBody.model_construct(data=response_data),
         status_code=status.HTTP_201_CREATED,
     )
 

diff --git a/robot-server/robot_server/service/dependencies.py b/robot-server/robot_server/service/dependencies.py
@@ -35,7 +35,27 @@ async def get_session_manager(
 
 async def get_unique_id() -> str:
     """Get a unique ID string to use as a resource identifier."""
-    return str(uuid4())
+    return UniqueIDFactory().get()
+
+
+class UniqueIDFactory:
+    """
+    This is equivalent to the `get_unique_id()` free function. Wrapping it in a factory
+    class makes things easier for FastAPI endpoint functions that need multiple unique
+    IDs. They can do:
+
+        unique_id_factory: UniqueIDFactory = fastapi.Depends(UniqueIDFactory)
+
+    And then:
+
+        unique_id_1 = await unique_id_factory.get()
+        unique_id_2 = await unique_id_factory.get()
+    """
+
+    @staticmethod
+    def get() -> str:
+        """Get a unique ID to use as a resource identifier."""
+        return str(uuid4())
 
 
 async def get_current_time() -> datetime: