Add --explain-type option to %%gremlin (#503)

* Add --explain-type option to %%gremlin

* Update changelog

---------

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

Author: michaelnchin

ChangeLog.md

@@ -3,6 +3,7 @@
 Starting with v1.31.6, this file will contain a record of major features and updates made in each release of graph-notebook.
 
 ## Upcoming
+- Added `--explain-type` option to `%%gremlin` ([Link to PR](https://github.com/aws/graph-notebook/pull/503))
 
 ## Release 3.8.2 (June 5, 2023)
 - New Sample Applications - Healthcare and Life Sciences notebooks ([Link to PR](https://github.com/aws/graph-notebook/pull/484))

src/graph_notebook/magics/graph_magic.py

@@ -44,7 +44,8 @@
     LOAD_JOB_MODES, MODE_AUTO, FINAL_LOAD_STATUSES, SPARQL_ACTION, FORMAT_CSV, FORMAT_OPENCYPHER, FORMAT_NTRIPLE, \
     FORMAT_NQUADS, FORMAT_RDFXML, FORMAT_TURTLE, STREAM_RDF, STREAM_PG, STREAM_ENDPOINTS, \
     NEPTUNE_CONFIG_HOST_IDENTIFIERS, is_allowed_neptune_host, \
-    STATISTICS_LANGUAGE_INPUTS, STATISTICS_MODES, SUMMARY_MODES
+    STATISTICS_LANGUAGE_INPUTS, STATISTICS_MODES, SUMMARY_MODES, \
+    SPARQL_EXPLAIN_MODES, OPENCYPHER_EXPLAIN_MODES
 from graph_notebook.network import SPARQLNetwork
 from graph_notebook.network.gremlin.GremlinNetwork import parse_pattern_list_str, GremlinNetwork
 from graph_notebook.visualization.rows_and_columns import sparql_get_rows_and_columns, opencypher_get_rows_and_columns
@@ -516,9 +517,9 @@ def sparql(self, line='', cell='', local_ns: dict = None):
                             help='prefix path to sparql endpoint. For example, if "foo/bar" were specified, '
                                  'the endpoint called would be host:port/foo/bar')
         parser.add_argument('--expand-all', action='store_true')
-        parser.add_argument('--explain-type', default='dynamic',
-                            help='explain mode to use when using the explain query mode',
-                            choices=['dynamic', 'static', 'details'])
+        parser.add_argument('--explain-type', type=str.lower, default='dynamic',
+                            help=f'Explain mode to use when using the explain query mode. '
+                                 f'Expected values: ${SPARQL_EXPLAIN_MODES}')
         parser.add_argument('--explain-format', default='text/html', help='response format for explain query mode',
                             choices=['text/csv', 'text/html'])
         parser.add_argument('-m', '--media-type', type=str, default='',
@@ -804,6 +805,8 @@ def gremlin(self, line, cell, local_ns: dict = None):
         parser = argparse.ArgumentParser()
         parser.add_argument('query_mode', nargs='?', default='query',
                             help='query mode (default=query) [query|explain|profile]')
+        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',
                             help='Property used to group nodes (e.g. code, T.region) default is T.label')
@@ -881,7 +884,8 @@ def gremlin(self, line, cell, local_ns: dict = None):
         transport_args = {'max_content_length': args.max_content_length}
 
         if mode == QueryMode.EXPLAIN:
-            res = self.client.gremlin_explain(cell)
+            res = self.client.gremlin_explain(cell,
+                                              args={'explain.mode': args.explain_type} if args.explain_type else {})
             res.raise_for_status()
             # Replace strikethrough character bytes, can't be encoded to ASCII
             explain_bytes = res.content.replace(b'\xcc', b'-')
@@ -2685,9 +2689,9 @@ def handle_opencypher_query(self, line, cell, local_ns):
         This method in its own handler so that the magics %%opencypher and %%oc can both call it
         """
         parser = argparse.ArgumentParser()
-        parser.add_argument('--explain-type', default='dynamic',
-                            help='explain mode to use when using the explain query mode',
-                            choices=['dynamic', 'static', 'details', 'debug'])
+        parser.add_argument('--explain-type', type=str.lower, default='dynamic',
+                            help=f'Explain mode to use when using the explain query mode. '
+                                 f'Accepted values: ${OPENCYPHER_EXPLAIN_MODES}')
         parser.add_argument('-qp', '--query-parameters', type=str, default='',
                             help='Parameter definitions to apply to the query. This option can accept a local variable '
                                  'name, or a string representation of the map.')

src/graph_notebook/neptune/client.py

@@ -114,6 +114,9 @@
 SUMMARY_MODES = ["", "basic", "detailed"]
 STATISTICS_LANGUAGE_INPUTS = ["propertygraph", "pg", "gremlin", "oc", "opencypher", "sparql", "rdf"]
 
+SPARQL_EXPLAIN_MODES = ['dynamic', 'static', 'details']
+OPENCYPHER_EXPLAIN_MODES = ['dynamic', 'static', 'details']
+
 
 def is_allowed_neptune_host(hostname: str, host_allowlist: list):
     for host_snippet in host_allowlist:
@@ -211,12 +214,8 @@ def do_sparql_request(self, data: dict, headers=None, explain: str = '', path: s
         if 'content-type' not in headers:
             headers['content-type'] = DEFAULT_SPARQL_CONTENT_TYPE
 
-        explain = explain.lower()
         if explain != '':
-            if explain not in ['static', 'dynamic', 'details']:
-                raise ValueError('explain mode not valid, must be one of "static", "dynamic", or "details"')
-            else:
-                data['explain'] = explain
+            data['explain'] = explain
 
         if path != '':
             sparql_path = f'/{path}'