[DATA-2034] Add file extensions to binary data capture upload (#466)

Author: njooma

src/viam/app/data_client.py

@@ -437,6 +437,7 @@ async def binary_data_capture_upload(
         method_parameters: Optional[Mapping[str, Any]] = None,
         tags: Optional[List[str]] = None,
         data_request_times: Optional[Tuple[datetime, datetime]] = None,
+        file_extension: Optional[str] = None,
     ) -> str:
         """Upload binary sensor data.
 
@@ -453,6 +454,8 @@ async def binary_data_capture_upload(
             tags (Optional[List[str]]): Optional list of tags to allow for tag-based data filtering when retrieving data.
             data_request_times (Optional[Tuple[datetime.datetime, datetime.datetime]]): Optional tuple containing `datetime`s objects
                 denoting the times this data was requested[0] by the robot and received[1] from the appropriate sensor.
+            file_extension (str): The file extension of binary data including the period, e.g. .jpg, .png, .pcd.
+                The backend will route the binary to its corresponding mime type based on this extension.
 
         Raises:
             GRPCError: If an invalid part ID is passed.
@@ -478,6 +481,8 @@ async def binary_data_capture_upload(
             method_parameters=method_parameters,
             tags=tags,
         )
+        if file_extension:
+            metadata.file_extension = file_extension if file_extension[0] == "." else f".{file_extension}"
         response = await self._data_capture_upload(metadata=metadata, sensor_contents=[sensor_contents])
         return response.file_id
 

tests/test_data_sync_client.py

@@ -1,6 +1,7 @@
 import pytest
 from datetime import datetime
-from typing import List, cast, Mapping, Any
+from typing import List, Tuple, cast, Mapping, Any
+from google.protobuf.timestamp_pb2 import Timestamp
 
 from grpclib.testing import ChannelFor
 
@@ -21,7 +22,7 @@
 BINARY_DATA = b"binary_data"
 METHOD_NAME = "method_name"
 DATETIMES = (datetime.now(), datetime.now())
-TIMESTAMPS = (datetime_to_timestamp(DATETIMES[0]), datetime_to_timestamp(DATETIMES[1]))
+TIMESTAMPS = cast(Tuple[Timestamp, Timestamp], (datetime_to_timestamp(DATETIMES[0]), datetime_to_timestamp(DATETIMES[1])))
 METHOD_PARAMETERS = {}
 TABULAR_DATA = [{"key": "value"}]
 FILE_NAME = "file_name"
@@ -51,11 +52,27 @@ async def test_binary_data_capture_upload(self, service: MockDataSync):
                 tags=TAGS,
                 data_request_times=DATETIMES,
                 binary_data=BINARY_DATA,
+                file_extension=".txt",
             )
             self.assert_sensor_contents(sensor_contents=list(service.sensor_contents), is_binary=True)
             self.assert_metadata(metadata=service.metadata)
+            assert service.metadata.file_extension == ".txt"
             assert file_id == FILE_UPLOAD_RESPONSE
 
+            # Test extension dot prepend
+            file_id = await client.binary_data_capture_upload(
+                part_id=PART_ID,
+                component_type=COMPONENT_TYPE,
+                component_name=COMPONENT_NAME,
+                method_name=METHOD_NAME,
+                method_parameters=METHOD_PARAMETERS,
+                tags=TAGS,
+                data_request_times=DATETIMES,
+                binary_data=BINARY_DATA,
+                file_extension="txt",
+            )
+            assert service.metadata.file_extension == ".txt"
+
     @pytest.mark.asyncio
     async def test_tabular_data_capture_upload(self, service: MockDataSync):
         async with ChannelFor([service]) as channel: