Enable query timeout for Neptune Analytics explains

aws · Sep 21, 2024 · f198957 · f198957
1 parent 1397a77
commit f198957
Showing 1 changed file with 7 additions and 3 deletions.
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