reg: several fixes for MetaCache, more tests, and a new wrapper method for adding many datasets

Author: paulmueller

dcoraid/dbmodel/meta_cache.py

@@ -47,6 +47,8 @@ class MetaCache:
     The implementation uses an SQLite database which is loaded upon
     init and edited whenever data changes. For compute-intensive
     tasks (searching), the metadata are loaded into memory.
+
+    Datasets are sorted according to "metadata_created", descending.
     """
     def __init__(self,
                  directory: str | pathlib.Path,
@@ -191,73 +193,6 @@ def close(self):
         for db in self._databases.values():
             db.close()
 
-    def insert_many(self, org_id, dataset_dicts):
-        """Insert multiple datasets at once for one organization
-
-        This is essentially the vectorization of
-        :meth:`MetaCache._upsert_dataset_insert`.
-        This method MUST NOT be used for "updating" datasets.
-        """
-        ds_ids = [ds_dict["id"] for ds_dict in dataset_dicts]
-        ms_created = [ds_dict["metadata_created"] for ds_dict in dataset_dicts]
-        blobs_new = [
-            _create_blob_for_search(ds_dict) for ds_dict in dataset_dicts]
-        blob_max_len = max(len(b) for b in blobs_new)
-
-        if blob_max_len > int(self._srt_blobs.dtype["blob"].str[2:]):
-            # Increase the search blob size.
-            new_dtype = [("id", "<U36"),
-                         ("created", "<U26"),
-                         ("blob", f"<U{blob_max_len + 10}")
-                         ]
-        else:
-            new_dtype = self._srt_blobs.dtype
-
-        # registry
-        self._registry_org.setdefault(org_id, []).__add__(ds_ids)
-
-        # search array
-        dates_cur = np.array(self._srt_blobs["created"], copy=True)
-        size_old = dates_cur.size
-        dates_new = np.array(ms_created)
-        dates_comb = np.concatenate((dates_cur, dates_new))
-        sorter = np.argsort(dates_comb)
-        sorter_cur = sorter[:size_old]
-        sorter_new = sorter[size_old:]
-
-        new_blobs = np.empty(dates_comb.size, dtype=new_dtype)
-        new_blobs[sorter_cur] = self._srt_blobs
-        new_blobs["blob"][sorter_new] = blobs_new
-        new_blobs["created"][sorter_new] = dates_new
-        new_blobs["id"][sorter_new] = ds_ids
-
-        self._srt_blobs = new_blobs
-
-        # datasets
-        datasets_unsrt = self.datasets + dataset_dicts
-        self.datasets = [datasets_unsrt[ii] for ii in sorter]
-
-        # persistent database
-        if org_id not in self._databases:
-            self._databases[org_id] = SQLiteKeyJSONDatabase(
-                db_name=self.base_dir / f"org_{org_id}.db")
-        self._databases[org_id].insert_many(dataset_dicts)
-
-        # user's dataset list
-        datasets_user_owned_unsrt = (
-            self.datasets_user_owned
-            + [ds["creator_user_id"] == self.user_id for ds in dataset_dicts]
-        )
-        self.datasets_user_owned = [
-            datasets_user_owned_unsrt[ii] for ii in sorter]
-
-        # dataset IDs
-        dataset_ids_unsr = self._dataset_ids + ds_ids
-        self._dataset_ids = [dataset_ids_unsr[ii] for ii in sorter]
-
-        for (idx, ds_id) in enumerate(self._dataset_ids):
-            self._dataset_index_dict[ds_id] = idx
-
     def reset(self):
         """Reset the entire database"""
         with self._lock:
@@ -435,6 +370,116 @@ def _upsert_dataset_update(self, ds_dict):
         # persistent database
         self._databases[org_id][ds_id] = ds_dict
 
+    def upsert_many(self, dataset_dicts, org_id=None):
+        """Insert or update multiple datasets at once
+
+        The implementation is faster when the organization ID `org_id`
+        is specified. The recommended workflow is to query the DCOR
+        server by organization.
+        """
+        if self._registry_org.get(org_id):
+            # Separate the datasets into new and existing datasets.
+            ds_list_update = [ds_dict for ds_dict in dataset_dicts
+                              if ds_dict["id"] in self._dataset_ids]
+        else:
+            # We have not seen this organization before
+            ds_list_update = []
+
+        if not ds_list_update:
+            # Nothing needs to be updated, and all datasets are inserted.
+            ds_list_insert = dataset_dicts
+        else:
+            ds_list_insert = [ds_dict for ds_dict in dataset_dicts
+                              if ds_dict["id"] not in self._dataset_ids]
+
+        # Update datasets
+        for ds_dict in ds_list_update:
+            if self._databases[org_id][ds_dict["id"]] != ds_dict:
+                self._upsert_dataset_update(ds_dict)
+
+        if ds_list_insert:
+            # Insert datasets by organization
+            if org_id is not None:
+                # All datasets belong to one organization
+                self._upsert_many_insert(org_id=org_id,
+                                         dataset_dicts=ds_list_insert)
+            else:
+                # Iterate over all organization IDs in dataset_dicts
+                org_dict = {}
+                for ds_dict in ds_list_insert:
+                    org_dict.setdefault(ds_dict["owner_org"],
+                                        []).append(ds_dict)
+                for org_id, ds_list in org_dict.items():
+                    self._upsert_many_insert(org_id=org_id,
+                                             dataset_dicts=ds_list)
+
+    def _upsert_many_insert(self, org_id, dataset_dicts):
+        """Insert multiple datasets at once for one organization
+
+        This is essentially the vectorization of
+        :meth:`MetaCache._upsert_dataset_insert`.
+        This method MUST NOT be used for "updating" datasets.
+        """
+        ds_ids = [ds_dict["id"] for ds_dict in dataset_dicts]
+        ms_created = [ds_dict["metadata_created"] for ds_dict in dataset_dicts]
+        blobs_new = [
+            _create_blob_for_search(ds_dict) for ds_dict in dataset_dicts]
+        blob_max_len = max(len(b) for b in blobs_new)
+
+        if blob_max_len > int(self._srt_blobs.dtype["blob"].str[2:]):
+            # Increase the search blob size.
+            new_dtype = [("id", "<U36"),
+                         ("created", "<U26"),
+                         ("blob", f"<U{blob_max_len + 10}")
+                         ]
+        else:
+            new_dtype = self._srt_blobs.dtype
+
+        # registry
+        self._registry_org.setdefault(org_id, []).__add__(ds_ids)
+
+        # search array
+        dates_cur = np.array(self._srt_blobs["created"], copy=True)
+        size_old = dates_cur.size
+        dates_new = np.array(ms_created)
+        dates_comb = np.concatenate((dates_cur, dates_new))
+        # sort according to dates descending
+        sorter = np.argsort(dates_comb)[::-1]
+
+        new_blobs = np.empty(dates_comb.size, dtype=new_dtype)
+        new_blobs[:size_old] = self._srt_blobs["blob"]
+        new_blobs["blob"][size_old:] = blobs_new
+        new_blobs["created"][size_old:] = dates_new
+        new_blobs["id"][size_old:] = ds_ids
+        new_blobs = new_blobs[sorter]
+
+        self._srt_blobs = new_blobs
+
+        # datasets
+        datasets_unsrt = self.datasets + dataset_dicts
+        self.datasets = [datasets_unsrt[ii] for ii in sorter]
+
+        # persistent database
+        if org_id not in self._databases:
+            self._databases[org_id] = SQLiteKeyJSONDatabase(
+                db_name=self.base_dir / f"org_{org_id}.db")
+        self._databases[org_id].insert_many(dataset_dicts)
+
+        # user's dataset list
+        datasets_user_owned_unsrt = (
+            self.datasets_user_owned
+            + [ds["creator_user_id"] == self.user_id for ds in dataset_dicts]
+        )
+        self.datasets_user_owned = [
+            datasets_user_owned_unsrt[ii] for ii in sorter]
+
+        # dataset IDs
+        dataset_ids_unsr = self._dataset_ids + ds_ids
+        self._dataset_ids = [dataset_ids_unsr[ii] for ii in sorter]
+
+        for (idx, ds_id) in enumerate(self._dataset_ids):
+            self._dataset_index_dict[ds_id] = idx
+
 
 def _create_blob_for_search(ds_dict: dict) -> str:
     """Create a string blob from a dataset dictionary for free text search"""

tests/test_dbmodel_meta_cache.py

@@ -90,7 +90,7 @@ def test_cache_dataset_index_dict_2(tmp_path):
         ds_list.append(make_dataset_dict_full_fake(time_created=t0 + ii))
 
     # sort dataset list according to ID
-    ds_list_edit = sorted(ds_list, key = lambda x: x["id"])
+    ds_list_edit = sorted(ds_list, key=lambda x: x["id"])
     assert ds_list_edit != ds_list
 
     with meta_cache.MetaCache(tmp_path) as mc:
@@ -109,15 +109,15 @@ def test_cache_dataset_index_dict_2_multi(tmp_path):
 
     ds_list = []
     for ii in range(200):
-        ds_list.append(make_dataset_dict_full_fake(time_created=t0 + ii,
+        ds_list.append(make_dataset_dict_full_fake(time_created=t0 - ii,
                                                    org_id=org_id))
 
     # sort dataset list according to ID
-    ds_list_edit = sorted(ds_list, key = lambda x: x["id"])
+    ds_list_edit = sorted(ds_list, key=lambda x: x["id"])
     assert ds_list_edit != ds_list
 
     with meta_cache.MetaCache(tmp_path) as mc:
-        mc.insert_many(org_id, ds_list_edit)
+        mc._upsert_many_insert(org_id, ds_list_edit)
 
     for ii in range(200):
         assert mc._dataset_index_dict[ds_list[ii]["id"]] == ii
@@ -151,7 +151,6 @@ def test_cache_datasets_user_owned(tmp_path):
             assert mc.datasets[idx] == ds
             assert not mc.datasets_user_owned[idx]
 
-
     with meta_cache.MetaCache(tmp_path, user_id=user_id) as mc:
         # Make sure that the user dictionaries are recovered correctly
         for ds in ds_user:
@@ -187,12 +186,12 @@ def test_cache_datasets_user_owned_full(tmp_path):
             for _ in range(10)
             ]
         ds_user += ds_dicts_a
-        mc.insert_many(org_id, ds_dicts_a)
+        mc._upsert_many_insert(org_id, ds_dicts_a)
         ds_dicts_b = [
             make_dataset_dict_full_fake(org_id=org_id)
             for _ in range(10)
             ]
-        mc.insert_many(org_id, ds_dicts_b)
+        mc._upsert_many_insert(org_id, ds_dicts_b)
         ds_other += ds_dicts_b
 
         # Make sure that the user dictionaries are stored correctly
@@ -206,7 +205,6 @@ def test_cache_datasets_user_owned_full(tmp_path):
             assert mc.datasets[idx] == ds
             assert not mc.datasets_user_owned[idx]
 
-
     with meta_cache.MetaCache(tmp_path, user_id=user_id) as mc:
         # Make sure that the user dictionaries are recovered correctly
         for ds in ds_user:
@@ -220,9 +218,86 @@ def test_cache_datasets_user_owned_full(tmp_path):
             assert not mc.datasets_user_owned[idx]
 
 
+def test_cache_upsert_many(tmp_path):
+    """Test the `upsert_many` wrapper"""
+    ds_list = []
+
+    org_id = str(uuid.uuid4())
+
+    # create datasets in ascending order (order must be reversed in MetaCache)
+    for ii in range(15):
+        ds_list.append(make_dataset_dict_full_fake(
+            org_id=org_id,
+            time_created=time.time() - 100 + ii,
+            ))
+
+    with meta_cache.MetaCache(tmp_path) as mc:
+        mc.upsert_many(ds_list, org_id=org_id)
+
+        # make sure datasets are in descending order
+        assert mc.datasets == ds_list[::-1]
+
+    with meta_cache.MetaCache(tmp_path) as mc:
+        assert mc.datasets == ds_list[::-1]
+
+        ds_list[0]["title"] = "Peter Pan"
+        ds_list[1]["title"] = "Peter Pan 3"
+        ds_list[-1]["title"] = "Peter Pan 2"
+
+        mc.upsert_many(ds_list, org_id=org_id)
+
+        assert mc.datasets == ds_list[::-1]
+        # descending order again
+        assert mc.datasets[-1]["title"] == "Peter Pan"
+        assert mc.datasets[0]["title"] == "Peter Pan 2"
+        assert mc.datasets[-2]["title"] == "Peter Pan 3"
+
+
+def test_cache_upsert_many_search(tmp_path):
+    """Test the `upsert_many` wrapper, searching data afterwards"""
+    ds_list = []
+
+    org_id = str(uuid.uuid4())
+
+    # create datasets in ascending order (order must be reversed in MetaCache)
+    for ii in range(15):
+        ds_list.append(make_dataset_dict_full_fake(
+            org_id=org_id,
+            time_created=time.time() - 100 + ii,
+            ))
+
+    ds_list.append(make_dataset_dict_full_fake(title="Mordor",
+                                               org_id=org_id,
+                                               time_created=time.time() - 98,
+                                               ))
+
+    org_id_2 = str(uuid.uuid4())
+
+    ds_list.append(make_dataset_dict_full_fake(title="Frodo",
+                                               org_id=org_id_2,
+                                               time_created=time.time() - 98,
+                                               ))
+
+    for ii in range(18):
+        ds_list.append(make_dataset_dict_full_fake(
+            org_id=org_id_2,
+            time_created=time.time() - 100 + ii,
+            ))
+
+    with meta_cache.MetaCache(tmp_path) as mc:
+        mc.upsert_many(ds_list, org_id=org_id)
+        # search for them in freshly-modified cache
+        assert mc.search(query="Frodo")[0]["title"] == "Frodo"
+        assert mc.search(query="Mordor")[0]["title"] == "Mordor"
+
+    with meta_cache.MetaCache(tmp_path) as mc:
+        # search for them in loaded cache
+        assert mc.search(query="Frodo")[0]["title"] == "Frodo"
+        assert mc.search(query="Mordor")[0]["title"] == "Mordor"
+
 
 @pytest.mark.parametrize("previous_datasets", [0, 10, 100])
-def test_cache_insert_many(tmp_path, previous_datasets):
+def test_cache_upsert_many_insert(tmp_path, previous_datasets):
     # generate datasets
     all_ds_dicts = []
     org_ids = [str(uuid.uuid4()), str(uuid.uuid4())]
@@ -241,7 +316,7 @@ def test_cache_insert_many(tmp_path, previous_datasets):
             mc.upsert_dataset(make_dataset_dict_full_fake())
 
         for (org_id, ds_dicts) in zip(org_ids, all_ds_dicts):
-            mc.insert_many(org_id, ds_dicts)
+            mc._upsert_many_insert(org_id, ds_dicts)
 
     # Check whether the databases exist
     for org_id in org_ids:
@@ -406,7 +481,8 @@ def test_cache_search_updated_dataset_case(tmp_path):
     [[2, {"other-key": "dop"}, [["sinn"]]], ["2", "sinn"]],
     ])
 def test_values_only(input, output):
-    assert list(meta_cache._values_only(input, only_keys=["peter", "rin"])) == output
+    assert list(meta_cache._values_only(input,
+                                        only_keys=["peter", "rin"])) == output
 
 
 def test_create_blob_for_search():