diff --git a/src/graph_notebook/magics/graph_magic.py b/src/graph_notebook/magics/graph_magic.py index 54c3b9d5..742de41e 100644 --- a/src/graph_notebook/magics/graph_magic.py +++ b/src/graph_notebook/magics/graph_magic.py @@ -65,6 +65,7 @@ GREMLIN_PROTOCOL_FORMATS, DEFAULT_HTTP_PROTOCOL, DEFAULT_WS_PROTOCOL, GRAPHSONV4_UNTYPED, \ GREMLIN_SERIALIZERS_WS, get_gremlin_serializer_mime, normalize_protocol_name, generate_snapshot_name) from graph_notebook.network import SPARQLNetwork +from graph_notebook.neptune.utils import serialize_query_params 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 from graph_notebook.visualization.template_retriever import retrieve_template @@ -1205,14 +1206,10 @@ def gremlin(self, line, cell, local_ns: dict = None): query_params_input = local_ns[args.query_parameters] else: query_params_input = args.query_parameters - if isinstance(query_params_input, dict): - query_params = json.dumps(query_params_input) - else: - try: - query_params_dict = json.loads(query_params_input.replace("'", '"')) - query_params = json.dumps(query_params_dict) - except Exception as e: - print(f"Invalid query parameter input, ignoring.") + query_params = serialize_query_params(query_params_input) + if query_params is None: + print(f"Invalid query parameter input, ignoring.") + query_params = None if args.no_scroll: gremlin_layout = UNRESTRICTED_LAYOUT @@ -3679,14 +3676,9 @@ def handle_opencypher_query(self, line, cell, local_ns, return_tabs=False): query_params_input = local_ns[args.query_parameters] else: query_params_input = args.query_parameters - if isinstance(query_params_input, dict): - query_params = json.dumps(query_params_input) - else: - try: - query_params_dict = json.loads(query_params_input.replace("'", '"')) - query_params = json.dumps(query_params_dict) - except Exception as e: - print(f"Invalid query parameter input, ignoring.") + query_params = serialize_query_params(query_params_input) + if query_params is None: + print(f"Invalid query parameter input, ignoring.") if args.no_scroll: oc_layout = UNRESTRICTED_LAYOUT diff --git a/src/graph_notebook/neptune/client.py b/src/graph_notebook/neptune/client.py index 35fa1602..cfea79cf 100644 --- a/src/graph_notebook/neptune/client.py +++ b/src/graph_notebook/neptune/client.py @@ -28,6 +28,7 @@ from networkx import is_valid_directed_joint_degree from graph_notebook.neptune.bolt_auth_token import NeptuneBoltAuthToken +from graph_notebook.neptune.utils import serialize_query_params # 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. @@ -489,7 +490,7 @@ def gremlin_http_query(self, query, headers=None, query_params: dict = None, data['language'] = 'gremlin' headers['content-type'] = 'application/json' if query_params: - data['parameters'] = str(query_params).replace("'", '"') + data['parameters'] = serialize_query_params(query_params) else: uri = f'{self.get_uri(use_websocket=False, use_proxy=use_proxy)}/gremlin' data['gremlin'] = query @@ -529,7 +530,7 @@ def _gremlin_query_plan(self, query: str, plan_type: str, args: dict, headers['content-type'] = 'application/json' if 'parameters' in args: query_params = args.pop('parameters') - data['parameters'] = str(query_params).replace("'", '"') + data['parameters'] = serialize_query_params(query_params) if plan_type == 'explain': # Remove explain.mode once HTTP is changed explain_mode = args.pop('explain.mode') @@ -586,7 +587,7 @@ def opencypher_http(self, query: str, headers: dict = None, explain: str = None, data['explain'] = explain headers['Accept'] = "text/html" if query_params: - data['parameters'] = str(query_params).replace("'", '"') # '{"AUS_code":"AUS","WLG_code":"WLG"}' + data['parameters'] = serialize_query_params(query_params) if query_timeout and self.is_analytics_domain(): data['queryTimeoutMilliseconds'] = str(query_timeout) else: diff --git a/src/graph_notebook/neptune/utils.py b/src/graph_notebook/neptune/utils.py new file mode 100644 index 00000000..6cb5d27f --- /dev/null +++ b/src/graph_notebook/neptune/utils.py @@ -0,0 +1,46 @@ +""" +Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +SPDX-License-Identifier: Apache-2.0 +""" + +import ast +import json +from functools import singledispatch + + +@singledispatch +def serialize_query_params(query_params): + """Serialize query parameters into a valid JSON string for Neptune. + + Dispatches based on input type: + - dict: Serializes directly via json.dumps. + - str: Validates as JSON and passes through, or parses as a Python + dict literal and converts to JSON. + + Args: + query_params: Query parameters as a dict, JSON string, or Python + dict literal string. + + Returns: + A valid JSON string, or None if the input cannot be parsed. + """ + return None + + +@serialize_query_params.register(dict) +def _(query_params): + """Serialize a dict to a JSON string.""" + return json.dumps(query_params) + + +@serialize_query_params.register(str) +def _(query_params): + """Validate a JSON string or parse a Python dict literal string to JSON.""" + try: + json.loads(query_params) + return query_params + except (json.JSONDecodeError, ValueError): + try: + return json.dumps(ast.literal_eval(query_params)) + except Exception: + return None diff --git a/test/unit/graph_magic/test_query_params_serialization.py b/test/unit/graph_magic/test_query_params_serialization.py new file mode 100644 index 00000000..7f6ee3f5 --- /dev/null +++ b/test/unit/graph_magic/test_query_params_serialization.py @@ -0,0 +1,63 @@ +""" +Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +SPDX-License-Identifier: Apache-2.0 +""" + +import json +import unittest + +from graph_notebook.neptune.utils import serialize_query_params + + +class TestSerializeQueryParams(unittest.TestCase): + """Tests for serialize_query_params. + See: https://github.com/aws/graph-notebook/issues/733 + """ + + # --- Dict input --- + + def test_dict_with_quotes_in_values(self): + params = {'key': """He said "I'm here" """} + result = serialize_query_params(params) + self.assertEqual(json.loads(result)['key'], """He said "I'm here" """) + + def test_dict_with_various_types(self): + params = {'name': 'AUS', 'limit': 10, 'active': True, 'label': None} + result = serialize_query_params(params) + self.assertEqual(json.loads(result), params) + + def test_dict_with_nested_structures(self): + params = {'filters': [{'field': 'name', 'value': "O'Brien"}]} + result = serialize_query_params(params) + self.assertEqual(json.loads(result)['filters'][0]['value'], "O'Brien") + + # --- JSON string input --- + + def test_json_string_passthrough(self): + json_str = '{"text": "I\'m happy", "num": 42}' + result = serialize_query_params(json_str) + self.assertEqual(result, json_str) + self.assertEqual(json.loads(result)['text'], "I'm happy") + + # --- Python dict literal string input --- + + def test_python_dict_literal_string(self): + result = serialize_query_params("{'text': \"I'm a string with 'quotes'\"}") + self.assertEqual(json.loads(result)['text'], "I'm a string with 'quotes'") + + # --- Invalid input --- + + def test_invalid_string_returns_none(self): + self.assertIsNone(serialize_query_params("not a dict at all {{{")) + + # --- End-to-end (magic layer → client layer, called twice) --- + + def test_end_to_end_idempotent(self): + params = {'inputs': [{'text': "I'm a string with 'single quotes'"}]} + magic_result = serialize_query_params(params) + client_result = serialize_query_params(magic_result) + self.assertEqual(json.loads(client_result)['inputs'][0]['text'], "I'm a string with 'single quotes'") + + +if __name__ == '__main__': + unittest.main()