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

@@ -115,115 +115,73 @@ def test_resolve_db_regions_with_none():
     assert composition_db_doc.as_dict() == resolved_data
 
 
-def test_resolve_local_regions():
-    full_recipe_data = {
-        "name": "one_sphere",
-        "objects": {
-            "sphere_25": {
-                "type": "single_sphere",
-                "max_jitter": [1, 1, 0],
-            },
-        },
-        "composition": {
-            "space": {"regions": {"interior": ["A"]}},
-            "A": {"object": "sphere_25", "count": 500},
-        },
-    }
-
-    local_data = CompositionDoc(
-        name="space",
-        count=None,
+def test_build_dependency_graph():
+    composition_db_doc = CompositionDoc(
+        name="test",
+        count=1,
         object=None,
-        regions={
-            "interior": [
-                "A",
-            ]
-        },
+        regions=None,
         molarity=None,
     )
-
-    resolved_data = {
-        "name": "space",
-        "object": None,
-        "count": None,
-        "molarity": None,
-        "regions": {
-            "interior": [
-                {
-                    "name": "A",
-                    "object": {"type": "single_sphere", "max_jitter": [1, 1, 0]},
-                    "count": 500,
-                    "molarity": None,
-                    "regions": {},
-                }
-            ]
-        },
+    compositions = {
+        "space": {"regions": {"interior": ["A"]}},
+        "A": {"object": "sphere_25", "count": 5},
     }
-    local_data.resolve_local_regions(local_data.as_dict(), full_recipe_data, mock_db)
-    assert local_data.as_dict() == resolved_data
+    dependency_map = composition_db_doc.build_dependency_graph(compositions)
+    assert dependency_map == {"space": ["A"], "A": []}
 
 
-def test_check_and_replace_references():
-    objects_to_path_map = {"test_obj": "firebase:objects/test_id"}
-    references_to_update = {}
-    composition_doc = CompositionDoc(
+def test_build_dependency_graph_with_complex_compositions():
+    composition_db_doc = CompositionDoc(
         name="test",
         count=1,
-        object="firebase:objects/test_id",
-        regions={"interior": [{"object": "test_obj", "count": 1}]},
+        object=None,
+        regions=None,
         molarity=None,
     )
-    composition_doc.check_and_replace_references(
-        objects_to_path_map, references_to_update, mock_db
-    )
-    assert composition_doc.as_dict()["object"] == "firebase:objects/test_id"
-    assert composition_doc.as_dict()["regions"] == {
-        "interior": [{"object": "firebase:objects/test_id", "count": 1}]
+    complex_compositions = {
+        "space": {"regions": {"interior": ["tree", "A", "B", "C"]}},
+        "tree": {"object": "sphere_tree_A", "molarity": 1},
+        "A": {
+            "object": "sphere_100",
+            "regions": {
+                "surface": [
+                    {"object": "sphere_50", "count": 5},
+                    {"object": "sphere_75", "count": 1},
+                ],
+                "interior": [{"object": "sphere_25", "count": 30, "priority": -1}],
+            },
+        },
+        "B": {
+            "object": "sphere_100",
+            "regions": {"surface": ["B_sphere_75"], "interior": ["B_sphere_50"]},
+        },
+        "B_sphere_50": {"object": "sphere_50", "count": 8, "priority": -1},
+        "B_sphere_75": {"object": "sphere_75", "count": 3},
     }
 
-
-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": {},
+    dependency_map = composition_db_doc.build_dependency_graph(complex_compositions)
+    assert dependency_map == {
+        "space": ["tree", "A", "B"],
+        "tree": [],
+        "A": [],
+        "B": ["B_sphere_75", "B_sphere_50"],
+        "B_sphere_75": [],
+        "B_sphere_50": [],
     }
-    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(
+def test_comp_upload_order():
+    composition_db_doc = CompositionDoc(
         name="test",
         count=1,
         object=None,
-        regions={},
+        regions=None,
         molarity=None,
     )
-
-    doc, doc_id = composition_doc.should_write(mock_db, recipe_data)
-    assert doc_id is not None
-    assert doc == existing_doc
+    compositions = {
+        "space": {"regions": {"interior": ["A"]}},
+        "A": {"object": "sphere_25", "count": 5},
+    }
+    upload_order = composition_db_doc.comp_upload_order(compositions)
+    assert upload_order == ["A", "space"]

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,28 +83,11 @@ 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_doc.upload_compositions(composition, recipe_to_save)
     assert composition_doc.comp_to_path_map == {
-        "space": {"id": "test_id", "path": "firebase:composition/test_id"},
-    }
-    assert references_to_update == {
-        "space": [{"comp_id": "test_id", "index": "regions.interior", "name": "A"}]
+        "space": "firebase:composition/test_id",
     }
 
 
@@ -185,7 +168,7 @@ def test_upload_recipe():
     recipe_doc = DBUploader(mock_db)
     recipe_doc.upload_recipe(recipe_meta_data, recipe_data)
     assert recipe_doc.comp_to_path_map == {
-        "space": {"path": "firebase:composition/test_id", "id": "test_id"},
-        "A": {"path": "firebase:composition/test_id", "id": "test_id"},
+        "space": "firebase:composition/test_id",
+        "A": "firebase:composition/test_id",
     }
     assert recipe_doc.objects_to_path_map == {"sphere_25": "firebase:objects/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"]