RSDK-9503 accept bson queries in MQL function (#803)

Author: purplenicole730

src/viam/app/data_client.py

@@ -84,7 +84,7 @@
     StreamingDataCaptureUploadResponse,
     UploadMetadata,
 )
-from viam.utils import ValueTypes, create_filter, datetime_to_timestamp, struct_to_dict
+from viam.utils import ValueTypes, _alias_param, create_filter, datetime_to_timestamp, struct_to_dict
 
 LOGGER = logging.getLogger(__name__)
 
@@ -334,33 +334,36 @@ async def tabular_data_by_sql(self, organization_id: str, sql_query: str) -> Lis
         response: TabularDataBySQLResponse = await self._data_client.TabularDataBySQL(request, metadata=self._metadata)
         return [bson.decode(bson_bytes) for bson_bytes in response.raw_data]
 
-    async def tabular_data_by_mql(self, organization_id: str, mql_binary: List[bytes]) -> List[Dict[str, Union[ValueTypes, datetime]]]:
+    @_alias_param("query", param_alias="mql_binary")
+    async def tabular_data_by_mql(
+        self, organization_id: str, query: Union[List[bytes], List[Dict[str, Any]]]
+    ) -> List[Dict[str, Union[ValueTypes, datetime]]]:
         """Obtain unified tabular data and metadata, queried with MQL.
 
         ::
 
             import bson
 
-            # using pymongo package (pip install pymongo)
-            tabular_data = await data_client.tabular_data_by_mql(organization_id="<YOUR-ORG-ID>", mql_binary=[
-                bson.encode({ '$match': { 'location_id': '<YOUR-LOCATION-ID>' } }),
-                bson.encode({ "$limit": 5 })
+            tabular_data = await data_client.tabular_data_by_mql(organization_id="<YOUR-ORG-ID>", mql_query=[
+                { '$match': { 'location_id': '<YOUR-LOCATION-ID>' } },
+                { "$limit": 5 }
             ])
 
             print(f"Tabular Data: {tabular_data}")
 
         Args:
             organization_id (str): The ID of the organization that owns the data.
                 You can obtain your organization ID from the Viam app's organization settings page.
-            mql_binary (List[bytes]): The MQL query to run as a list of BSON queries. You can encode your bson queries using a library like
-                `pymongo`.
+            query (Union[List[bytes], List[Dict[str, Any]]]): The MQL query to run as a list of BSON queries.
+                Note: Support for bytes will be removed in the future, so using a dictionary is preferred.
 
         Returns:
             List[Dict[str, Union[ValueTypes, datetime]]]: An array of decoded BSON data objects.
 
         For more information, see `Data Client API <https://docs.viam.com/appendix/apis/data-client/>`_.
         """
-        request = TabularDataByMQLRequest(organization_id=organization_id, mql_binary=mql_binary)
+        binary: List[bytes] = [bson.encode(query) for query in query] if isinstance(query[0], dict) else query  # type: ignore
+        request = TabularDataByMQLRequest(organization_id=organization_id, mql_binary=binary)
         response: TabularDataByMQLResponse = await self._data_client.TabularDataByMQL(request, metadata=self._metadata)
         return [bson.decode(bson_bytes) for bson_bytes in response.raw_data]
 

src/viam/utils.py

@@ -4,7 +4,7 @@
 import sys
 import threading
 from datetime import datetime
-from typing import Any, Dict, List, Mapping, Optional, SupportsBytes, SupportsFloat, Type, TypeVar, Union
+from typing import Any, Callable, Dict, List, Mapping, Optional, SupportsBytes, SupportsFloat, Type, TypeVar, Union
 
 from google.protobuf.json_format import MessageToDict, ParseDict
 from google.protobuf.message import Message
@@ -339,3 +339,28 @@ def create_filter(
         bbox_labels=bbox_labels,
         dataset_id=dataset_id if dataset_id else "",
     )
+
+
+def _alias_param(param_name: str, param_alias: str) -> Callable:
+    """
+    Decorator for aliasing a param in a function. Intended for providing backwards compatibility on params with name changes.
+
+    Args:
+        param_name: name of param in function to alias
+        param_alias: alias that can be used for this param
+    Returns:
+        The input function, plus param alias.
+    """
+    def decorator(func: Callable):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            alias_param_value = kwargs.get(param_alias)
+            if alias_param_value:
+                # Only use alias value if param is not given.
+                if not kwargs.get(param_name):
+                    kwargs[param_name] = alias_param_value
+                del kwargs[param_alias]
+            result = func(*args, **kwargs)
+            return result
+        return wrapper
+    return decorator

tests/test_data_client.py

@@ -75,7 +75,7 @@
 )
 BBOXES = [BBOX]
 SQL_QUERY = "sql_query"
-MQL_BINARY = [b"mql_binary"]
+MQL_BINARY = [{"binary": "mql_binary"}]
 TABULAR_DATA = {"key": "value"}
 TABULAR_METADATA = CaptureMetadata(
     organization_id=ORG_ID,
@@ -180,6 +180,9 @@ async def test_tabular_data_by_mql(self, service: MockData):
             response = await client.tabular_data_by_mql(ORG_ID, MQL_BINARY)
             assert isinstance(response[0]["key1"], datetime)
             assert response == TABULAR_QUERY_RESPONSE
+            response = await client.tabular_data_by_mql(ORG_ID, mql_binary=[b"mql_binary"])
+            assert isinstance(response[0]["key1"], datetime)
+            assert response == TABULAR_QUERY_RESPONSE
 
     async def test_get_latest_tabular_data(self, service: MockData):
         async with ChannelFor([service]) as channel: