unifying structure of returned history for db backend and file backends

pepkit · Apr 9, 2024 · 816e2f1 · 816e2f1
1 parent 346435f
commit 816e2f1
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 5 deletions.
diff --git a/pipestat/backends/db_backend/dbbackend.py b/pipestat/backends/db_backend/dbbackend.py
@@ -532,7 +532,7 @@ def retrieve_history_db(self, record_identifier, result_identifier=None):
                 columns = result_identifier
             else:
                 raise ValueError("Result identifier must be a str or list[str]")
-            for i in ["id", "record_identifier", MODIFIED_TIME]:
+            for i in ["id", MODIFIED_TIME]:
                 if i not in columns:
                     columns = [i] + columns
 
@@ -584,8 +584,25 @@ def retrieve_history_db(self, record_identifier, result_identifier=None):
                 record_dict = dict(record._mapping)
                 end_results.append(record_dict)
 
+        # This next step is to process the results such that they will match output similar to the filebackend
+
+        collected_keys = []
+        new_history_dict = {}
+        for result in end_results:
+            for key, value in result.items():
+                if key == MODIFIED_TIME:
+                    continue
+                elif value:
+                    if key not in new_history_dict:
+                        collected_keys.append(key)
+                        new_history_dict[key] = {result[MODIFIED_TIME]: {"reported_result": value}}
+                    else:
+                        new_history_dict[key].update(
+                            {result[MODIFIED_TIME]: {"reported_result": value}}
+                        )
+
         records_dict = {
-            "history": end_results,
+            "history": new_history_dict,
         }
 
         return records_dict

diff --git a/tests/test_pipestat.py b/tests/test_pipestat.py
@@ -2420,12 +2420,12 @@ def test_select_history_basic(
 
             all_history_result = psm.retrieve_history(record_identifier="sample1")
 
+            assert len(history_result.keys()) == 2
+
             if backend == "file":
-                assert len(history_result.keys()) == 2
                 assert len(all_history_result.keys()) == 4
 
             if backend == "db":
-                assert len(history_result) == 2 and isinstance(history_result, list)
-                assert len(all_history_result) == 2
+                assert len(all_history_result.keys()) == 6
 
             print("Done")