Re-establish JDBC connection if it has been closed (#347)

microsoft · Oct 11, 2024 · 12a2e76 · 12a2e76
1 parent a7b202d
commit 12a2e76
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 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
@@ -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
@@ -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);
   }
 }