Enable DrioDataSource to fetch data from warm path (#61)

* allo fetching from warmpath
* add tests
* black
* set archive as default
* update docstring, add experimental for warm storage

Author: heidi-holm-4ss

fourinsight/engineroom/utils/_datamanage.py

@@ -510,8 +510,11 @@ class DrioDataSource(BaseDataSource):
     ----------
     drio_client : obj
         DataReservoir.io client.
-    lables : dict
+    labels : dict
         Labels and timeseries IDs as key/value pairs.
+    storage : str, optional
+        Where to fetch data from. If 'archive' (default), data is fetched from the archive.
+        If 'warm' (experimental), data is fetched from warm storage, which supports aggregated queries.
     index_type : str or obj
         Index type (see Notes). Should be 'datetime', 'integer' or an `index converter`
         object.
@@ -558,12 +561,16 @@ class DrioDataSource(BaseDataSource):
       be given as a dtype that the :meth:`index_converter.to_universal_delta` can
       parse.
 
+    - The 'warm' storage option is experimental and its behavior may change in
+      future releases.
+
     """
 
     def __init__(
         self,
         drio_client,
         labels,
+        storage="archive",
         index_type="datetime",
         index_sync=False,
         tolerance=None,
@@ -572,6 +579,7 @@ def __init__(
         **get_kwargs,
     ):
         self._drio_client = drio_client
+        self.storage = storage
         self._get_kwargs = get_kwargs
 
         self._labels = {lab: id.strip() for lab, id in labels.items()}
@@ -586,6 +594,9 @@ def __init__(
         else:
             raise ValueError("'index_type' should be 'datetime' or 'integer'.")
 
+        if storage not in ["warm", "archive"]:
+            raise ValueError("storage must be either 'warm' or 'archive'")
+
         super().__init__(
             index_converter,
             index_sync=index_sync,
@@ -622,10 +633,24 @@ def _get(self, start, end):
             Label and data as key/value pairs. The data is returned as ``pandas.Series``
             objects.
         """
+        if self.storage == "warm":
+
+            def get_fun(ts_id, start, end, **kwargs):
+                aggregation_period = kwargs.pop("aggregation_period", "tick")
+                aggregation_function = kwargs.pop("aggregation_function", "mean")
+                return self._drio_client.get_samples_aggregate(
+                    ts_id,
+                    start=start,
+                    end=end,
+                    aggregation_period=aggregation_period,
+                    aggregation_function=aggregation_function,
+                )
+
+        elif self.storage == "archive":
+            get_fun = self._drio_client.get
+
         return {
-            label: self._drio_client.get(
-                ts_id, start=start, end=end, **self._get_kwargs
-            )
+            label: get_fun(ts_id, start=start, end=end, **self._get_kwargs)
             for label, ts_id in self._labels.items()
         }
 

tests/test_datamanage.py

@@ -1079,7 +1079,7 @@ def test__get(self):
             "c": "timeseriesid-c",
         }
         source = DrioDataSource(
-            drio_client, labels, convert_date=True, raise_empty=False
+            drio_client, labels, storage="archive", convert_date=True, raise_empty=False
         )
 
         data_out = source._get("<start-time>", "<end-time>")
@@ -1120,6 +1120,121 @@ def test__get(self):
             ]
         )
 
+    def test__get_warm_storage(self):
+        # Mock the drio_client
+        drio_client = Mock()
+        drio_client.get_samples_aggregate.return_value = pd.Series(
+            data=[1.1, 1.2, 1.3], index=[1.0, 2.0, 3.0]
+        )
+
+        # Labels to request
+        labels = {
+            "a": "timeseriesid-a",
+            "b": "timeseriesid-b",
+            "c": "timeseriesid-c",
+        }
+
+        # Instantiate DrioDataSource with storage="warm" and constructor kwargs
+        source = DrioDataSource(
+            drio_client,
+            labels,
+            storage="warm",
+            convert_date=True,
+            raise_empty=False,
+            aggregation_period="1h",
+            aggregation_function="min",
+        )
+
+        # Call _get (no kwargs allowed per your current design)
+        data_out = source._get("<start-time>", "<end-time>")
+
+        # Expected output
+        data_expect = {
+            "a": pd.Series(data=[1.1, 1.2, 1.3], index=[1.0, 2.0, 3.0]),
+            "b": pd.Series(data=[1.1, 1.2, 1.3], index=[1.0, 2.0, 3.0]),
+            "c": pd.Series(data=[1.1, 1.2, 1.3], index=[1.0, 2.0, 3.0]),
+        }
+
+        # Assert returned data is correct
+        assert data_out.keys() == data_expect.keys()
+        for key, series_expect in data_expect.items():
+            pd.testing.assert_series_equal(series_expect, data_out[key])
+
+        # Assert get_samples_aggregate called with correct parameters
+        drio_client.get_samples_aggregate.assert_has_calls(
+            [
+                call(
+                    "timeseriesid-a",
+                    start="<start-time>",
+                    end="<end-time>",
+                    aggregation_period="1h",
+                    aggregation_function="min",
+                ),
+                call(
+                    "timeseriesid-b",
+                    start="<start-time>",
+                    end="<end-time>",
+                    aggregation_period="1h",
+                    aggregation_function="min",
+                ),
+                call(
+                    "timeseriesid-c",
+                    start="<start-time>",
+                    end="<end-time>",
+                    aggregation_period="1h",
+                    aggregation_function="min",
+                ),
+            ]
+        )
+
+    def test__get_warm_defaults(self):
+        # Mock the drio_client
+        drio_client = Mock()
+        drio_client.get_samples_aggregate.return_value = pd.Series(
+            data=[10, 20, 30], index=[1, 2, 3]
+        )
+
+        labels = {
+            "x": "ts-x",
+            "y": "ts-y",
+        }
+
+        # Instantiate DrioDataSource with storage="warm", no overrides
+        source = DrioDataSource(
+            drio_client, labels, storage="warm", convert_date=True, raise_empty=False
+        )
+
+        # Call _get with defaults
+        data_out = source._get("<start>", "<end>")
+
+        # Expected series
+        expected = pd.Series([10, 20, 30], index=[1, 2, 3])
+
+        # Assert keys and values
+        assert data_out.keys() == labels.keys()
+        for key in labels:
+            pd.testing.assert_series_equal(data_out[key], expected)
+
+        # Assert get_samples_aggregate called with default aggregation values
+        drio_client.get_samples_aggregate.assert_has_calls(
+            [
+                call(
+                    "ts-x",
+                    start="<start>",
+                    end="<end>",
+                    aggregation_period="tick",
+                    aggregation_function="mean",
+                ),
+                call(
+                    "ts-y",
+                    start="<start>",
+                    end="<end>",
+                    aggregation_period="tick",
+                    aggregation_function="mean",
+                ),
+            ]
+        )
+
 
 class Test_NullDataSource:
     def test__init__(self):