RSDK-5731: add dataset apis (#538)

Author: purplenicole730

src/viam/app/data_client.py

@@ -1,13 +1,14 @@
 import warnings
 from datetime import datetime
 from pathlib import Path
-from typing import Any, List, Mapping, Optional, Tuple
+from typing import Any, List, Mapping, Optional, Sequence, Tuple
 
 from google.protobuf.struct_pb2 import Struct
 from grpclib.client import Channel, Stream
 
 from viam import logging
 from viam.proto.app.data import (
+    AddBinaryDataToDatasetByIDsRequest,
     AddBoundingBoxToImageByIDRequest,
     AddBoundingBoxToImageByIDResponse,
     AddTagsToBinaryDataByFilterRequest,
@@ -32,6 +33,7 @@
     Filter,
     GetDatabaseConnectionRequest,
     GetDatabaseConnectionResponse,
+    RemoveBinaryDataFromDatasetByIDsRequest,
     RemoveBoundingBoxFromImageByIDRequest,
     RemoveTagsFromBinaryDataByFilterRequest,
     RemoveTagsFromBinaryDataByFilterResponse,
@@ -42,6 +44,18 @@
     TagsByFilterRequest,
     TagsByFilterResponse,
 )
+from viam.proto.app.dataset import (
+    CreateDatasetRequest,
+    CreateDatasetResponse,
+    Dataset,
+    DatasetServiceStub,
+    DeleteDatasetRequest,
+    ListDatasetsByIDsRequest,
+    ListDatasetsByIDsResponse,
+    ListDatasetsByOrganizationIDRequest,
+    ListDatasetsByOrganizationIDResponse,
+    RenameDatasetRequest,
+)
 from viam.proto.app.datasync import (
     DataCaptureUploadMetadata,
     DataCaptureUploadRequest,
@@ -98,6 +112,7 @@ def __eq__(self, other: object) -> bool:
                 return str(self) == str(other)
             return False
 
+    # TODO (RSDK-6684): Revisit if this shadow type is necessary
     class BinaryData:
         """Class representing a piece of binary data and associated metadata.
 
@@ -131,10 +146,12 @@ def __init__(self, channel: Channel, metadata: Mapping[str, str]):
         self._metadata = metadata
         self._data_client = DataServiceStub(channel)
         self._data_sync_client = DataSyncServiceStub(channel)
+        self._dataset_client = DatasetServiceStub(channel)
         self._channel = channel
 
     _data_client: DataServiceStub
     _data_sync_client: DataSyncServiceStub
+    _dataset_client: DatasetServiceStub
     _metadata: Mapping[str, str]
     _channel: Channel
 
@@ -478,6 +495,93 @@ async def get_database_connection(self, organization_id: str) -> str:
     async def configure_database_user(self, organization_id: str, password: str) -> None:
         raise NotImplementedError()
 
+    async def create_dataset(self, name: str, organization_id: str) -> str:
+        """Create a new dataset.
+
+        Args:
+            name (str): The name of the dataset being created.
+            organization_id (str): The ID of the organization where the dataset is being created.
+
+        Returns:
+            str: The dataset ID of the created dataset.
+        """
+        request = CreateDatasetRequest(name=name, organization_id=organization_id)
+        response: CreateDatasetResponse = await self._dataset_client.CreateDataset(request, metadata=self._metadata)
+        return response.id
+
+    async def list_dataset_by_ids(self, ids: List[str]) -> Sequence[Dataset]:
+        """Get a list of datasets using their IDs.
+
+        Args:
+            ids (List[str]): The IDs of the datasets being called for.
+
+        Returns:
+            Sequence[Dataset]: The list of datasets.
+        """
+        request = ListDatasetsByIDsRequest(ids=ids)
+        response: ListDatasetsByIDsResponse = await self._dataset_client.ListDatasetsByIDs(request, metadata=self._metadata)
+
+        return response.datasets
+
+    async def list_datasets_by_organization_id(self, organization_id: str) -> Sequence[Dataset]:
+        """Get the datasets in an organization.
+
+        Args:
+            organization_id (str): The ID of the organization.
+
+        Returns:
+            Sequence[Dataset]: The list of datasets in the organization.
+        """
+        request = ListDatasetsByOrganizationIDRequest(organization_id=organization_id)
+        response: ListDatasetsByOrganizationIDResponse = await self._dataset_client.ListDatasetsByOrganizationID(
+            request, metadata=self._metadata
+        )
+
+        return response.datasets
+
+    async def rename_dataset(self, id: str, name: str) -> None:
+        """Rename a dataset specified by the dataset ID.
+
+        Args:
+            id (str): The ID of the dataset.
+            name (str): The new name of the dataset.
+        """
+        request = RenameDatasetRequest(id=id, name=name)
+        await self._dataset_client.RenameDataset(request, metadata=self._metadata)
+
+    async def delete_dataset(self, id: str) -> None:
+        """Delete a dataset.
+
+        Args:
+            id (str): The ID of the dataset.
+        """
+        request = DeleteDatasetRequest(id=id)
+        await self._dataset_client.DeleteDataset(request, metadata=self._metadata)
+
+    async def add_binary_data_to_dataset_by_ids(self, binary_ids: List[BinaryID], dataset_id: str) -> None:
+        """Add the BinaryData to the provided dataset.
+
+        This BinaryData will be tagged with the VIAM_DATASET_{id} label.
+
+        Args:
+            binary_ids (List[BinaryID]): The IDs of binary data to add to dataset.
+            dataset_id (str): The ID of the dataset to be added to.
+        """
+        request = AddBinaryDataToDatasetByIDsRequest(binary_ids=binary_ids, dataset_id=dataset_id)
+        await self._data_client.AddBinaryDataToDatasetByIDs(request, metadata=self._metadata)
+
+    async def remove_binary_data_from_dataset_by_ids(self, binary_ids: List[BinaryID], dataset_id: str) -> None:
+        """Remove the BinaryData from the provided dataset.
+
+        This BinaryData will lose the VIAM_DATASET_{id} tag.
+
+        Args:
+            binary_ids (List[BinaryID]): The IDs of binary data to remove from dataset.
+            dataset_id (str): The ID of the dataset to be removed from.
+        """
+        request = RemoveBinaryDataFromDatasetByIDsRequest(binary_ids=binary_ids, dataset_id=dataset_id)
+        await self._data_client.RemoveBinaryDataFromDatasetByIDs(request, metadata=self._metadata)
+
     async def binary_data_capture_upload(
         self,
         binary_data: bytes,
@@ -806,8 +910,9 @@ def create_filter(
         end_time: Optional[datetime] = None,
         tags: Optional[List[str]] = None,
         bbox_labels: Optional[List[str]] = None,
+        dataset_id: Optional[str] = None,
     ) -> Filter:
-        warnings.warn("DataClient.create_filter is deprecated. Use AppClient.create_filter instead.", DeprecationWarning, stacklevel=2)
+        warnings.warn("DataClient.create_filter is deprecated. Use utils.create_filter instead.", DeprecationWarning, stacklevel=2)
         return create_filter(
             component_name,
             component_type,
@@ -823,4 +928,5 @@ def create_filter(
             end_time,
             tags,
             bbox_labels,
+            dataset_id,
         )

src/viam/utils.py

@@ -284,6 +284,7 @@ def create_filter(
     end_time: Optional[datetime] = None,
     tags: Optional[List[str]] = None,
     bbox_labels: Optional[List[str]] = None,
+    dataset_id: Optional[str] = None,
 ) -> Filter:
     """Create a `Filter`.
 
@@ -303,6 +304,7 @@ def create_filter(
         tags (Optional[List[str]]): Optional list of tags attached to the data being filtered (e.g., ["test"]).
         bbox_labels (Optional[List[str]]): Optional list of bounding box labels attached to the data being filtered (e.g., ["square",
             "circle"]).
+        dataset_id (Optional[str]): Optional ID of dataset associated with data being filtered
 
     Returns:
         viam.proto.app.data.Filter: The `Filter` object.
@@ -328,4 +330,5 @@ def create_filter(
         else None,
         tags_filter=TagsFilter(tags=tags),
         bbox_labels=bbox_labels,
+        dataset_id=dataset_id if dataset_id else "",
     )

tests/mocks/services.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List, Mapping, Optional, Tuple, Union
+from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple, Union
 
 import numpy as np
 from grpclib.server import Stream
@@ -217,6 +217,20 @@
     TagsByFilterRequest,
     TagsByFilterResponse,
 )
+from viam.proto.app.dataset import (
+    CreateDatasetRequest,
+    CreateDatasetResponse,
+    Dataset,
+    DatasetServiceBase,
+    DeleteDatasetRequest,
+    DeleteDatasetResponse,
+    ListDatasetsByIDsRequest,
+    ListDatasetsByIDsResponse,
+    ListDatasetsByOrganizationIDRequest,
+    ListDatasetsByOrganizationIDResponse,
+    RenameDatasetRequest,
+    RenameDatasetResponse,
+)
 from viam.proto.app.datasync import (
     DataCaptureUploadRequest,
     DataCaptureUploadResponse,
@@ -822,12 +836,20 @@ async def ConfigureDatabaseUser(self, stream: Stream[ConfigureDatabaseUserReques
     async def AddBinaryDataToDatasetByIDs(
         self, stream: Stream[AddBinaryDataToDatasetByIDsRequest, AddBinaryDataToDatasetByIDsResponse]
     ) -> None:
-        raise NotImplementedError()
+        request = await stream.recv_message()
+        assert request is not None
+        self.added_data_ids = request.binary_ids
+        self.dataset_id = request.dataset_id
+        await stream.send_message(AddBinaryDataToDatasetByIDsResponse())
 
     async def RemoveBinaryDataFromDatasetByIDs(
         self, stream: Stream[RemoveBinaryDataFromDatasetByIDsRequest, RemoveBinaryDataFromDatasetByIDsResponse]
     ) -> None:
-        raise NotImplementedError()
+        request = await stream.recv_message()
+        assert request is not None
+        self.removed_data_ids = request.binary_ids
+        self.dataset_id = request.dataset_id
+        await stream.send_message(RemoveBinaryDataFromDatasetByIDsResponse())
 
     async def TabularDataBySQL(self, stream: Stream[TabularDataBySQLRequest, TabularDataBySQLResponse]) -> None:
         raise NotImplementedError()
@@ -836,6 +858,46 @@ async def TabularDataByMQL(self, stream: Stream[TabularDataByMQLRequest, Tabular
         raise NotImplementedError()
 
 
+class MockDataset(DatasetServiceBase):
+    def __init__(self, create_response: str, datasets_response: Sequence[Dataset]):
+        self.create_response = create_response
+        self.datasets_response = datasets_response
+
+    async def CreateDataset(self, stream: Stream[CreateDatasetRequest, CreateDatasetResponse]) -> None:
+        request = await stream.recv_message()
+        assert request is not None
+        self.name = request.name
+        self.org_id = request.organization_id
+        await stream.send_message(CreateDatasetResponse(id=self.create_response))
+
+    async def DeleteDataset(self, stream: Stream[DeleteDatasetRequest, DeleteDatasetResponse]) -> None:
+        request = await stream.recv_message()
+        assert request is not None
+        self.deleted_id = request.id
+        await stream.send_message(DeleteDatasetResponse())
+
+    async def ListDatasetsByIDs(self, stream: Stream[ListDatasetsByIDsRequest, ListDatasetsByIDsResponse]) -> None:
+        request = await stream.recv_message()
+        assert request is not None
+        self.ids = request.ids
+        await stream.send_message(ListDatasetsByIDsResponse(datasets=self.datasets_response))
+
+    async def ListDatasetsByOrganizationID(
+        self, stream: Stream[ListDatasetsByOrganizationIDRequest, ListDatasetsByOrganizationIDResponse]
+    ) -> None:
+        request = await stream.recv_message()
+        assert request is not None
+        self.org_id = request.organization_id
+        await stream.send_message(ListDatasetsByOrganizationIDResponse(datasets=self.datasets_response))
+
+    async def RenameDataset(self, stream: Stream[RenameDatasetRequest, RenameDatasetResponse]) -> None:
+        request = await stream.recv_message()
+        assert request is not None
+        self.id = request.id
+        self.name = request.name
+        await stream.send_message((RenameDatasetResponse()))
+
+
 class MockDataSync(DataSyncServiceBase):
     def __init__(self, file_upload_response: str):
         self.file_upload_response = file_upload_response

tests/test_data_client.py

@@ -36,6 +36,7 @@
 TAGS = ["tag"]
 BBOX_LABEL = "bbox_label"
 BBOX_LABELS = [BBOX_LABEL]
+DATASET_ID = "VIAM_DATASET_1"
 FILTER = create_filter(
     component_name=COMPONENT_NAME,
     component_type=COMPONENT_TYPE,
@@ -51,6 +52,7 @@
     end_time=END_DATETIME,
     tags=TAGS,
     bbox_labels=BBOX_LABELS,
+    dataset_id=DATASET_ID,
 )
 
 FILE_ID = "file_id"
@@ -248,6 +250,22 @@ async def test_get_database_connection(self, service: MockData):
     async def test_configure_database_user(self, service: MockData):
         assert True
 
+    @pytest.mark.asyncio
+    async def test_add_binary_data_to_dataset_by_ids(self, service: MockData):
+        async with ChannelFor([service]) as channel:
+            client = DataClient(channel, DATA_SERVICE_METADATA)
+            await client.add_binary_data_to_dataset_by_ids(binary_ids=BINARY_IDS, dataset_id=DATASET_ID)
+            assert service.added_data_ids == BINARY_IDS
+            assert service.dataset_id == DATASET_ID
+
+    @pytest.mark.asyncio
+    async def test_remove_binary_data_to_dataset_by_ids(self, service: MockData):
+        async with ChannelFor([service]) as channel:
+            client = DataClient(channel, DATA_SERVICE_METADATA)
+            await client.remove_binary_data_from_dataset_by_ids(binary_ids=BINARY_IDS, dataset_id=DATASET_ID)
+            assert service.removed_data_ids == BINARY_IDS
+            assert service.dataset_id == DATASET_ID
+
     def assert_filter(self, filter: Filter) -> None:
         assert filter.component_name == COMPONENT_NAME
         assert filter.component_type == COMPONENT_TYPE

tests/test_dataset.py

@@ -0,0 +1,59 @@
+import pytest
+from google.protobuf.timestamp_pb2 import Timestamp
+from grpclib.testing import ChannelFor
+
+from viam.app.data_client import DataClient
+from viam.proto.app.dataset import Dataset
+
+from .mocks.services import MockDataset
+
+CREATED_ID = "VIAM_DATASET_0"
+ID = "VIAM_DATASET_1"
+NAME = "dataset"
+ORG_ID = "org_id"
+SECONDS = 978310861
+NANOS = 0
+TIME = Timestamp(seconds=SECONDS, nanos=NANOS)
+DATASET = Dataset(id=ID, name=NAME, organization_id=ORG_ID, time_created=TIME)
+DATASETS = [DATASET]
+AUTH_TOKEN = "auth_token"
+DATA_SERVICE_METADATA = {"authorization": f"Bearer {AUTH_TOKEN}"}
+
+
+@pytest.fixture(scope="function")
+def service() -> MockDataset:
+    return MockDataset(CREATED_ID, DATASETS)
+
+
+class TestClient:
+    @pytest.mark.asyncio
+    async def test_create_dataset(self, service: MockDataset):
+        async with ChannelFor([service]) as channel:
+            client = DataClient(channel, DATA_SERVICE_METADATA)
+            id = await client.create_dataset(NAME, ORG_ID)
+            assert service.name == NAME
+            assert service.org_id == ORG_ID
+            assert id == CREATED_ID
+
+    @pytest.mark.asyncio
+    async def test_delete_dataset(self, service: MockDataset):
+        async with ChannelFor([service]) as channel:
+            client = DataClient(channel, DATA_SERVICE_METADATA)
+            await client.delete_dataset(ID)
+            assert service.deleted_id == ID
+
+    @pytest.mark.asyncio
+    async def test_list_datasets_by_ids(self, service: MockDataset):
+        async with ChannelFor([service]) as channel:
+            client = DataClient(channel, DATA_SERVICE_METADATA)
+            datasets = await client.list_dataset_by_ids([ID])
+            assert service.ids == [ID]
+            assert datasets == DATASETS
+
+    @pytest.mark.asyncio
+    async def test_list_datasets_by_organization_id(self, service: MockDataset):
+        async with ChannelFor([service]) as channel:
+            client = DataClient(channel, DATA_SERVICE_METADATA)
+            datasets = await client.list_datasets_by_organization_id(ORG_ID)
+            assert service.org_id == ORG_ID
+            assert datasets == DATASETS