From 12a2e7695f8bb5f968400e8c1c69a8427e0b253d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jes=C3=BAs=20Camacho=20Rodr=C3=ADguez?=
 <jesusca@microsoft.com>
Date: Fri, 11 Oct 2024 11:55:23 -0700
Subject: [PATCH] Re-establish JDBC connection if it has been closed (#347)

---
 .../lst_bench/client/JDBCConnection.java      | 37 +++++++++++++++++--
 .../client/JDBCConnectionManager.java         | 20 ++--------
 2 files changed, 36 insertions(+), 21 deletions(-)

diff --git a/core/src/main/java/com/microsoft/lst_bench/client/JDBCConnection.java b/core/src/main/java/com/microsoft/lst_bench/client/JDBCConnection.java
index a15e5f86..ece26016 100644
--- a/core/src/main/java/com/microsoft/lst_bench/client/JDBCConnection.java
+++ b/core/src/main/java/com/microsoft/lst_bench/client/JDBCConnection.java
@@ -15,12 +15,15 @@
  */
 package com.microsoft.lst_bench.client;
 
+import java.sql.DriverManager;
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.SQLWarning;
 import java.sql.Statement;
 import java.util.ArrayList;
 import java.util.List;
+import javax.annotation.Nullable;
+import org.apache.commons.lang3.StringUtils;
 import org.apache.commons.lang3.exception.ExceptionUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -30,14 +33,22 @@ public class JDBCConnection implements Connection {
 
   private static final Logger LOGGER = LoggerFactory.getLogger(JDBCConnection.class);
 
-  private final java.sql.Connection connection;
+  private final String url;
   private final int maxNumRetries;
   private final boolean showWarnings;
+  @Nullable private final String username;
+  @Nullable private final String password;
 
-  public JDBCConnection(java.sql.Connection connection, int maxNumRetries, boolean showWarnings) {
-    this.connection = connection;
+  @Nullable private java.sql.Connection connection;
+
+  public JDBCConnection(
+      String url, int maxNumRetries, boolean showWarnings, String username, String password) {
+    this.url = url;
+    this.username = username;
+    this.password = password;
     this.maxNumRetries = maxNumRetries;
     this.showWarnings = showWarnings;
+    this.connection = null;
   }
 
   @Override
@@ -51,6 +62,8 @@ public QueryResult executeQuery(String sqlText) throws ClientException {
   }
 
   private QueryResult execute(String sqlText, boolean ignoreResults) throws ClientException {
+    refreshConnection();
+
     QueryResult queryResult = null;
     int errorCount = 0;
 
@@ -129,7 +142,23 @@ private QueryResult execute(String sqlText, boolean ignoreResults) throws Client
   @Override
   public void close() throws ClientException {
     try {
-      connection.close();
+      if (connection != null && !connection.isClosed()) {
+        connection.close();
+      }
+    } catch (SQLException e) {
+      throw new ClientException(e);
+    }
+  }
+
+  private synchronized void refreshConnection() throws ClientException {
+    try {
+      if (connection == null || connection.isClosed()) {
+        if (StringUtils.isEmpty(username)) {
+          connection = DriverManager.getConnection(url);
+        } else {
+          connection = DriverManager.getConnection(url, username, password);
+        }
+      }
     } catch (SQLException e) {
       throw new ClientException(e);
     }
diff --git a/core/src/main/java/com/microsoft/lst_bench/client/JDBCConnectionManager.java b/core/src/main/java/com/microsoft/lst_bench/client/JDBCConnectionManager.java
index 9f3bdf4b..e15bb7b0 100644
--- a/core/src/main/java/com/microsoft/lst_bench/client/JDBCConnectionManager.java
+++ b/core/src/main/java/com/microsoft/lst_bench/client/JDBCConnectionManager.java
@@ -15,10 +15,7 @@
  */
 package com.microsoft.lst_bench.client;
 
-import java.sql.DriverManager;
-import java.sql.SQLException;
 import javax.annotation.Nullable;
-import org.apache.commons.lang3.StringUtils;
 
 /** Simple JDBC connection manager. */
 public class JDBCConnectionManager implements ConnectionManager {
@@ -41,19 +38,8 @@ public JDBCConnectionManager(
   }
 
   @Override
-  public Connection createConnection() throws ClientException {
-    try {
-      if (StringUtils.isEmpty(username)) {
-        return new JDBCConnection(
-            DriverManager.getConnection(url), this.maxNumRetries, this.showWarnings);
-      } else {
-        return new JDBCConnection(
-            DriverManager.getConnection(url, username, password),
-            this.maxNumRetries,
-            this.showWarnings);
-      }
-    } catch (SQLException e) {
-      throw new ClientException(e);
-    }
+  public Connection createConnection() {
+    return new JDBCConnection(
+        this.url, this.maxNumRetries, this.showWarnings, this.username, this.password);
   }
 }