Support parameterized queries for openCypher (#496)

* Support parameterized queries for openCypher

* update changelog

---------

Co-authored-by: Michael Chin <[email protected]>

Author: michaelnchin

ChangeLog.md

@@ -7,6 +7,7 @@ Starting with v1.31.6, this file will contain a record of major features and upd
   - Path: 03-Sample-Applications > 05-Healthcare-and-Life-Sciences-Graphs
 - Added openCypher and local file path support to `%seed` ([Link to PR](https://github.com/aws/graph-notebook/pull/292))
 - Added S3 support to `%seed` ([Link to PR](https://github.com/aws/graph-notebook/pull/488))
+- Added support for openCypher parameterized queries ([Link to PR](https://github.com/aws/graph-notebook/pull/496))
 - Added `%toggle_traceback` line magic ([Link to PR](https://github.com/aws/graph-notebook/pull/486))
 - Added support for setting `%graph_notebook_vis_options` from a variable ([Link to PR](https://github.com/aws/graph-notebook/pull/487))
 - Pinned JupyterLab<4.x to fix Python 3.8/3.10 builds ([Link to PR](https://github.com/aws/graph-notebook/pull/490))

src/graph_notebook/magics/graph_magic.py

@@ -2619,6 +2619,9 @@ def handle_opencypher_query(self, line, cell, local_ns):
         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('-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.')
         parser.add_argument('-g', '--group-by', type=str, default='~labels',
                             help='Property used to group nodes (e.g. code, ~id) default is ~labels')
         parser.add_argument('-gd', '--group-by-depth', action='store_true', default=False,
@@ -2660,6 +2663,20 @@ def handle_opencypher_query(self, line, cell, local_ns):
         res_format = None
         results_df = None
 
+        query_params = None
+        if args.query_parameters:
+            if args.query_parameters in local_ns:
+                query_params_input = local_ns[args.query_parameters]
+            else:
+                query_params_input = args.query_parameters
+            if isinstance(query_params_input, dict):
+                query_params = query_params_input
+            else:
+                try:
+                    query_params = json.loads(query_params_input.replace("'", '"'))
+                except Exception as e:
+                    print(f"Invalid query parameter input, ignoring.")
+
         if args.no_scroll:
             oc_layout = UNRESTRICTED_LAYOUT
             oc_scrollY = True
@@ -2680,7 +2697,7 @@ def handle_opencypher_query(self, line, cell, local_ns):
 
         if args.mode == 'explain':
             query_start = time.time() * 1000  # time.time() returns time in seconds w/high precision; x1000 to get in ms
-            res = self.client.opencypher_http(cell, explain=args.explain_type)
+            res = self.client.opencypher_http(cell, explain=args.explain_type, query_params=query_params)
             query_time = time.time() * 1000 - query_start
             explain = res.content.decode("utf-8")
             res.raise_for_status()
@@ -2696,7 +2713,7 @@ def handle_opencypher_query(self, line, cell, local_ns):
                                                                     link=f"data:text/html;base64,{base64_str}")
         elif args.mode == 'query':
             query_start = time.time() * 1000  # time.time() returns time in seconds w/high precision; x1000 to get in ms
-            oc_http = self.client.opencypher_http(cell)
+            oc_http = self.client.opencypher_http(cell, query_params=query_params)
             query_time = time.time() * 1000 - query_start
             oc_http.raise_for_status()
 
@@ -2754,7 +2771,10 @@ def handle_opencypher_query(self, line, cell, local_ns):
         elif args.mode == 'bolt':
             res_format = 'bolt'
             query_start = time.time() * 1000
-            res = self.client.opencyper_bolt(cell)
+            if query_params:
+                res = self.client.opencyper_bolt(cell, **query_params)
+            else:
+                res = self.client.opencyper_bolt(cell)
             query_time = time.time() * 1000 - query_start
             if not args.silent:
                 oc_metadata = build_opencypher_metadata_from_query(query_type='bolt', results=res,

src/graph_notebook/neptune/client.py

@@ -327,7 +327,8 @@ def _gremlin_query_plan(self, query: str, plan_type: str, args: dict, ) -> reque
         res = self._http_session.send(req, verify=self.ssl_verify)
         return res
 
-    def opencypher_http(self, query: str, headers: dict = None, explain: str = None) -> requests.Response:
+    def opencypher_http(self, query: str, headers: dict = None, explain: str = None,
+                        query_params: dict = None) -> requests.Response:
         if headers is None:
             headers = {}
 
@@ -343,6 +344,8 @@ def opencypher_http(self, query: str, headers: dict = None, explain: str = None)
             if explain:
                 data['explain'] = explain
                 headers['Accept'] = "text/html"
+            if query_params:
+                data['parameters'] = str(query_params).replace("'", '"')  # '{"AUS_code":"AUS","WLG_code":"WLG"}'
         else:
             url += 'db/neo4j/tx/commit'
             headers['content-type'] = 'application/json'

test/integration/iam/notebook/test_open_cypher_graph_notebook.py

@@ -48,6 +48,66 @@ def test_opencypher_bolt(self):
         assert len(res) == 1
         assert 'b' in res[0]
 
+    @pytest.mark.jupyter
+    @pytest.mark.opencypher
+    def test_opencypher_query_parameterized_with_var_input(self):
+        expected_league_name = "English Premier League"
+        query = 'MATCH (l:League {nickname: $LEAGUE_NICKNAME}) RETURN l.name'
+
+        store_to_var = 'res'
+        self.ip.user_ns['params_var'] = {'LEAGUE_NICKNAME': 'EPL'}
+        cell = f'''%%oc --query-parameters params_var --store-to {store_to_var}
+                        {query}'''
+        self.ip.run_cell(cell)
+        res = self.ip.user_ns[store_to_var]
+
+        assert len(res['results']) == 1
+        assert expected_league_name == res['results'][0]['l.name']
+
+    @pytest.mark.jupyter
+    @pytest.mark.opencypher
+    def test_opencypher_query_parameterized_with_str_input(self):
+        expected_league_name = "English Premier League"
+        query = 'MATCH (l:League {nickname: $LEAGUE_NICKNAME}) RETURN l.name'
+
+        store_to_var = 'res'
+        params_str = '{"LEAGUE_NICKNAME":"EPL"}'
+        cell = f'''%%oc --query-parameters {params_str} --store-to {store_to_var}
+                {query}'''
+        self.ip.run_cell(cell)
+        res = self.ip.user_ns[store_to_var]
+
+        assert len(res['results']) == 1
+        assert expected_league_name == res['results'][0]['l.name']
+
+    @pytest.mark.jupyter
+    @pytest.mark.opencypher
+    def test_opencypher_query_parameterized_invalid(self):
+        query = 'MATCH (l:League {nickname: $LEAGUE_NICKNAME}) RETURN l.name'
+
+        self.ip.user_ns['params_var'] = ['LEAGUE_NICKNAME']
+        cell = f'''%%oc --query-parameters params_var 
+                    {query}'''
+        self.ip.run_cell(cell)
+        self.assertTrue('graph_notebook_error' in self.ip.user_ns)
+
+    @pytest.mark.jupyter
+    @pytest.mark.opencypher
+    def test_opencypher_bolt_parameterized(self):
+        expected_league_name = "English Premier League"
+        query = 'MATCH (l:League {nickname: $LEAGUE_NICKNAME}) RETURN l.name'
+
+        store_to_var = 'res'
+        params_var = '{"LEAGUE_NICKNAME":"EPL"}'
+        cell = f'''%%oc bolt --query-parameters {params_var} --store-to {store_to_var}
+                    {query}'''
+        self.ip.run_cell(cell)
+        self.assertFalse('graph_notebook_error' in self.ip.user_ns)
+        res = self.ip.user_ns[store_to_var]
+
+        assert len(res) == 1
+        assert expected_league_name == res[0]['l.name']
+
     @pytest.mark.jupyter
     def test_load_opencypher_config(self):
         config = '''{