Refactor uploading system deduplication using hash

cellpack/autopack/FirebaseHandler.py

@@ -101,7 +101,7 @@ def set_doc(self, collection, id, data):
             return doc_ref
         else:
             logging.error(
-                f"ERROR: {doc_ref.path} already exists. If uploading new data, provide a unique recipe name."
+                f"WARNING: {doc_ref.path} already exists. If uploading new data, provide a unique recipe name."
             )
             return
 
@@ -244,3 +244,15 @@ def is_firebase_obj(obj):
         return isinstance(
             obj, (firestore.DocumentReference, firestore.DocumentSnapshot)
         )
+
+    def check_doc_existence(self, collection, doc_data):
+        doc_hash = doc_data.get("dedup_hash")
+        doc_name = doc_data.get("name")
+        query = (
+            self.db.collection(collection).where("dedup_hash", "==", doc_hash).limit(1)
+        )
+        docs = list(query.stream())
+        if any(docs):
+            logging.info(f"{doc_name} already exists in database with id {docs[0].id}.")
+            return docs[0].reference
+        return None

cellpack/tests/mocks/mock_db.py

@@ -8,6 +8,7 @@ def __init__(self, data) -> None:
             self.obj["id"] = index
         self.data = data
         self.name = "test_db"
+        self._existing_doc = []
 
     @staticmethod
     def is_firebase_obj(obj):
@@ -69,3 +70,15 @@ def update_reference_on_doc(self, doc_ref, index, new_item_ref):
 
     def update_elements_in_array(self, doc_ref, index, new_item_ref, remove_item):
         return True
+
+    def check_doc_existence(self, collection, doc_data):
+        doc_hash = doc_data.get("dedup_hash")
+        if len(self._existing_doc) == 0:
+            self._existing_doc.append({"dedup_hash": doc_hash})
+
+        if self.data.get("dedup_hash") == doc_hash:
+            mock_ref = Mock()
+            mock_ref.id = "existing_doc_id"
+            return mock_ref
+
+        return None

cellpack/tests/test_comp_doc.py

@@ -180,50 +180,3 @@ def test_check_and_replace_references():
     assert composition_doc.as_dict()["regions"] == {
         "interior": [{"object": "firebase:objects/test_id", "count": 1}]
     }
-
-
-def test_composition_doc_should_write_with_no_existing_doc():
-    recipe_data = {
-        "bounding_box": [[0, 0, 0], [10, 10, 10]],
-        "name": "test",
-        "count": 1,
-        "objects": None,
-        "regions": {},
-    }
-    composition_doc = CompositionDoc(
-        name="test",
-        count=1,
-        object=None,
-        regions={},
-        molarity=None,
-    )
-    doc, doc_id = composition_doc.should_write(mock_db, recipe_data)
-    assert doc_id is None
-    assert doc is None
-
-
-def test_composition_doc_should_write_with_existing_doc():
-    existing_doc = {
-        "name": "test",
-        "count": 1,
-        "object": None,
-        "regions": {},
-        "molarity": None,
-    }
-    mock_db.data = existing_doc
-    recipe_data = {
-        "name": "test",
-        "count": 1,
-        "objects": None,
-    }
-    composition_doc = CompositionDoc(
-        name="test",
-        count=1,
-        object=None,
-        regions={},
-        molarity=None,
-    )
-
-    doc, doc_id = composition_doc.should_write(mock_db, recipe_data)
-    assert doc_id is not None
-    assert doc == existing_doc

cellpack/tests/test_data_doc.py

@@ -36,3 +36,12 @@ def test_is_nested_list():
 def test_is_obj():
     assert DataDoc.is_obj(object_example) is True
     assert DataDoc.is_obj(composition_example) is False
+
+
+def test_generate_hash():
+    test_cases = [object_example, composition_example, None]
+
+    for input_data in test_cases:
+        generated_hash = DataDoc.generate_hash(input_data)
+        assert isinstance(generated_hash, str)
+        assert generated_hash == DataDoc.generate_hash(input_data)

cellpack/tests/test_db_uploader.py

@@ -83,22 +83,10 @@ def test_upload_compositions():
         "space": {"regions": {"interior": ["A"]}},
     }
     recipe_to_save = {"format_version": "2.1", "name": "one_sphere", "composition": {}}
-    recipe_data = {
-        "name": "one_sphere",
-        "objects": {
-            "sphere_25": {
-                "type": "single_sphere",
-                "max_jitter": [1, 1, 0],
-            },
-        },
-        "composition": {
-            "space": {"regions": {"interior": ["A"]}},
-        },
-    }
 
     composition_doc = DBUploader(mock_db)
     references_to_update = composition_doc.upload_compositions(
-        composition, recipe_to_save, recipe_data
+        composition, recipe_to_save
     )
     assert composition_doc.comp_to_path_map == {
         "space": {"id": "test_id", "path": "firebase:composition/test_id"},

cellpack/tests/test_gradient_doc.py

@@ -1,21 +1,25 @@
-from cellpack.autopack.DBRecipeHandler import GradientDoc
+from cellpack.autopack.DBRecipeHandler import DataDoc
 from cellpack.tests.mocks.mock_db import MockDB
 
 mock_db = MockDB({})
 
 
-def test_should_write_with_no_existing_doc():
-    gradient_doc = GradientDoc({"name": "test_grad_name", "test_key": "test_value"})
-    doc, doc_id = gradient_doc.should_write(mock_db, "test_grad_name")
-    assert doc_id is None
-    assert doc is None
+def test_check_doc_existence_for_gradient():
+    data_doc = DataDoc()
+    gradient_doc = {
+        "nucleus_surface_gradient": {
+            "mode": "surface",
+            "name": "nucleus_surface_gradient",
+        }
+    }
 
+    test_data = {
+        "doc": gradient_doc,
+        "dedup_hash": data_doc.generate_hash({"doc": gradient_doc}),
+    }
 
-def test_should_write_with_existing_doc():
-    existing_doc = {"name": "test_grad_name", "test_key": "test_value"}
-    mock_db.data = existing_doc
-    gradient_doc = GradientDoc({"name": "test_grad_name", "test_key": "test_value"})
+    result = mock_db.check_doc_existence("gradients", test_data)
+    assert result is None
 
-    doc, doc_id = gradient_doc.should_write(mock_db, "test_grad_name")
-    assert doc_id is not None
-    assert doc is not None
+    existing_doc = mock_db._existing_doc[0]
+    assert existing_doc["dedup_hash"] == test_data["dedup_hash"]

cellpack/tests/test_object_doc.py

@@ -44,18 +44,24 @@ def test_convert_representation():
     assert result == expected_result
 
 
-def test_object_doc_should_write_with_no_existing_doc():
-    object_doc = ObjectDoc("test", {"test_key": "test_value"})
-    doc, doc_id = object_doc.should_write(mock_db)
-    assert doc_id is None
-    assert doc is None
+def test_check_doc_existence_for_object():
+    object_doc = {
+        "peroxisome": {
+            "gradient": {
+                "mode": "surface",
+                "name": "nucleus_surface_gradient",
+            },
+            "name": "peroxisome",
+        }
+    }
 
+    test_data = {
+        "doc": object_doc,
+        "dedup_hash": ObjectDoc.generate_hash({"doc": object_doc}),
+    }
 
-def test_object_doc_should_write_with_existing_doc():
-    existing_doc = {"name": "test", "test_key": "test_value"}
-    mock_db.data = existing_doc
-    object_doc = ObjectDoc("test", {"test_key": "test_value"})
+    result = mock_db.check_doc_existence("objects", test_data)
+    assert result is None
 
-    doc, doc_id = object_doc.should_write(mock_db)
-    assert doc_id is not None
-    assert doc == existing_doc
+    existing_doc = mock_db._existing_doc[0]
+    assert existing_doc["dedup_hash"] == test_data["dedup_hash"]