Fix Gremlin visualizations failing to render with UUID type IDs (#475)

* Fix Gremlin visualizations with UUID type IDs

* remove dev comment

* update changelog

---------

Co-authored-by: Michael Chin <chnmch@amazon.com>

Author: michaelnchin

ChangeLog.md

@@ -4,10 +4,11 @@ Starting with v1.31.6, this file will contain a record of major features and upd
 
 ## Upcoming
 - Added support for Python 3.10 ([Link to PR](https://github.com/aws/graph-notebook/pull/476))
+- Deprecated Python 3.7 support ([PR #1](https://github.com/aws/graph-notebook/pull/453)) ([PR #2](https://github.com/aws/graph-notebook/pull/473))
 - Fixed Dockerfile builds breaking with AL2023 ([Link to PR](https://github.com/aws/graph-notebook/pull/466))
 - Fixed `--store-to` option for several magics ([Link to PR](https://github.com/aws/graph-notebook/pull/463))
 - Fixed broken documentation links in Neptune ML notebooks ([PR #1](https://github.com/aws/graph-notebook/pull/467)) ([PR #2](https://github.com/aws/graph-notebook/pull/468))
-- Deprecated Python 3.7 support ([PR #1](https://github.com/aws/graph-notebook/pull/453)) ([PR #2](https://github.com/aws/graph-notebook/pull/473))
+- Fixed Gremlin graph tab not rendering with UUID type IDs ([Link to PR](https://github.com/aws/graph-notebook/pull/475))
 
 ## Release 3.7.3 (March 14, 2023)
 - Fixed detailed mode output for graph summary requests ([Link to PR](https://github.com/aws/graph-notebook/pull/461))

src/graph_notebook/magics/graph_magic.py

@@ -936,6 +936,7 @@ def gremlin(self, line, cell, local_ns: dict = None):
                                                                      query_time=query_time)
                 titles.append('Console')
 
+                gremlin_network = None
                 try:
                     logger.debug(f'groupby: {args.group_by}')
                     logger.debug(f'display_property: {args.display_property}')
@@ -957,18 +958,24 @@ def gremlin(self, line, cell, local_ns: dict = None):
                     else:
                         pattern = parse_pattern_list_str(args.path_pattern)
                         gn.add_results_with_pattern(query_res, pattern)
+                    gremlin_network = gn
                     logger.debug(f'number of nodes is {len(gn.graph.nodes)}')
-                    if len(gn.graph.nodes) > 0:
+                except ValueError as value_error:
+                    logger.debug(
+                        f'Unable to create graph network from result due to error: {value_error}. '
+                        f'Skipping from result set.')
+                if gremlin_network and len(gremlin_network.graph.nodes) > 0:
+                    try:
                         self.graph_notebook_vis_options['physics']['disablePhysicsAfterInitialSimulation'] \
                             = args.stop_physics
                         self.graph_notebook_vis_options['physics']['simulationDuration'] = args.simulation_duration
-                        f = Force(network=gn, options=self.graph_notebook_vis_options)
+                        f = Force(network=gremlin_network, options=self.graph_notebook_vis_options)
                         titles.append('Graph')
                         children.append(f)
                         logger.debug('added gremlin network to tabs')
-                except ValueError as value_error:
-                    logger.debug(
-                        f'unable to create gremlin network from result. Skipping from result set: {value_error}')
+                    except Exception as force_error:
+                        logger.debug(
+                            f'Unable to render visualization from graph network due to error: {force_error}. Skipping.')
 
                 # Check if we can access the CDNs required by itables library.
                 # If not, then render our own HTML template.

src/graph_notebook/network/gremlin/GremlinNetwork.py

@@ -78,16 +78,21 @@ def get_id(element):
     extract id from a given element for use in the GremlinNetwork
     """
     if isinstance(element, Vertex):
-        return str(element.id)
+        element_id = str(element.id)
     elif isinstance(element, Edge):
-        return str(element.label)
+        element_id = str(element.label)
     elif isinstance(element, dict):
         if T.id in element:
-            return element[T.id]
+            element_id = element[T.id]
         else:
-            return generate_id_from_dict(element)
+            element_id = generate_id_from_dict(element)
     else:
-        return str(element)
+        element_id = str(element)
+
+    if isinstance(element_id, uuid.UUID):
+        element_id = str(element_id)
+
+    return element_id
 
 
 class GremlinNetwork(EventfulNetwork):
@@ -342,7 +347,10 @@ def add_vertex(self, v, path_index: int = -1):
         depth_group = "__DEPTH-" + str(path_index) + "__"
         node_id = ''
         if type(v) is Vertex:
-            node_id = v.id
+            if isinstance(v.id, uuid.UUID):
+                node_id = str(v.id)
+            else:
+                node_id = v.id
             title = v.label
             label_full = ''
             tooltip_display_is_set = True
@@ -447,6 +455,8 @@ def add_vertex(self, v, path_index: int = -1):
                     properties[k] = copy_val
                 elif isinstance(v[k], Decimal):
                     properties[k] = float(v[k])
+                elif isinstance(v[k], uuid.UUID):
+                    properties[k] = str(v[k])
                 else:
                     properties[k] = v[k]
 
@@ -529,15 +539,20 @@ def add_path_edge(self, edge, from_id='', to_id='', data=None):
 
         if type(edge) is Edge:
             from_id = from_id if from_id != '' else edge.outV.id
+            if isinstance(from_id, uuid.UUID):
+                from_id = str(from_id)
             to_id = to_id if to_id != '' else edge.inV.id
+            if isinstance(to_id, uuid.UUID):
+                to_id = str(to_id)
+            edge_id = str(edge.id) if isinstance(edge.id, uuid.UUID) else edge.id
             edge_label_full = ''
             using_custom_tooltip = False
             tooltip_display_is_set = True
             edge_title = edge.label
             if self.edge_tooltip_property and self.edge_tooltip_property != self.edge_display_property:
                 using_custom_tooltip = True
                 tooltip_display_is_set = False
-            data['properties'] = {'id': edge.id, 'label': edge.label, 'outV': str(edge.outV), 'inV': str(edge.inV)}
+            data['properties'] = {'id': edge_id, 'label': edge.label, 'outV': str(edge.outV), 'inV': str(edge.inV)}
 
             edge_label = edge_title if len(edge_title) <= self.edge_label_max_length \
                 else edge_title[:self.edge_label_max_length - 3] + '...'
@@ -564,7 +579,7 @@ def add_path_edge(self, edge, from_id='', to_id='', data=None):
             self.get_explicit_edge_property_value(data, edge, self.edge_tooltip_property)
 
             data['title'] = edge_title
-            self.add_edge(from_id=from_id, to_id=to_id, edge_id=edge.id, label=edge_label, title=edge_title,
+            self.add_edge(from_id=from_id, to_id=to_id, edge_id=edge_id, label=edge_label, title=edge_title,
                           data=data)
         elif type(edge) is dict:
             properties = {}
@@ -592,6 +607,8 @@ def add_path_edge(self, edge, from_id='', to_id='', data=None):
                     properties[k] = get_id(edge[k])
                 elif isinstance(edge[k], Decimal):
                     properties[k] = float(edge[k])
+                elif isinstance(edge[k], uuid.UUID):
+                    properties[k] = str(edge[k])
                 else:
                     properties[k] = edge[k]
 

test/unit/network/gremlin/test_gremlin_network.py

@@ -5,6 +5,7 @@
 
 import unittest
 from decimal import Decimal
+from uuid import uuid4
 from gremlin_python.structure.graph import Path, Edge, Vertex
 from gremlin_python.process.traversal import T, Direction
 from graph_notebook.network.EventfulNetwork import EVENT_ADD_NODE
@@ -191,6 +192,75 @@ def test_add_explicit_type_vertex_with_valid_label_and_invalid_tooltip(self):
         self.assertEqual(node1['label'], '1')
         self.assertEqual(node1['title'], '1')
 
+    def test_add_explicit_type_vertex_with_string_id(self):
+        v_id = 'a_id'
+        vertex = Vertex(id=v_id)
+
+        gn = GremlinNetwork()
+        gn.add_vertex(vertex)
+        node = gn.graph.nodes.get(v_id)
+        self.assertIsNotNone(node)
+        self.assertEqual(node['properties']['id'], v_id)
+
+    def test_add_explicit_type_vertex_with_uuid_id(self):
+        v_id = uuid4()
+        vertex = Vertex(id=v_id)
+
+        gn = GremlinNetwork()
+        gn.add_vertex(vertex)
+        node = gn.graph.nodes.get(str(v_id))
+        self.assertIsNotNone(node)
+        self.assertEqual(node['properties']['id'], str(v_id))
+
+    def test_add_explicit_type_vertex_with_integer_id(self):
+        v_id = 1
+        vertex = Vertex(id=v_id)
+
+        gn = GremlinNetwork()
+        gn.add_vertex(vertex)
+        node = gn.graph.nodes.get(v_id)
+        self.assertIsNotNone(node)
+        self.assertEqual(node['properties']['id'], v_id)
+
+    def test_add_vertex_with_string_id(self):
+        v_id = 'a_id'
+        vertex = {
+            T.id: v_id,
+            T.label: 'label'
+        }
+
+        gn = GremlinNetwork()
+        gn.add_vertex(vertex)
+        node = gn.graph.nodes.get(v_id)
+        self.assertIsNotNone(node)
+        self.assertEqual(node['properties'][T.id], v_id)
+
+    def test_add_vertex_with_uuid_id(self):
+        v_id = uuid4()
+        vertex = {
+            T.id: v_id,
+            T.label: 'label'
+        }
+
+        gn = GremlinNetwork()
+        gn.add_vertex(vertex)
+        node = gn.graph.nodes.get(str(v_id))
+        self.assertIsNotNone(node)
+        self.assertEqual(node['properties'][T.id], str(v_id))
+
+    def test_add_vertex_with_integer_id(self):
+        v_id = 99
+        vertex = {
+            T.id: v_id,
+            T.label: 'label'
+        }
+
+        gn = GremlinNetwork()
+        gn.add_vertex(vertex)
+        node = gn.graph.nodes.get(str(v_id))
+        self.assertIsNotNone(node)
+        self.assertEqual(node['properties'][T.id], v_id)
+
     def test_add_vertex_without_node_property(self):
         vertex = {
             T.id: '1234',
@@ -753,6 +823,51 @@ def test_add_vertex_with_Decimal_type_property_in_list(self):
         self.assertIsInstance(final_lon_value, float)
         self.assertIsInstance(final_lat_value, float)
 
+    def test_add_explicit_type_single_edge_with_string_id(self):
+        vertex1 = Vertex(id='1')
+        vertex2 = Vertex(id='2')
+        e_id = '1'
+
+        edge1 = Edge(id=e_id, outV=vertex1, inV=vertex2, label='route')
+
+        gn = GremlinNetwork()
+        gn.add_vertex(vertex1)
+        gn.add_vertex(vertex2)
+        gn.add_path_edge(edge1)
+        edge = gn.graph.get_edge_data('1', '2')
+        self.assertIsNotNone(edge)
+        self.assertEqual(edge[e_id]['properties']['id'], e_id)
+
+    def test_add_explicit_type_single_edge_with_uuid_id(self):
+        vertex1 = Vertex(id='1')
+        vertex2 = Vertex(id='2')
+        e_id = uuid4()
+
+        edge1 = Edge(id=e_id, outV=vertex1, inV=vertex2, label='route')
+
+        gn = GremlinNetwork()
+        gn.add_vertex(vertex1)
+        gn.add_vertex(vertex2)
+        gn.add_path_edge(edge1)
+        edge = gn.graph.get_edge_data('1', '2')
+        self.assertIsNotNone(edge)
+        self.assertEqual(edge[str(e_id)]['properties']['id'], str(e_id))
+
+    def test_add_explicit_type_single_edge_with_integer_id(self):
+        vertex1 = Vertex(id='1')
+        vertex2 = Vertex(id='2')
+        e_id = 1
+
+        edge1 = Edge(id=e_id, outV=vertex1, inV=vertex2, label='route')
+
+        gn = GremlinNetwork()
+        gn.add_vertex(vertex1)
+        gn.add_vertex(vertex2)
+        gn.add_path_edge(edge1)
+        edge = gn.graph.get_edge_data('1', '2')
+        self.assertIsNotNone(edge)
+        self.assertEqual(edge[e_id]['properties']['id'], e_id)
+
     def test_add_explicit_type_single_edge_without_edge_property(self):
         vertex1 = Vertex(id='1')
         vertex2 = Vertex(id='2')
@@ -981,6 +1096,51 @@ def test_add_explicit_type_single_edge_with_valid_label_and_invalid_tooltip(self
         self.assertEqual(edge['1']['label'], 'v[2]')
         self.assertEqual(edge['1']['title'], 'v[2]')
 
+    def test_add_single_edge_with_string_id(self):
+        vertex1 = Vertex(id='1')
+        vertex2 = Vertex(id='2')
+        e_id = '1'
+
+        edge1 = {T.id: e_id, T.label: 'route', 'outV': 'v[1]', 'inV': 'v[2]'}
+
+        gn = GremlinNetwork()
+        gn.add_vertex(vertex1)
+        gn.add_vertex(vertex2)
+        gn.add_path_edge(edge1, from_id='1', to_id='2')
+        edge = gn.graph.get_edge_data('1', '2')
+        self.assertIsNotNone(edge)
+        self.assertEqual(edge[e_id]['properties'][T.id], e_id)
+
+    def test_add_single_edge_with_uuid_id(self):
+        vertex1 = Vertex(id='1')
+        vertex2 = Vertex(id='2')
+        e_id = uuid4()
+
+        edge1 = {T.id: e_id, T.label: 'route', 'outV': 'v[1]', 'inV': 'v[2]'}
+
+        gn = GremlinNetwork()
+        gn.add_vertex(vertex1)
+        gn.add_vertex(vertex2)
+        gn.add_path_edge(edge1, from_id='1', to_id='2')
+        edge = gn.graph.get_edge_data('1', '2')
+        self.assertIsNotNone(edge)
+        self.assertEqual(edge[str(e_id)]['properties'][T.id], str(e_id))
+
+    def test_add_single_edge_with_integer_id(self):
+        vertex1 = Vertex(id='1')
+        vertex2 = Vertex(id='2')
+        e_id = 1
+
+        edge1 = {T.id: e_id, T.label: 'route', 'outV': 'v[1]', 'inV': 'v[2]'}
+
+        gn = GremlinNetwork()
+        gn.add_vertex(vertex1)
+        gn.add_vertex(vertex2)
+        gn.add_path_edge(edge1, from_id='1', to_id='2')
+        edge = gn.graph.get_edge_data('1', '2')
+        self.assertIsNotNone(edge)
+        self.assertEqual(edge[str(e_id)]['properties'][T.id], e_id)
+
     def test_add_single_edge_without_edge_property(self):
         vertex1 = Vertex(id='1')
         vertex2 = Vertex(id='2')
@@ -2213,6 +2373,26 @@ def test_add_elementmap_edge(self):
         self.assertEqual(inv_data['properties'], edge_map[Direction.IN])
         self.assertEqual(outv_data['properties'], edge_map[Direction.OUT])
 
+    def test_add_elementmap_edge_with_direction_uuid_ids(self):
+        in_node_id = uuid4()
+        out_node_id = uuid4()
+        edge_map = {
+            T.id: '5298',
+            T.label: 'route',
+            Direction.IN: {T.id: in_node_id, T.label: 'airport'},
+            Direction.OUT: {T.id: out_node_id, T.label: 'airport'},
+            'dist': 763
+        }
+
+        gn = GremlinNetwork()
+        gn.insert_elementmap(edge_map, index=1)
+        inv_data = gn.graph.nodes.get(str(in_node_id))
+        outv_data = gn.graph.nodes.get(str(out_node_id))
+        self.assertIsNotNone(inv_data)
+        self.assertIsNotNone(outv_data)
+        self.assertEqual(inv_data['properties'][T.id], str(edge_map[Direction.IN][T.id]))
+        self.assertEqual(outv_data['properties'][T.id], str(edge_map[Direction.OUT][T.id]))
+
     def test_add_elementmap_edge_groupby_depth(self):
         edge_map = {
             T.id: '5298',