DATA 4228: Update SDKs and CLI (#934)

Author: Emerald-Z

src/viam/app/data_client.py

@@ -270,6 +270,9 @@ class DataPipeline:
         enabled: bool
         """Whether the data pipeline is enabled"""
 
+        data_source_type: TabularDataSourceType.ValueType
+        """The type of data source for the data pipeline"""
+
         @classmethod
         def from_proto(cls, data_pipeline: ProtoDataPipeline) -> Self:
             return cls(
@@ -281,6 +284,7 @@ def from_proto(cls, data_pipeline: ProtoDataPipeline) -> Self:
                 created_on=data_pipeline.created_on.ToDatetime(),
                 updated_at=data_pipeline.updated_at.ToDatetime(),
                 enabled=data_pipeline.enabled,
+                data_source_type=data_pipeline.data_source_type,
             )
 
     @dataclass
@@ -1883,7 +1887,14 @@ async def list_data_pipelines(self, organization_id: str) -> List[DataPipeline]:
         response: ListDataPipelinesResponse = await self._data_pipelines_client.ListDataPipelines(request, metadata=self._metadata)
         return [DataClient.DataPipeline.from_proto(pipeline) for pipeline in response.data_pipelines]
 
-    async def create_data_pipeline(self, organization_id: str, name: str, mql_binary: List[Dict[str, Any]], schedule: str) -> str:
+    async def create_data_pipeline(
+        self,
+        organization_id: str,
+        name: str,
+        mql_binary: List[Dict[str, Any]],
+        schedule: str,
+        data_source_type: TabularDataSourceType.ValueType = TabularDataSourceType.TABULAR_DATA_SOURCE_TYPE_STANDARD,
+    ) -> str:
         """Create a new data pipeline.
 
         ::
@@ -1892,7 +1903,8 @@ async def create_data_pipeline(self, organization_id: str, name: str, mql_binary
                 organization_id="<YOUR-ORGANIZATION-ID>",
                 name="<YOUR-PIPELINE-NAME>",
                 mql_binary=[<YOUR-MQL-PIPELINE-AGGREGATION>],
-                schedule="<YOUR-SCHEDULE>"
+                schedule="<YOUR-SCHEDULE>",
+                data_source_type=TabularDataSourceType.TABULAR_DATA_SOURCE_TYPE_STANDARD,
             )
 
         Args:
@@ -1902,12 +1914,14 @@ async def create_data_pipeline(self, organization_id: str, name: str, mql_binary
             mql_binary (List[Dict[str, Any]]):The MQL pipeline to run, as a list of MongoDB aggregation pipeline stages.
             schedule (str): A cron expression representing the expected execution schedule in UTC (note this also
                 defines the input time window; an hourly schedule would process 1 hour of data at a time).
+            data_source_type (TabularDataSourceType): The type of data source to use for the pipeline.
+                Defaults to TabularDataSourceType.TABULAR_DATA_SOURCE_TYPE_STANDARD.
 
         Returns:
             str: The ID of the newly created pipeline.
         """
         binary: List[bytes] = [bson.encode(query) for query in mql_binary]
-        request = CreateDataPipelineRequest(organization_id=organization_id, name=name, mql_binary=binary, schedule=schedule)
+        request = CreateDataPipelineRequest(organization_id=organization_id, name=name, mql_binary=binary, schedule=schedule, data_source_type=data_source_type)
         response: CreateDataPipelineResponse = await self._data_pipelines_client.CreateDataPipeline(request, metadata=self._metadata)
         return response.id
 

tests/mocks/services.py

@@ -1147,6 +1147,7 @@ async def CreateDataPipeline(self, stream: Stream[CreateDataPipelineRequest, Cre
         self.mql_binary = request.mql_binary
         self.schedule = request.schedule
         self.org_id = request.organization_id
+        self.data_source_type = request.data_source_type
         await stream.send_message(CreateDataPipelineResponse(id=self.create_response))
 
     async def GetDataPipeline(self, stream: Stream[GetDataPipelineRequest, GetDataPipelineResponse]) -> None:

tests/test_data_pipelines.py

@@ -6,6 +6,7 @@
 from viam.app.data_client import DataClient
 from viam.proto.app.datapipelines import DataPipeline, DataPipelineRun, DataPipelineRunStatus
 from viam.utils import datetime_to_timestamp
+from viam.proto.app.data import TabularDataSourceType
 
 from .mocks.services import MockDataPipelines
 
@@ -15,6 +16,8 @@
 SCHEDULE = "0 0 * * *"
 UPDATED_SCHEDULE = "0 1 * * *"
 MQL_BINARY = []
+DATA_SOURCE_TYPE = TabularDataSourceType.TABULAR_DATA_SOURCE_TYPE_UNSPECIFIED
+STANDARD_DATA_SOURCE_TYPE = TabularDataSourceType.TABULAR_DATA_SOURCE_TYPE_STANDARD
 
 TIMESTAMP = datetime.fromtimestamp(0)
 TIMESTAMP_PROTO = datetime_to_timestamp(TIMESTAMP)
@@ -28,6 +31,7 @@
     enabled=True,
     created_on=TIMESTAMP_PROTO,
     updated_at=TIMESTAMP_PROTO,
+    data_source_type=DATA_SOURCE_TYPE,
 )
 PROTO_DATA_PIPELINES = [PROTO_DATA_PIPELINE]
 
@@ -84,6 +88,7 @@ async def test_create_data_pipeline(self, service: MockDataPipelines):
             assert service.org_id == ORG_ID
             assert service.schedule == SCHEDULE
             assert service.mql_binary == MQL_BINARY
+            assert service.data_source_type == STANDARD_DATA_SOURCE_TYPE
 
     async def test_get_data_pipeline(self, service: MockDataPipelines):
         async with ChannelFor([service]) as channel: