Add support for Gremlin proxy host and Neptune HTTP query visualization (#530)

* Add Gremlin proxy host fix and Neptune HTTP query support

* update changelog

Author: michaelnchin

ChangeLog.md

@@ -5,6 +5,7 @@ Starting with v1.31.6, this file will contain a record of major features and upd
 ## Upcoming
 - Added `--explain-type` option to `%%gremlin` ([Link to PR](https://github.com/aws/graph-notebook/pull/503))
 - Added general documentation for `%%graph_notebook_config` options ([Link to PR](https://github.com/aws/graph-notebook/pull/504))
+- Added support for Gremlin proxy hosts and visualization of Neptune HTTP results ([Link to PR](https://github.com/aws/graph-notebook/pull/530))
 - Modified Dockerfile to support Python 3.10 ([Link to PR](https://github.com/aws/graph-notebook/pull/519))
 - Updated Docker documentation with platform-specific run commands ([Link to PR](https://github.com/aws/graph-notebook/pull/502))
 - Fixed deprecation warnings in GitHub workflows ([Link to PR](https://github.com/aws/graph-notebook/pull/506))

src/graph_notebook/magics/graph_magic.py

@@ -29,6 +29,7 @@
 from SPARQLWrapper import SPARQLWrapper
 from botocore.session import get_session
 from gremlin_python.driver.protocol import GremlinServerError
+from gremlin_python.structure.graph import Path
 from IPython.core.display import HTML, display_html, display
 from IPython.core.magic import (Magics, magics_class, cell_magic, line_magic, line_cell_magic, needs_local_scope)
 from ipywidgets.widgets.widget_description import DescriptionStyle
@@ -809,15 +810,15 @@ def gremlin(self, line, cell, local_ns: dict = None):
         parser.add_argument('--explain-type', type=str.lower, default='',
                             help='Explain mode to use when using the explain query mode.')
         parser.add_argument('-p', '--path-pattern', default='', help='path pattern')
-        parser.add_argument('-g', '--group-by', type=str, default='T.label',
+        parser.add_argument('-g', '--group-by', type=str, default='',
                             help='Property used to group nodes (e.g. code, T.region) default is T.label')
         parser.add_argument('-gd', '--group-by-depth', action='store_true', default=False,
                             help="Group nodes based on path hierarchy")
         parser.add_argument('-gr', '--group-by-raw', action='store_true', default=False,
                             help="Group nodes by the raw result")
-        parser.add_argument('-d', '--display-property', type=str, default='T.label',
+        parser.add_argument('-d', '--display-property', type=str, default='',
                             help='Property to display the value of on each node, default is T.label')
-        parser.add_argument('-de', '--edge-display-property', type=str, default='T.label',
+        parser.add_argument('-de', '--edge-display-property', type=str, default='',
                             help='Property to display the value of on each edge, default is T.label')
         parser.add_argument('-t', '--tooltip-property', type=str, default='',
                             help='Property to display the value of on each node tooltip. If not specified, tooltip '
@@ -937,8 +938,16 @@ def gremlin(self, line, cell, local_ns: dict = None):
                 else:
                     first_tab_html = pre_container_template.render(content='No profile found')
         else:
+            using_http = False
             query_start = time.time() * 1000  # time.time() returns time in seconds w/high precision; x1000 to get in ms
-            query_res = self.client.gremlin_query(cell, transport_args=transport_args)
+            if self.graph_notebook_config.proxy_host != '' and self.client.is_neptune_domain():
+                using_http = True
+                query_res_http = self.client.gremlin_http_query(cell, headers={'Accept': 'application/vnd.gremlin-v1.0+json;types=false'})
+                query_res_http.raise_for_status()
+                query_res_http_json = query_res_http.json()
+                query_res = query_res_http_json['result']['data']
+            else:
+                query_res = self.client.gremlin_query(cell, transport_args=transport_args)
             query_time = time.time() * 1000 - query_start
             if not args.silent:
                 gremlin_metadata = build_gremlin_metadata_from_query(query_type='query', results=query_res,
@@ -952,18 +961,30 @@ def gremlin(self, line, cell, local_ns: dict = None):
                     logger.debug(f'edge_display_property: {args.edge_display_property}')
                     logger.debug(f'label_max_length: {args.label_max_length}')
                     logger.debug(f'ignore_groups: {args.ignore_groups}')
-                    gn = GremlinNetwork(group_by_property=args.group_by, display_property=args.display_property,
+                    gn = GremlinNetwork(group_by_property=args.group_by,
+                                        display_property=args.display_property,
                                         group_by_raw=args.group_by_raw,
                                         group_by_depth=args.group_by_depth,
                                         edge_display_property=args.edge_display_property,
                                         tooltip_property=args.tooltip_property,
                                         edge_tooltip_property=args.edge_tooltip_property,
                                         label_max_length=args.label_max_length,
                                         edge_label_max_length=args.edge_label_max_length,
-                                        ignore_groups=args.ignore_groups)
+                                        ignore_groups=args.ignore_groups,
+                                        using_http=using_http)
+
+                    if using_http and 'path()' in cell and query_res:
+                        first_path = query_res[0]
+                        if isinstance(first_path, dict) and first_path.keys() == {'labels', 'objects'}:
+                            query_res_to_path_type = []
+                            for path in query_res:
+                                new_path_list = path['objects']
+                                new_path = Path(labels=[], objects=new_path_list)
+                                query_res_to_path_type.append(new_path)
+                            query_res = query_res_to_path_type
 
                     if args.path_pattern == '':
-                        gn.add_results(query_res)
+                        gn.add_results(query_res, is_http=using_http)
                     else:
                         pattern = parse_pattern_list_str(args.path_pattern)
                         gn.add_results_with_pattern(query_res, pattern)

src/graph_notebook/neptune/client.py

@@ -17,14 +17,14 @@
 from botocore.awsrequest import AWSRequest
 from gremlin_python.driver import client, serializer
 from gremlin_python.driver.protocol import GremlinServerError
+from gremlin_python.driver.aiohttp.transport import AiohttpTransport
 from neo4j import GraphDatabase, DEFAULT_DATABASE
 from neo4j.exceptions import AuthError
 from base64 import b64encode
 import nest_asyncio
 
 from graph_notebook.neptune.bolt_auth_token import NeptuneBoltAuthToken
 
-
 # This patch is no longer needed when graph_notebook is using the a Gremlin Python
 # client >= 3.5.0 as the HashableDict is now part of that client driver.
 # import graph_notebook.neptune.gremlin.graphsonV3d0_MapType_objectify_patch  # noqa F401
@@ -45,7 +45,7 @@
 # TODO: add doc links to each command
 
 FORMAT_CSV = 'csv'
-FORMAT_OPENCYPHER='opencypher'
+FORMAT_OPENCYPHER = 'opencypher'
 FORMAT_NTRIPLE = 'ntriples'
 FORMAT_NQUADS = 'nquads'
 FORMAT_RDFXML = 'rdfxml'
@@ -191,11 +191,19 @@ def is_neptune_domain(self):
         return is_allowed_neptune_host(hostname=self.target_host, host_allowlist=self.neptune_hosts)
 
     def get_uri_with_port(self, use_websocket=False, use_proxy=False):
-        protocol = self._http_protocol
         if use_websocket is True:
             protocol = self._ws_protocol
+        else:
+            protocol = self._http_protocol
 
-        uri = f'{protocol}://{self.host}:{self.port}'
+        if use_proxy is True:
+            uri_host = self.proxy_host
+            uri_port = self.proxy_port
+        else:
+            uri_host = self.target_host
+            uri_port = self.target_port
+
+        uri = f'{protocol}://{uri_host}:{uri_port}'
         return uri
 
     def sparql_query(self, query: str, headers=None, explain: str = '', path: str = '') -> requests.Response:
@@ -267,11 +275,20 @@ def sparql_cancel(self, query_id: str, silent: bool = False):
     def get_gremlin_connection(self, transport_kwargs) -> client.Client:
         nest_asyncio.apply()
 
-        ws_url = f'{self.get_uri_with_port(use_websocket=True)}/gremlin'
-        request = self._prepare_request('GET', ws_url)
+        ws_url = f'{self.get_uri_with_port(use_websocket=True, use_proxy=False)}/gremlin'
+        if self.proxy_host != '':
+            proxy_http_url = f'{self.get_uri_with_port(use_websocket=False, use_proxy=True)}/gremlin'
+            transport_factory_args = lambda: AiohttpTransport(call_from_event_loop=True, proxy=proxy_http_url,
+                                                              **transport_kwargs)
+            request = self._prepare_request('GET', proxy_http_url)
+        else:
+            transport_factory_args = lambda: AiohttpTransport(**transport_kwargs)
+            request = self._prepare_request('GET', ws_url)
+
         traversal_source = 'g' if self.is_neptune_domain() else self.gremlin_traversal_source
-        return client.Client(ws_url, traversal_source, username=self.gremlin_username,
-                             password=self.gremlin_password, message_serializer=self.gremlin_serializer,
+        return client.Client(ws_url, traversal_source, transport_factory=transport_factory_args,
+                             username=self.gremlin_username, password=self.gremlin_password,
+                             message_serializer=self.gremlin_serializer,
                              headers=dict(request.headers), **transport_kwargs)
 
     def gremlin_query(self, query, transport_args=None, bindings=None):
@@ -298,7 +315,8 @@ def gremlin_http_query(self, query, headers=None) -> requests.Response:
         if headers is None:
             headers = {}
 
-        uri = f'{self.get_uri_with_port()}/gremlin'
+        use_proxy = True if self.proxy_host != '' else False
+        uri = f'{self.get_uri_with_port(use_websocket=False, use_proxy=use_proxy)}/gremlin'
         data = {'gremlin': query}
         req = self._prepare_request('POST', uri, data=json.dumps(data), headers=headers)
         res = self._http_session.send(req, verify=self.ssl_verify)
@@ -431,7 +449,7 @@ def stream(self, url, **kwargs) -> requests.Response:
         params = {}
         for k, v in kwargs.items():
             params[k] = v
-        req = self._prepare_request('GET', url, params=params,data='')
+        req = self._prepare_request('GET', url, params=params, data='')
         res = self._http_session.send(req, verify=self.ssl_verify)
         return res.json()
 
@@ -850,7 +868,7 @@ def with_sparql_path(self, path: str):
     def with_gremlin_traversal_source(self, traversal_source: str):
         self.args['gremlin_traversal_source'] = traversal_source
         return ClientBuilder(self.args)
-        
+
     def with_gremlin_login(self, username: str, password: str):
         self.args['gremlin_username'] = username
         self.args['gremlin_password'] = password
@@ -859,7 +877,7 @@ def with_gremlin_login(self, username: str, password: str):
     def with_gremlin_serializer(self, message_serializer: str):
         self.args['gremlin_serializer'] = message_serializer
         return ClientBuilder(self.args)
-    
+
     def with_neo4j_login(self, username: str, password: str, auth: bool, database: str):
         self.args['neo4j_username'] = username
         self.args['neo4j_password'] = password