allo fetching from warmpath

Author: heidi-holm-4ss

fourinsight/engineroom/utils/_datamanage.py

@@ -512,6 +512,8 @@ class DrioDataSource(BaseDataSource):
         DataReservoir.io client.
     lables : dict
         Labels and timeseries IDs as key/value pairs.
+    storage : str, optional
+        If 'warm' (default), drio is fetched from warm storage. If 'archive', data is fetched from the archive.
     index_type : str or obj
         Index type (see Notes). Should be 'datetime', 'integer' or an `index converter`
         object.
@@ -564,6 +566,7 @@ def __init__(
         self,
         drio_client,
         labels,
+        storage="warm",
         index_type="datetime",
         index_sync=False,
         tolerance=None,
@@ -572,6 +575,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 +590,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 +629,20 @@ 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>")