fixes to allow getting a diff for a single branch (#5842)

Author: ajtmccarty

backend/infrahub/core/diff/model/path.py

@@ -409,8 +409,8 @@ class EnrichedDiffRootMetadata(BaseSummary):
     from_time: Timestamp
     to_time: Timestamp
     uuid: str
-    partner_uuid: str
     tracking_id: TrackingId
+    partner_uuid: str | None = field(default=None)
     exists_on_database: bool = field(default=False)
 
     def __hash__(self) -> int:

backend/infrahub/core/diff/query/save.py

@@ -42,6 +42,7 @@ async def query_init(self, db: InfrahubDatabase, **kwargs: Any) -> None:
 }
 WITH DISTINCT diff_root AS diff_root
 WITH collect(diff_root) AS diff_roots
+WHERE SIZE(diff_roots) = 2
 CALL {
     WITH diff_roots
     WITH diff_roots[0] AS base_diff_node, diff_roots[1] AS branch_diff_node

backend/infrahub/core/diff/repository/deserializer.py

@@ -164,6 +164,7 @@ def _deserialize_diff_root(self, root_node: Neo4jNode) -> EnrichedDiffRoot:
     def build_diff_root_metadata(cls, root_node: Neo4jNode) -> EnrichedDiffRootMetadata:
         from_time = Timestamp(str(root_node.get("from_time")))
         to_time = Timestamp(str(root_node.get("to_time")))
+        partner_uuid = cls._get_str_or_none_property_value(node=root_node, property_name="partner_uuid")
         tracking_id_str = str(root_node.get("tracking_id"))
         tracking_id = deserialize_tracking_id(tracking_id_str=tracking_id_str)
         return EnrichedDiffRootMetadata(
@@ -172,7 +173,7 @@ def build_diff_root_metadata(cls, root_node: Neo4jNode) -> EnrichedDiffRootMetad
             from_time=from_time,
             to_time=to_time,
             uuid=str(root_node.get("uuid")),
-            partner_uuid=str(root_node.get("partner_uuid")),
+            partner_uuid=partner_uuid,
             tracking_id=tracking_id,
             num_added=int(root_node.get("num_added", 0)),
             num_updated=int(root_node.get("num_updated", 0)),

backend/infrahub/core/diff/repository/repository.py

@@ -154,18 +154,23 @@ async def get_pairs(
             diff_branch_names=[base_branch_name],
             max_depth=max_depth,
             batch_size_limit=batch_size_limit,
-            diff_ids=[d.partner_uuid for d in diffs_by_uuid.values()],
+            diff_ids=[d.partner_uuid for d in diffs_by_uuid.values() if d.partner_uuid],
         )
         diffs_by_uuid.update({bbr.uuid: bbr for bbr in base_branch_roots})
-        return [
-            EnrichedDiffs(
-                base_branch_name=base_branch_name,
-                diff_branch_name=diff_branch_name,
-                base_branch_diff=diffs_by_uuid[dbr.partner_uuid],
-                diff_branch_diff=dbr,
+        diff_pairs = []
+        for dbr in diff_branch_roots:
+            if dbr.partner_uuid is None:
+                continue
+            base_branch_diff = diffs_by_uuid[dbr.partner_uuid]
+            diff_pairs.append(
+                EnrichedDiffs(
+                    base_branch_name=base_branch_name,
+                    diff_branch_name=diff_branch_name,
+                    base_branch_diff=base_branch_diff,
+                    diff_branch_diff=dbr,
+                )
             )
-            for dbr in diff_branch_roots
-        ]
+        return diff_pairs
 
     async def hydrate_diff_pair(
         self,
@@ -325,7 +330,7 @@ async def get_diff_pairs_metadata(
         roots_by_id = {root.uuid: root for root in empty_roots}
         pairs: list[EnrichedDiffsMetadata] = []
         for branch_root in empty_roots:
-            if branch_root.base_branch_name == branch_root.diff_branch_name:
+            if branch_root.base_branch_name == branch_root.diff_branch_name or branch_root.partner_uuid is None:
                 continue
             base_root = roots_by_id[branch_root.partner_uuid]
             pairs.append(

backend/tests/unit/core/diff/test_coordinator.py

@@ -8,7 +8,7 @@
 from infrahub.core.diff.calculator import DiffCalculator
 from infrahub.core.diff.combiner import DiffCombiner
 from infrahub.core.diff.coordinator import DiffCoordinator
-from infrahub.core.diff.model.path import BranchTrackingId, EnrichedDiffRootMetadata
+from infrahub.core.diff.model.path import BranchTrackingId, EnrichedDiffRootMetadata, NameTrackingId
 from infrahub.core.diff.repository.repository import DiffRepository
 from infrahub.core.initialization import create_branch
 from infrahub.core.manager import NodeManager
@@ -240,3 +240,64 @@ async def test_unrelated_changes_skip_some_expensive_operations(
         )
         assert no_changes_diff.from_time == no_changes_diff_metadata.from_time == diff_with_data.from_time
         assert no_changes_diff.to_time == no_changes_diff_metadata.to_time
+
+    async def test_diff_on_default_branch_only(
+        self,
+        db: InfrahubDatabase,
+        default_branch: Branch,
+        person_john_main: Node,
+        person_jane_main: Node,
+        person_alfred_main: Node,
+        car_camry_main: Node,
+        car_accord_main: Node,
+    ) -> None:
+        branch = await create_branch(db=db, branch_name="branch1")
+
+        updated_person = await NodeManager.get_one(db=db, id=person_john_main.id)
+        updated_person.height.value = 200
+        await updated_person.save(db=db)
+
+        new_person = await Node.init(db=db, schema="TestPerson", branch=default_branch)
+        await new_person.new(db=db, name="Jeff", height=170)
+        await new_person.save(db=db)
+
+        deleted_person = await NodeManager.get_one(db=db, id=person_alfred_main.id)
+        await deleted_person.delete(db=db)
+
+        updated_car = await NodeManager.get_one(db=db, id=car_accord_main.id)
+        await updated_car.owner.update(db=db, data=new_person)
+        await updated_car.save(db=db)
+
+        from_time = Timestamp(branch.get_branched_from())
+        to_time = Timestamp()
+        name = str(uuid4())
+        diff_coordinator = await self.get_wrapped_diff_coordinator(db=db, branch=default_branch)
+        component_registry = get_component_registry()
+        diff_repository = await component_registry.get_component(DiffRepository, db=db, branch=default_branch)
+        main_diff_metadata = await diff_coordinator.create_or_update_arbitrary_timeframe_diff(
+            base_branch=default_branch, diff_branch=default_branch, from_time=from_time, to_time=to_time, name=name
+        )
+        main_diff = await diff_repository.get_one(diff_branch_name=default_branch.name, diff_id=main_diff_metadata.uuid)
+
+        assert main_diff.base_branch_name == default_branch.name
+        assert main_diff.diff_branch_name == default_branch.name
+        assert main_diff.from_time == from_time
+        assert main_diff.to_time == to_time
+        assert main_diff.tracking_id == NameTrackingId(name=name)
+        assert len(main_diff.nodes) == 4
+        nodes_by_id = {n.uuid: n for n in main_diff.nodes}
+        assert set(nodes_by_id.keys()) == {updated_person.id, new_person.id, deleted_person.id, updated_car.id}
+        new_person_diff = nodes_by_id[new_person.id]
+        assert new_person_diff.action is DiffAction.ADDED
+        deleted_person_diff = nodes_by_id[deleted_person.id]
+        assert deleted_person_diff.action is DiffAction.REMOVED
+        updated_car_diff = nodes_by_id[updated_car.id]
+        assert updated_car_diff.action is DiffAction.UPDATED
+        assert updated_car_diff.attributes == set()
+        rel_diffs = {(r.name, r.action) for r in updated_car_diff.relationships}
+        assert rel_diffs == {("owner", DiffAction.UPDATED)}
+        updated_person_diff = nodes_by_id[updated_person.id]
+        rel_diffs = {(r.name, r.action) for r in updated_person_diff.relationships}
+        assert rel_diffs == {("cars", DiffAction.UPDATED)}
+        attr_diffs = {(a.name, a.action) for a in updated_person_diff.attributes}
+        assert attr_diffs == {("height", DiffAction.UPDATED)}