From 077c11855bf023b8d6239b33746040856b3cedff Mon Sep 17 00:00:00 2001
From: Michael Chin <mchin188@yahoo.com>
Date: Tue, 10 Sep 2024 22:27:56 -0700
Subject: [PATCH] Update Gremlin HTTP requests and error handling (#692)

* Update Gremlin HTTP requests and error handling

* update changelog
---
 ChangeLog.md                             | 1 +
 requirements.txt                         | 1 +
 setup.py                                 | 3 ++-
 src/graph_notebook/magics/graph_magic.py | 7 ++++++-
 src/graph_notebook/neptune/client.py     | 9 +++++----
 5 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/ChangeLog.md b/ChangeLog.md
index cb5ba9db..addca271 100644
--- a/ChangeLog.md
+++ b/ChangeLog.md
@@ -11,6 +11,7 @@ Starting with v1.31.6, this file will contain a record of major features and upd
 - Improved iPython config directory retrieval logic ([Link to PR](https://github.com/aws/graph-notebook/pull/687))
 - Fixed `%db_reset` output for token modes ([Link to PR](https://github.com/aws/graph-notebook/pull/691))
 - Use `extras_require` to specify tests ([Link to PR](https://github.com/aws/graph-notebook/pull/688))
+- Updated Gremlin HTTP requests, fixed handling of internal error responses ([Link to PR](https://github.com/aws/graph-notebook/pull/692))
 
 ## Release 4.5.2 (August 15, 2024)
 
diff --git a/requirements.txt b/requirements.txt
index 43b206af..4dc3232e 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -24,6 +24,7 @@ pandas>=2.1.0,<=2.2.2
 numpy<1.24.0
 nest_asyncio>=1.5.5,<=1.5.6
 async-timeout>=4.0,<5.0
+json-repair==0.29.2
 
 # requirements for testing
 botocore~=1.34.74
diff --git a/setup.py b/setup.py
index 0d422c21..5254fffe 100644
--- a/setup.py
+++ b/setup.py
@@ -91,7 +91,8 @@ def get_version():
         'pandas>=2.1.0,<=2.2.2',
         'numpy<1.24.0',
         'nest_asyncio>=1.5.5,<=1.5.6',
-        'async-timeout>=4.0,<5.0'
+        'async-timeout>=4.0,<5.0',
+        'json-repair==0.29.2'
     ],
     package_data={
         'graph_notebook': ['graph_notebook/widgets/nbextensions/**',
diff --git a/src/graph_notebook/magics/graph_magic.py b/src/graph_notebook/magics/graph_magic.py
index bff554df..44423455 100644
--- a/src/graph_notebook/magics/graph_magic.py
+++ b/src/graph_notebook/magics/graph_magic.py
@@ -36,6 +36,7 @@
 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
 from requests import HTTPError
+from json_repair import repair_json
 
 import graph_notebook
 from graph_notebook.configuration.generate_config import generate_default_config, DEFAULT_CONFIG_LOCATION, \
@@ -1257,7 +1258,11 @@ def gremlin(self, line, cell, local_ns: dict = None):
                         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()
+                        try:
+                            query_res_http_json = query_res_http.json()
+                        except JSONDecodeError:
+                            query_res_fixed = repair_json(query_res_http.text)
+                            query_res_http_json = json.loads(query_res_fixed)
                         query_res = query_res_http_json['result']['data']
                     else:
                         query_res = self.client.gremlin_query(cell, transport_args=transport_args)
diff --git a/src/graph_notebook/neptune/client.py b/src/graph_notebook/neptune/client.py
index be9cef34..8508e753 100644
--- a/src/graph_notebook/neptune/client.py
+++ b/src/graph_notebook/neptune/client.py
@@ -406,8 +406,8 @@ def gremlin_http_query(self, query, headers=None) -> requests.Response:
         use_proxy = True if self.proxy_host != '' else False
         if self.is_analytics_domain():
             uri = f'{self.get_uri(use_websocket=False, use_proxy=use_proxy, include_port=False)}/queries'
+            data['query'] = query
             data['language'] = 'gremlin'
-            data['gremlin'] = query
             headers['content-type'] = 'application/json'
         else:
             uri = f'{self.get_uri(use_websocket=False, use_proxy=use_proxy)}/gremlin'
@@ -440,13 +440,14 @@ def _gremlin_query_plan(self, query: str, plan_type: str, args: dict, ) -> reque
         url = f'{self._http_protocol}://{self.host}'
         if self.is_analytics_domain():
             url += '/queries'
-            data['gremlin'] = query
+            data['query'] = query
             data['language'] = 'gremlin'
             headers['content-type'] = 'application/json'
             if plan_type == 'explain':
-                data['explain.mode'] = args.pop('explain.mode')
+                data['explain-mode'] = args.pop('explain.mode')
             elif plan_type == 'profile':
-                data['profile.debug'] = args.pop('profile.debug')
+                for param, value in args.items():
+                    data[param] = value
         else:
             url += f':{self.port}/gremlin/{plan_type}'
             data['gremlin'] = query