support properties on attributes in the new merge (#4585)

* support properties on attributes

* support for attribute properties merge

Author: ajtmccarty

backend/infrahub/core/diff/merger/merger.py

@@ -3,7 +3,7 @@
 from typing import TYPE_CHECKING
 
 from infrahub.core.diff.model.path import BranchTrackingId
-from infrahub.core.diff.query.merge import DiffMergeQuery
+from infrahub.core.diff.query.merge import DiffMergePropertiesQuery, DiffMergeQuery
 
 if TYPE_CHECKING:
     from infrahub.core.branch import Branch
@@ -33,15 +33,23 @@ async def merge_graph(self, at: Timestamp) -> None:
         enriched_diff = await self.diff_repository.get_one(
             diff_branch_name=self.source_branch.name, tracking_id=BranchTrackingId(name=self.source_branch.name)
         )
-        node_diff_dicts = await self.serializer.serialize(diff=enriched_diff)
-        merge_query = await DiffMergeQuery.init(
-            db=self.db,
-            branch=self.source_branch,
-            at=at,
-            target_branch=self.destination_branch,
-            node_diff_dicts=node_diff_dicts,
-        )
-        await merge_query.execute(db=self.db)
+        async for node_diff_dicts, property_diff_dicts in self.serializer.serialize_diff(diff=enriched_diff):
+            merge_query = await DiffMergeQuery.init(
+                db=self.db,
+                branch=self.source_branch,
+                at=at,
+                target_branch=self.destination_branch,
+                node_diff_dicts=node_diff_dicts,
+            )
+            await merge_query.execute(db=self.db)
+            merge_properties_query = await DiffMergePropertiesQuery.init(
+                db=self.db,
+                branch=self.source_branch,
+                at=at,
+                target_branch=self.destination_branch,
+                property_diff_dicts=property_diff_dicts,
+            )
+            await merge_properties_query.execute(db=self.db)
 
         self.source_branch.branched_from = at.to_string()
         await self.source_branch.save(db=self.db)

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

@@ -1,4 +1,4 @@
-from typing_extensions import TypedDict
+from typing import TypedDict
 
 
 class RelationshipMergeDict(TypedDict):
@@ -17,3 +17,22 @@ class NodeMergeDict(TypedDict):
     action: str
     attributes: list[AttributeMergeDict]
     relationships: list[RelationshipMergeDict]
+
+
+class PropertyMergeDict(TypedDict):
+    property_type: str
+    action: str
+    value: str | bool | int | float | None
+
+
+class AttributePropertyMergeDict(TypedDict):
+    node_uuid: str
+    attribute_name: str
+    properties: list[PropertyMergeDict]
+
+
+class RelationshipPropertyMergeDict(TypedDict):
+    node_uuid: str
+    relationship_id: str
+    peer_uuid: str
+    properties: list[PropertyMergeDict]

backend/infrahub/core/diff/merger/serializer.py

@@ -1,21 +1,40 @@
+from typing import AsyncGenerator
+
 from infrahub.core.constants import DiffAction
 from infrahub.core.constants.database import DatabaseEdgeType
 from infrahub.core.schema.schema_branch import SchemaBranch
+from infrahub.types import ATTRIBUTE_PYTHON_TYPES
 
 from ..model.path import (
     ConflictSelection,
     EnrichedDiffAttribute,
     EnrichedDiffConflict,
+    EnrichedDiffProperty,
     EnrichedDiffRoot,
     EnrichedDiffSingleRelationship,
 )
-from .model import AttributeMergeDict, NodeMergeDict, RelationshipMergeDict
+from .model import (
+    AttributeMergeDict,
+    AttributePropertyMergeDict,
+    NodeMergeDict,
+    PropertyMergeDict,
+    RelationshipMergeDict,
+    RelationshipPropertyMergeDict,
+)
+
+Primitives = str | bool | int | float
 
 
 class DiffMergeSerializer:
-    def __init__(self, schema_branch: SchemaBranch) -> None:
+    def __init__(self, schema_branch: SchemaBranch, max_batch_size: int) -> None:
         self.schema_branch = schema_branch
+        self.max_batch_size = max_batch_size
         self._relationship_id_cache: dict[tuple[str, str], str] = {}
+        self._attribute_type_cache: dict[tuple[str, str], type] = {}
+
+    def _reset_caches(self) -> None:
+        self._relationship_id_cache = {}
+        self._attribute_type_cache = {}
 
     def _get_action(self, action: DiffAction, conflict: EnrichedDiffConflict | None) -> DiffAction:
         if not conflict:
@@ -29,32 +48,75 @@ def _get_action(self, action: DiffAction, conflict: EnrichedDiffConflict | None)
     def _to_action_str(self, action: DiffAction) -> str:
         return str(action.value).upper()
 
-    def _get_relationship_identifier(self, schema_kind_str: str, relationship_name: str) -> str:
-        cache_key = (schema_kind_str, relationship_name)
+    def _get_relationship_identifier(self, schema_kind: str, relationship_name: str) -> str:
+        cache_key = (schema_kind, relationship_name)
         if cache_key in self._relationship_id_cache:
             return self._relationship_id_cache[cache_key]
-        node_schema = self.schema_branch.get(name=schema_kind_str, duplicate=False)
+        node_schema = self.schema_branch.get(name=schema_kind, duplicate=False)
         relationship_schema = node_schema.get_relationship(name=relationship_name)
         relationship_identifier = relationship_schema.get_identifier()
         self._relationship_id_cache[cache_key] = relationship_identifier
         return relationship_identifier
 
-    async def serialize(self, diff: EnrichedDiffRoot) -> list[NodeMergeDict]:
+    def _get_property_type_for_attribute_value(self, schema_kind: str, attribute_name: str) -> type:
+        cache_key = (schema_kind, attribute_name)
+        if cache_key in self._attribute_type_cache:
+            return self._attribute_type_cache[cache_key]
+        node_schema = self.schema_branch.get(name=schema_kind, duplicate=False)
+        attribute_schema = node_schema.get_attribute(name=attribute_name)
+        python_type = ATTRIBUTE_PYTHON_TYPES[attribute_schema.kind]
+        final_python_type: type = str
+        if python_type in (str, int, float, bool):
+            final_python_type = python_type
+        self._attribute_type_cache[cache_key] = final_python_type
+        return final_python_type
+
+    def _convert_property_value(
+        self, property_type: DatabaseEdgeType, raw_value: str | None, value_type: type | None = None
+    ) -> Primitives | None:
+        if raw_value is None:
+            if property_type is DatabaseEdgeType.HAS_VALUE:
+                return "NULL"
+            return None
+        # peer IDs are strings
+        if property_type in (DatabaseEdgeType.HAS_OWNER, DatabaseEdgeType.HAS_SOURCE, DatabaseEdgeType.IS_RELATED):
+            return raw_value
+        # these are boolean
+        if property_type in (DatabaseEdgeType.IS_VISIBLE, DatabaseEdgeType.IS_PROTECTED):
+            return raw_value.lower() == "true"
+        # this must be HAS_VALUE
+        if value_type:
+            return value_type(raw_value)
+        return raw_value
+
+    async def serialize_diff(
+        self, diff: EnrichedDiffRoot
+    ) -> AsyncGenerator[
+        tuple[list[NodeMergeDict], list[AttributePropertyMergeDict | RelationshipPropertyMergeDict]], None
+    ]:
+        self._reset_caches()
         serialized_node_diffs = []
+        serialized_property_diffs: list[AttributePropertyMergeDict | RelationshipPropertyMergeDict] = []
         for node in diff.nodes:
             node_action = self._get_action(action=node.action, conflict=node.conflict)
-            attribute_diffs = [self._serialize_attribute(attribute_diff=attr_diff) for attr_diff in node.attributes]
+            attribute_diffs = []
+            for attr_diff in node.attributes:
+                attribute_diff, attribute_property_diff = self._serialize_attribute(
+                    attribute_diff=attr_diff, node_uuid=node.uuid, node_kind=node.kind
+                )
+                attribute_diffs.append(attribute_diff)
+                serialized_property_diffs.append(attribute_property_diff)
             relationship_diffs = []
             for rel_diff in node.relationships:
                 relationship_identifier = self._get_relationship_identifier(
-                    schema_kind_str=node.kind, relationship_name=rel_diff.name
+                    schema_kind=node.kind, relationship_name=rel_diff.name
                 )
                 for relationship_element_diff in rel_diff.relationships:
-                    relationship_diffs.extend(
-                        self._serialize_relationship_element(
-                            relationship_diff=relationship_element_diff, relationship_identifier=relationship_identifier
-                        )
+                    element_diffs, relationship_property_diffs = self._serialize_relationship_element(
+                        relationship_diff=relationship_element_diff, relationship_identifier=relationship_identifier
                     )
+                    relationship_diffs.extend(element_diffs)
+                    serialized_property_diffs.extend(relationship_property_diffs)
             serialized_node_diffs.append(
                 NodeMergeDict(
                     uuid=node.uuid,
@@ -63,41 +125,82 @@ async def serialize(self, diff: EnrichedDiffRoot) -> list[NodeMergeDict]:
                     relationships=relationship_diffs,
                 )
             )
-        return serialized_node_diffs
+            if len(serialized_node_diffs) == self.max_batch_size:
+                yield (serialized_node_diffs, serialized_property_diffs)
+                serialized_node_diffs, serialized_property_diffs = [], []
+        yield (serialized_node_diffs, serialized_property_diffs)
+
+    def _get_property_actions_and_values(
+        self, property_diff: EnrichedDiffProperty, python_value_type: type
+    ) -> list[tuple[DiffAction, Primitives]]:
+        action = property_diff.action
+        new_value = property_diff.new_value
+        if property_diff.conflict and property_diff.conflict.selected_branch is ConflictSelection.BASE_BRANCH:
+            action = property_diff.conflict.base_branch_action
+            if property_diff.conflict.base_branch_value:
+                new_value = property_diff.conflict.base_branch_value
+        actions = [action]
+        if property_diff.action is DiffAction.UPDATED:
+            actions = [DiffAction.ADDED, DiffAction.REMOVED]
+        actions_and_values: list[tuple[DiffAction, Primitives]] = []
+        for action in actions:
+            if action not in (DiffAction.ADDED, DiffAction.REMOVED):
+                continue
+            if action is DiffAction.ADDED:
+                raw_value = new_value
+            else:
+                raw_value = property_diff.previous_value
+            final_value = self._convert_property_value(
+                property_type=property_diff.property_type, raw_value=raw_value, value_type=python_value_type
+            )
+            if final_value:
+                actions_and_values.append((action, final_value))
+        return actions_and_values
 
-    def _serialize_attribute(self, attribute_diff: EnrichedDiffAttribute) -> AttributeMergeDict:
-        return AttributeMergeDict(
+    def _serialize_attribute(
+        self, attribute_diff: EnrichedDiffAttribute, node_uuid: str, node_kind: str
+    ) -> tuple[AttributeMergeDict, AttributePropertyMergeDict]:
+        prop_dicts: list[PropertyMergeDict] = []
+        python_type = self._get_property_type_for_attribute_value(
+            schema_kind=node_kind, attribute_name=attribute_diff.name
+        )
+        for property_diff in attribute_diff.properties:
+            actions_and_values = self._get_property_actions_and_values(
+                property_diff=property_diff, python_value_type=python_type
+            )
+            for action, value in actions_and_values:
+                prop_dicts.append(
+                    PropertyMergeDict(
+                        property_type=property_diff.property_type.value,
+                        action=self._to_action_str(action=action),
+                        value=value,
+                    )
+                )
+        attr_dict = AttributeMergeDict(
             name=attribute_diff.name,
             action=self._to_action_str(action=attribute_diff.action),
         )
+        attr_prop_dict = AttributePropertyMergeDict(
+            node_uuid=node_uuid, attribute_name=attribute_diff.name, properties=prop_dicts
+        )
+        return attr_dict, attr_prop_dict
 
     def _serialize_relationship_element(
         self, relationship_diff: EnrichedDiffSingleRelationship, relationship_identifier: str
-    ) -> list[RelationshipMergeDict]:
+    ) -> tuple[list[RelationshipMergeDict], list[RelationshipPropertyMergeDict]]:
         relationship_dicts = []
         for property_diff in relationship_diff.properties:
             if property_diff.property_type is not DatabaseEdgeType.IS_RELATED:
                 continue
-            action = property_diff.action
-            new_value = relationship_diff.peer_id
-            if property_diff.conflict and property_diff.conflict.selected_branch is ConflictSelection.BASE_BRANCH:
-                action = property_diff.conflict.base_branch_action
-                if property_diff.conflict.base_branch_value:
-                    new_value = property_diff.conflict.base_branch_value
-            actions = [action]
-            if property_diff.action is DiffAction.UPDATED:
-                actions = [DiffAction.ADDED, DiffAction.REMOVED]
-            actions_and_values: list[tuple[DiffAction, str]] = []
-            for action in actions:
-                if action is DiffAction.ADDED:
-                    actions_and_values.append((action, new_value))
-                elif action is DiffAction.REMOVED and property_diff.previous_value:
-                    actions_and_values.append((action, property_diff.previous_value))
-
-        for action, value in actions_and_values:
-            relationship_dicts.append(
-                RelationshipMergeDict(
-                    peer_id=value, name=relationship_identifier, action=self._to_action_str(action=action)
-                )
+            # TODO: next PR will add properties to this method
+            actions_and_values = self._get_property_actions_and_values(
+                property_diff=property_diff, python_value_type=str
             )
-        return relationship_dicts
+
+            for action, value in actions_and_values:
+                relationship_dicts.append(
+                    RelationshipMergeDict(
+                        peer_id=str(value), name=relationship_identifier, action=self._to_action_str(action=action)
+                    )
+                )
+        return relationship_dicts, []