Add merge_datasets to DataClient (#1008)

Author: dhritinaidu

src/viam/app/data_client.py

@@ -91,6 +91,8 @@
     ListDatasetsByOrganizationIDRequest,
     ListDatasetsByOrganizationIDResponse,
     RenameDatasetRequest,
+    MergeDatasetsRequest,
+    MergeDatasetsResponse,
 )
 from viam.proto.app.datasync import (
     DataCaptureUploadMetadata,
@@ -1289,6 +1291,33 @@ async def create_dataset(self, name: str, organization_id: str) -> str:
         response: CreateDatasetResponse = await self._dataset_client.CreateDataset(request, metadata=self._metadata)
         return response.id
 
+
+    async def merge_datasets(self, name: str, organization_id: str, dataset_ids: List[str]) -> str:
+        """Merge multiple datasets into a new dataset.
+
+        ::
+
+            dataset_id = await data_client.merge_datasets(
+                name="<DATASET-NAME>",
+                organization_id="<YOUR-ORG-ID>",
+                dataset_ids=["<YOUR-DATASET-ID-1>", "<YOUR-DATASET-ID-2>"]
+            )
+            print(dataset_id)
+
+        Args:
+            name (str): The name of the dataset being created.
+            organization_id (str): The ID of the organization where the dataset is being created.
+                To find your organization ID, visit the organization settings page.
+            dataset_ids (List[str]): The IDs of the datasets that you would like to merge.
+        Returns:
+            str: The dataset ID of the created dataset.
+
+        For more information, see `Data Client API <https://docs.viam.com/dev/reference/apis/data-client/#mergedatasets>`_.
+        """
+        request = MergeDatasetsRequest(name=name, organization_id=organization_id, dataset_ids=dataset_ids)
+        response: MergeDatasetsResponse = await self._dataset_client.MergeDatasets(request, metadata=self._metadata)
+        return response.dataset_id
+
     async def list_dataset_by_ids(self, ids: List[str]) -> Sequence[Dataset]:
         """Get a list of datasets using their IDs.
 

tests/mocks/services.py

@@ -316,7 +316,6 @@
     PointCloudObject,
     Pose,
     PoseInFrame,
-    ResourceName,
     Transform,
 )
 from viam.proto.provisioning import (
@@ -616,11 +615,11 @@ def __init__(
     async def Move(self, stream: Stream[MoveRequest, MoveResponse]) -> None:
         request = await stream.recv_message()
         assert request is not None
-        name: ResourceName = request.component_name
+        name: str = request.component_name
         self.constraints = request.constraints
         self.extra = struct_to_dict(request.extra)
         self.timeout = stream.deadline.time_remaining() if stream.deadline else None
-        success = self.move_responses[name.name]
+        success = self.move_responses[name]
         response = MoveResponse(success=success)
         await stream.send_message(response)
 
@@ -655,10 +654,10 @@ async def MoveOnGlobe(self, stream: Stream[MoveOnGlobeRequest, MoveOnGlobeRespon
     async def GetPose(self, stream: Stream[GetPoseRequest, GetPoseResponse]) -> None:
         request = await stream.recv_message()
         assert request is not None
-        name: ResourceName = request.component_name
+        name: str = request.component_name
         self.extra = struct_to_dict(request.extra)
         self.timeout = stream.deadline.time_remaining() if stream.deadline else None
-        pose = self.get_pose_responses[name.name]
+        pose = self.get_pose_responses[name]
         response = GetPoseResponse(pose=pose)
         await stream.send_message(response)
 
@@ -1071,9 +1070,10 @@ async def ExportTabularData(self, stream: Stream[ExportTabularDataRequest, Expor
 
 
 class MockDataset(DatasetServiceBase):
-    def __init__(self, create_response: str, datasets_response: Sequence[Dataset]):
+    def __init__(self, create_response: str, datasets_response: Sequence[Dataset], merged_response: Optional[str] = None):
         self.create_response = create_response
         self.datasets_response = datasets_response
+        self.merged_response = merged_response
 
     async def CreateDataset(self, stream: Stream[CreateDatasetRequest, CreateDatasetResponse]) -> None:
         request = await stream.recv_message()
@@ -1105,7 +1105,11 @@ async def ListDatasetsByOrganizationID(
     async def MergeDatasets(self, stream: Stream[MergeDatasetsRequest, MergeDatasetsResponse]) -> None:
         request = await stream.recv_message()
         assert request is not None
-        await stream.send_message(MergeDatasetsResponse())
+        self.name = request.name
+        self.org_id = request.organization_id
+        self.dataset_ids = request.dataset_ids
+        self.merged_response = "".join(self.dataset_ids)
+        await stream.send_message(MergeDatasetsResponse(dataset_id=self.merged_response))
 
     async def RenameDataset(self, stream: Stream[RenameDatasetRequest, RenameDatasetResponse]) -> None:
         request = await stream.recv_message()

tests/test_dataset.py

@@ -8,7 +8,10 @@
 from .mocks.services import MockDataset
 
 CREATED_ID = "VIAM_DATASET_0"
+MERGED_NAME = "VIAM_DATASET_MERGED"
 ID = "VIAM_DATASET_1"
+ID2 = "VIAM_DATASET_2"
+MERGED_ID = f'{ID}{ID2}'
 NAME = "dataset"
 ORG_ID = "org_id"
 SECONDS = 978310861
@@ -22,7 +25,7 @@
 
 @pytest.fixture(scope="function")
 def service() -> MockDataset:
-    return MockDataset(CREATED_ID, DATASETS)
+    return MockDataset(CREATED_ID, DATASETS, MERGED_ID)
 
 
 class TestClient:
@@ -34,6 +37,14 @@ async def test_create_dataset(self, service: MockDataset):
             assert service.org_id == ORG_ID
             assert id == CREATED_ID
 
+    async def test_merge_datasets(self, service: MockDataset):
+        async with ChannelFor([service]) as channel:
+            client = DataClient(channel, DATA_SERVICE_METADATA)
+            id = await client.merge_datasets(MERGED_NAME, ORG_ID, [ID, ID2])
+            assert service.name == MERGED_NAME
+            assert service.org_id == ORG_ID
+            assert id == MERGED_ID
+
     async def test_delete_dataset(self, service: MockDataset):
         async with ChannelFor([service]) as channel:
             client = DataClient(channel, DATA_SERVICE_METADATA)