Merge d59a34a into 7ae7e7a

intake_dal/dal_online.py

@@ -75,7 +75,12 @@ def get_metadata(key: str, default):
 
     def _get_partition(self, _) -> pd.DataFrame:
         self._get_schema()
-        data = _http_get_avro_data_set(self._url, self._canonical_name, self._key_value)
+        if isinstance(self._key_value, list):
+            massaged_key = ",".join(map(str, self._key_value))
+        else:
+            massaged_key = self._key_value
+
+        data = _http_get_avro_data_set(self._url, self._canonical_name, massaged_key)
         for row in data:
             for key, field in row.items():
                 if isinstance(field, dict) and "format" in field:

intake_dal/tests/test_dal_online.py

@@ -107,3 +107,23 @@ def test_post_in_chunks(mock_put: MagicMock, serving_cat: DalCatalog, user_event
 
     assert len(mock_put.call_args_list) == 2
     assert len(ret) == 2
+
+
+@mock.patch("intake_dal.dal_online._http_get_avro_data_set")
+def test_dal_online_multi_key_read(
+    mock_get: MagicMock,
+    serving_cat: DalCatalog,
+    user_events_df: pd.DataFrame,
+    user_events_json: List[Dict],
+):
+    canonical_name = "entity.user.user_events"
+    mock_get.return_value = user_events_json
+
+    assert_frame_equal(user_events_df, serving_cat.entity.user.user_events(key=[100, 101]).read(), check_dtype=False)
+    mock_get.assert_called()
+
+    assert serialize_panda_df_to_str(
+        user_events_df, schema=json.loads(serving_cat.metadata["data_schema"][canonical_name])
+    ) != serialize_panda_df_to_str(
+        user_events_df, schema=json.loads(serving_cat.metadata["data_schema"][canonical_name])
+    )