Add --use-port option to %%gremlin and %%oc (#712)

* Add enable port option to %%gremlin and %%oc

* update changelog

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 Neptune `--use-port` option to `%%gremlin` and `%%oc` ([Link to PR](https://github.com/aws/graph-notebook/pull/712))
 
 ## Release 4.6.1 (October 1, 2024)
 

src/graph_notebook/magics/graph_magic.py

@@ -1139,6 +1139,9 @@ def gremlin(self, line, cell, local_ns: dict = None):
                             help='Enable debug mode.')
         parser.add_argument('--profile-misc-args', type=str, default='{}',
                             help='Additional profile options, passed in as a map.')
+        parser.add_argument('--use-port', action='store_true', default=False,
+                            help='Includes the port in the URI for applicable Neptune HTTP requests where it is '
+                                 'excluded by default.')
         parser.add_argument('-sp', '--stop-physics', action='store_true', default=False,
                             help="Disable visualization physics after the initial simulation stabilizes.")
         parser.add_argument('-sd', '--simulation-duration', type=int, default=1500,
@@ -1208,7 +1211,8 @@ def gremlin(self, line, cell, local_ns: dict = None):
                 if self.client.is_analytics_domain() and query_params:
                     explain_args['parameters'] = query_params
                 res = self.client.gremlin_explain(cell,
-                                                  args=explain_args)
+                                                  args=explain_args,
+                                                  use_port=args.use_port)
                 res.raise_for_status()
             except Exception as e:
                 if self.client.is_analytics_domain():
@@ -1251,7 +1255,9 @@ def gremlin(self, line, cell, local_ns: dict = None):
                 print('--profile-misc-args received invalid input, please check that you are passing in a valid '
                       'string representation of a map, ex. "{\'profile.x\':\'true\'}"')
             try:
-                res = self.client.gremlin_profile(query=cell, args=profile_args)
+                res = self.client.gremlin_profile(query=cell,
+                                                  args=profile_args,
+                                                  use_port=args.use_port)
                 res.raise_for_status()
             except Exception as e:
                 if self.client.is_analytics_domain():
@@ -1302,7 +1308,8 @@ def gremlin(self, line, cell, local_ns: dict = None):
                         passed_params = query_params if self.client.is_analytics_domain() else None
                         query_res_http = self.client.gremlin_http_query(cell,
                                                                         headers=headers,
-                                                                        query_params=passed_params)
+                                                                        query_params=passed_params,
+                                                                        use_port=args.use_port)
                         query_res_http.raise_for_status()
                         try:
                             query_res_http_json = query_res_http.json()
@@ -3550,6 +3557,9 @@ def handle_opencypher_query(self, line, cell, local_ns):
         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('--use-port', action='store_true', default=False,
+                            help='Includes the port in the URI for applicable Neptune HTTP requests where it is '
+                                 'excluded by default.')
         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,
@@ -3638,7 +3648,8 @@ def handle_opencypher_query(self, line, cell, local_ns):
                                               explain=args.explain_type,
                                               query_params=query_params,
                                               plan_cache=args.plan_cache,
-                                              query_timeout=args.query_timeout)
+                                              query_timeout=args.query_timeout,
+                                              use_port=args.use_port)
             query_time = time.time() * 1000 - query_start
             res_replace_chars = res.content.replace(b'$', b'\$')
             explain = res_replace_chars.decode("utf-8")
@@ -3660,7 +3671,8 @@ def handle_opencypher_query(self, line, cell, local_ns):
             oc_http = self.client.opencypher_http(cell,
                                                   query_params=query_params,
                                                   plan_cache=args.plan_cache,
-                                                  query_timeout=args.query_timeout)
+                                                  query_timeout=args.query_timeout,
+                                                  use_port=args.use_port)
             query_time = time.time() * 1000 - query_start
             if oc_http.status_code == 400 and not self.client.is_analytics_domain() and args.plan_cache != "auto":
                 try:

src/graph_notebook/neptune/client.py

@@ -466,14 +466,15 @@ def gremlin_query(self, query, transport_args=None, bindings=None):
             c.close()
             raise e
 
-    def gremlin_http_query(self, query, headers=None, query_params: dict = None) -> requests.Response:
+    def gremlin_http_query(self, query, headers=None, query_params: dict = None,
+                           use_port: bool = False) -> requests.Response:
         if headers is None:
             headers = {}
 
         data = {}
         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'
+            uri = f'{self.get_uri(use_websocket=False, use_proxy=use_proxy, include_port=use_port)}/queries'
             data['query'] = query
             data['language'] = 'gremlin'
             headers['content-type'] = 'application/json'
@@ -498,17 +499,20 @@ def gremlin_cancel(self, query_id: str):
             raise ValueError('query_id must be a non-empty string')
         return self._query_status('gremlin', query_id=query_id, cancelQuery=True)
 
-    def gremlin_explain(self, query: str, args={}) -> requests.Response:
-        return self._gremlin_query_plan(query=query, plan_type='explain', args=args)
+    def gremlin_explain(self, query: str, use_port: bool = False, args={}) -> requests.Response:
+        return self._gremlin_query_plan(query=query, plan_type='explain', args=args, use_port=use_port)
 
-    def gremlin_profile(self, query: str, args={}) -> requests.Response:
-        return self._gremlin_query_plan(query=query, plan_type='profile', args=args)
+    def gremlin_profile(self, query: str, use_port: bool = False, args={}) -> requests.Response:
+        return self._gremlin_query_plan(query=query, plan_type='profile', args=args, use_port=use_port)
 
-    def _gremlin_query_plan(self, query: str, plan_type: str, args: dict, ) -> requests.Response:
+    def _gremlin_query_plan(self, query: str, plan_type: str, args: dict,
+                            use_port: bool = False) -> requests.Response:
         data = {}
         headers = {}
         url = f'{self._http_protocol}://{self.host}'
         if self.is_analytics_domain():
+            if use_port:
+                url += f':{self.port}'
             url += '/queries'
             data['query'] = query
             data['language'] = 'gremlin'
@@ -537,7 +541,8 @@ def _gremlin_query_plan(self, query: str, plan_type: str, args: dict, ) -> reque
     def opencypher_http(self, query: str, headers: dict = None, explain: str = None,
                         query_params: dict = None,
                         plan_cache: str = None,
-                        query_timeout: int = None) -> requests.Response:
+                        query_timeout: int = None,
+                        use_port: bool = False) -> requests.Response:
         if headers is None:
             headers = {}
 
@@ -546,6 +551,8 @@ def opencypher_http(self, query: str, headers: dict = None, explain: str = None,
         if self.is_neptune_domain():
             data = {}
             if self.is_analytics_domain():
+                if use_port:
+                    url += f':{self.port}'
                 url += f'/queries'
                 data['language'] = 'opencypher'
             else: