Enable query timeout for Neptune Analytics explains (#701)

* Enable query timeout for Neptune Analytics explains * update changelog
aws · Sep 21, 2024 · f087157 · f087157
1 parent 1397a77
commit f087157
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/ChangeLog.md b/ChangeLog.md
@@ -4,6 +4,8 @@ Starting with v1.31.6, this file will contain a record of major features and upd
 
 ## Upcoming
 
+- Enabled `--query-timeout` on `%%oc explain` for Neptune Analytics ([Link to PR](https://github.com/aws/graph-notebook/pull/701))
+
 ## Release 4.6.0 (September 19, 2024)
 
 - Updated Gremlin config `message_serializer` to accept all TinkerPop serializers ([Link to PR](https://github.com/aws/graph-notebook/pull/685))

diff --git a/src/graph_notebook/magics/graph_magic.py b/src/graph_notebook/magics/graph_magic.py
@@ -3597,8 +3597,11 @@ 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, query_params=query_params,
-                                              plan_cache=args.plan_cache)
+            res = self.client.opencypher_http(cell, 
+                                              explain=args.explain_type,
+                                              query_params=query_params,
+                                              plan_cache=args.plan_cache,
+                                              query_timeout=args.query_timeout)
             query_time = time.time() * 1000 - query_start
             res_replace_chars = res.content.replace(b'$', b'\&#36;')
             explain = res_replace_chars.decode("utf-8")
@@ -3617,7 +3620,8 @@ def handle_opencypher_query(self, line, cell, local_ns):
                 print("queryTimeoutMilliseconds is not supported for Neptune DB, ignoring.")
 
             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, query_params=query_params,
+            oc_http = self.client.opencypher_http(cell,
+                                                  query_params=query_params,
                                                   plan_cache=args.plan_cache,
                                                   query_timeout=args.query_timeout)
             query_time = time.time() * 1000 - query_start