From 768a7b8c004f2ab0a4284ab29479cd917f486c1b Mon Sep 17 00:00:00 2001
From: Mike Thomsen <mikerthomsen@gmail.com>
Date: Thu, 4 Jul 2019 07:50:21 -0400
Subject: [PATCH] NIFI-6425 Made executeQuery able to reconnect to the Gremlin
 cluster.

Signed-off-by: Matthew Burgess <mattyb149@apache.org>

This closes #3572
---
 .../nifi/graph/GremlinClientService.java       | 18 ++++++++++++++++--
 1 file changed, 16 insertions(+), 2 deletions(-)

diff --git a/nifi-nar-bundles/nifi-graph-bundle/nifi-other-graph-services/src/main/java/org/apache/nifi/graph/GremlinClientService.java b/nifi-nar-bundles/nifi-graph-bundle/nifi-other-graph-services/src/main/java/org/apache/nifi/graph/GremlinClientService.java
index a79094f427..e001fd0493 100644
--- a/nifi-nar-bundles/nifi-graph-bundle/nifi-other-graph-services/src/main/java/org/apache/nifi/graph/GremlinClientService.java
+++ b/nifi-nar-bundles/nifi-graph-bundle/nifi-other-graph-services/src/main/java/org/apache/nifi/graph/GremlinClientService.java
@@ -37,9 +37,11 @@ public class GremlinClientService extends AbstractTinkerpopClientService impleme
     private Cluster cluster;
     protected Client client;
     public static final String NOT_SUPPORTED = "NOT_SUPPORTED";
+    private ConfigurationContext context;
 
     @OnEnabled
     public void onEnabled(ConfigurationContext context) {
+        this.context = context;
         cluster = buildCluster(context);
         client = cluster.connect();
     }
@@ -52,8 +54,7 @@ public class GremlinClientService extends AbstractTinkerpopClientService impleme
         cluster = null;
     }
 
-    @Override
-    public Map<String, String> executeQuery(String query, Map<String, Object> parameters, GraphQueryResultCallback handler) {
+    public Map<String, String> doQuery(String query, Map<String, Object> parameters, GraphQueryResultCallback handler) {
         try {
             Iterator<Result> iterator = client.submit(query, parameters).iterator();
             long count = 0;
@@ -86,6 +87,19 @@ public class GremlinClientService extends AbstractTinkerpopClientService impleme
         }
     }
 
+    @Override
+    public Map<String, String> executeQuery(String query, Map<String, Object> parameters, GraphQueryResultCallback handler) {
+        try {
+            return doQuery(query, parameters, handler);
+        } catch (Exception ex) {
+            cluster.close();
+            client.close();
+            cluster = buildCluster(context);
+            client = cluster.connect();
+            return doQuery(query, parameters, handler);
+        }
+    }
+
     @Override
     public String getTransitUrl() {
         return transitUrl;