diff --git a/vitrage/graph/driver/networkx_graph.py b/vitrage/graph/driver/networkx_graph.py index cb76641ab..2e9d82903 100644 --- a/vitrage/graph/driver/networkx_graph.py +++ b/vitrage/graph/driver/networkx_graph.py @@ -13,13 +13,14 @@ # under the License. -from oslo_log import log as logging - import copy import json + import networkx as nx from networkx.readwrite import json_graph +from oslo_log import log as logging + from elements import Edge from elements import Vertex from graph import Direction @@ -29,7 +30,6 @@ from vitrage.common.constants import VertexProperties as VProps from vitrage.graph.filter import check_filter from vitrage.graph.query import create_predicate - LOG = logging.getLogger(__name__) @@ -152,10 +152,10 @@ class NXGraph(Graph): :type v: Vertex """ - if hard_update: - properties = self._g.node.get(v.vertex_id, None) - if properties: - properties.clear() + orig_prop = self._g.node.get(v.vertex_id, None) + if orig_prop and hard_update: + orig_prop.clear() + v.properties = self._update_properties(orig_prop, v.properties) self._add_vertex(v) @Notifier.update_notify @@ -164,14 +164,27 @@ class NXGraph(Graph): :type e: Edge """ - if hard_update: - properties = self._get_edge_properties(e.source_id, - e.target_id, - e.label) - if properties: - properties.clear() + orig_prop = self._g.edge.get( + e.source_id, {}).get( + e.target_id, {}).get( + e.label, None) + if orig_prop and hard_update: + orig_prop.clear() + e.properties = self._update_properties(orig_prop, e.properties) self._add_edge(e) + @staticmethod + def _update_properties(orig_props, new_props): + if orig_props is None: + orig_props = dict() + keys_to_remove = [key for key, val in new_props.items() if val is None] + for key in keys_to_remove: + del new_props[key] + if key in orig_props: + del orig_props[key] + orig_props.update(new_props) + return orig_props + def remove_vertex(self, v): """Remove Vertex v and its edges from the graph diff --git a/vitrage/tests/unit/graph/test_graph.py b/vitrage/tests/unit/graph/test_graph.py index 496d2e1ee..b6e1bbafc 100644 --- a/vitrage/tests/unit/graph/test_graph.py +++ b/vitrage/tests/unit/graph/test_graph.py @@ -18,7 +18,6 @@ test_vitrage graph Tests for `vitrage` graph driver """ -from oslo_log import log as logging from vitrage.common.constants import EdgeProperties as EProps from vitrage.common.constants import VertexProperties as VProps @@ -91,6 +90,14 @@ class GraphTest(GraphTestBase): self.assertEqual(updated_v[VProps.CATEGORY], v[VProps.CATEGORY], 'Graph item should change after update') + # Update the graph item and see changes take place + updated_v['KUKU'] = None + g.update_vertex(updated_v) + # Get it again + v = g.get_vertex(v_node.vertex_id) + self.assertFalse('KUKU' in v.properties, + 'Update value to None should entirly remove the key') + # check metadata another_vertex = utils.create_vertex( vitrage_id='123', entity_id='456', entity_category=INSTANCE, @@ -170,6 +177,14 @@ class GraphTest(GraphTestBase): e[EProps.UPDATE_TIMESTAMP], 'Graph item should change after update') + # Update the graph item and see changes take place + updated_e[EProps.IS_DELETED] = None + g.update_edge(updated_e) + # Get it again + e = g.get_edge(v_node.vertex_id, v_host.vertex_id, label) + self.assertFalse(EProps.IS_DELETED in e.properties, + 'Update value to None should entirly remove the key') + # check metadata another_label = 'ANOTHER_LABEL' another_edge = utils.create_edge(