Repository: googleapis/java-spanner-jdbc (java)

Issue:
Optimize the isValid method to use a cached value
The `isValid` method in a JDBC connection is used by connection pools to determine whether a connection is safe to hand out to an application. Executing a query each time for this is not very efficient for Cloud Spanner, as the Cloud Spanner JDBC connection is just a wrapper around the client library, and does not implement a physical connection with Cloud Spanner. This means that loss of network connectivity is 'all-or-nothing'; if the network connection to Cloud Spanner has been lost, then it will apply to all JDBC connections, and not only a specific one. There is therefore no point in repeatedly executing a query on Cloud Spanner to check whether a connection is valid or not.

Pull request:
perf: optimize isValid implementation
Cloud Spanner JDBC connections do not maintain a physical connection to Cloud Spanner, but are merely a wrapper around the underlying Java client library. This again uses a pool of gRPC channels to communicate with Cloud Spanner. This means that a single JDBC connection will never lose its network connection with Cloud Spanner, and checking whether it is valid or not by executing a query every time is not useful. Instead, the check should:
1. Verify that a connection has successfully been established with Cloud Spanner. The result should be cached for all JDBC connections using the same Cloud Spanner client library instance.
2. Verify that the connection has not been closed.

The above can be achieved by checking that the dialect of the database that the connection is connected to has been successfully fetched. This result is cached in the client library, and being able to get that means that there has been a valid connection.

This means that the isValid method will still return true if the network connectivity has been lost completely between the client and Cloud Spanner. However, this check is mostly used by connection pools to determine whether a connection is safe to be handed out to an application, and when all network connectivity has been lost, this will apply to all JDBC connections and not just one, meaning that the check is void.

The original isValid check can be enabled by setting the System property spanner.jdbc.use_legacy_is_valid_check to true or setting the Environment variable SPANNER_JDBC_USE_LEGACY_IS_VALID_CHECK to true.

Fixes #1443


Patch:
diff --git a/src/main/java/com/google/cloud/spanner/jdbc/JdbcConnection.java b/src/main/java/com/google/cloud/spanner/jdbc/JdbcConnection.java
index 9248dac..4cef6b3 100644
--- a/src/main/java/com/google/cloud/spanner/jdbc/JdbcConnection.java
+++ b/src/main/java/com/google/cloud/spanner/jdbc/JdbcConnection.java
@@ -29,6 +29,7 @@ import com.google.cloud.spanner.connection.Connection;
 import com.google.cloud.spanner.connection.ConnectionOptions;
 import com.google.cloud.spanner.connection.SavepointSupport;
 import com.google.cloud.spanner.connection.TransactionMode;
+import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterators;
 import java.sql.Array;
@@ -57,14 +58,42 @@ class JdbcConnection extends AbstractJdbcConnection {
       "Only result sets with concurrency CONCUR_READ_ONLY are supported";
   private static final String ONLY_CLOSE_CURSORS_AT_COMMIT =
       "Only result sets with holdability CLOSE_CURSORS_AT_COMMIT are supported";
-  static final String IS_VALID_QUERY = "SELECT 1";
+
+  /**
+   * This query is used to check the aliveness of the connection if legacy alive check has been
+   * enabled. As Cloud Spanner JDBC connections do not maintain a physical or logical connection to
+   * Cloud Spanner, there is also no point in repeatedly executing a simple query to check whether a
+   * connection is alive. Instead, we rely on the result from the initial query to Spanner that
+   * determines the dialect to determine whether the connection is alive or not. This result is
+   * cached for all JDBC connections using the same {@link com.google.cloud.spanner.Spanner}
+   * instance.
+   *
+   * <p>The legacy {@link #isValid(int)} check using a SELECT 1 statement can be enabled by setting
+   * the System property spanner.jdbc.use_legacy_is_valid_check to true or setting the environment
+   * variable SPANNER_JDBC_USE_LEGACY_IS_VALID_CHECK to true.
+   */
+  static final String LEGACY_IS_VALID_QUERY = "SELECT 1";
 
   static final ImmutableList<String> NO_GENERATED_KEY_COLUMNS = ImmutableList.of();
 
   private Map<String, Class<?>> typeMap = new HashMap<>();
 
+  private final boolean useLegacyIsValidCheck;
+
   JdbcConnection(String connectionUrl, ConnectionOptions options) throws SQLException {
     super(connectionUrl, options);
+    this.useLegacyIsValidCheck = useLegacyValidCheck();
+  }
+
+  static boolean useLegacyValidCheck() {
+    String value = System.getProperty("spanner.jdbc.use_legacy_is_valid_check");
+    if (Strings.isNullOrEmpty(value)) {
+      value = System.getenv("SPANNER_JDBC_USE_LEGACY_IS_VALID_CHECK");
+    }
+    if (!Strings.isNullOrEmpty(value)) {
+      return Boolean.parseBoolean(value);
+    }
+    return false;
   }
 
   @Override
@@ -347,23 +376,38 @@ class JdbcConnection extends AbstractJdbcConnection {
     this.typeMap = new HashMap<>(map);
   }
 
+  boolean isUseLegacyIsValidCheck() {
+    return useLegacyIsValidCheck;
+  }
+
   @Override
   public boolean isValid(int timeout) throws SQLException {
     JdbcPreconditions.checkArgument(timeout >= 0, "timeout must be >= 0");
     if (!isClosed()) {
+      if (isUseLegacyIsValidCheck()) {
+        return legacyIsValid(timeout);
+      }
       try {
-        Statement statement = createStatement();
-        statement.setQueryTimeout(timeout);
-        try (ResultSet rs = statement.executeQuery(IS_VALID_QUERY)) {
-          if (rs.next()) {
-            if (rs.getLong(1) == 1L) {
-              return true;
-            }
+        return getDialect() != null;
+      } catch (Exception ignore) {
+        // ignore and fall through.
+      }
+    }
+    return false;
+  }
+
+  private boolean legacyIsValid(int timeout) throws SQLException {
+    try (Statement statement = createStatement()) {
+      statement.setQueryTimeout(timeout);
+      try (ResultSet rs = statement.executeQuery(LEGACY_IS_VALID_QUERY)) {
+        if (rs.next()) {
+          if (rs.getLong(1) == 1L) {
+            return true;
           }
         }
-      } catch (SQLException e) {
-        // ignore
       }
+    } catch (SQLException e) {
+      // ignore and fall through.
     }
     return false;
   }


Test patch:
diff --git a/src/test/java/com/google/cloud/spanner/jdbc/JdbcConnectionTest.java b/src/test/java/com/google/cloud/spanner/jdbc/JdbcConnectionTest.java
index ff811f3..8d26bb3 100644
--- a/src/test/java/com/google/cloud/spanner/jdbc/JdbcConnectionTest.java
+++ b/src/test/java/com/google/cloud/spanner/jdbc/JdbcConnectionTest.java
@@ -502,7 +502,7 @@ public class JdbcConnectionTest {
         mock(com.google.cloud.spanner.connection.Connection.class);
     when(spannerConnection.getDialect()).thenReturn(dialect);
     when(options.getConnection()).thenReturn(spannerConnection);
-    Statement statement = Statement.of(JdbcConnection.IS_VALID_QUERY);
+    Statement statement = Statement.of(JdbcConnection.LEGACY_IS_VALID_QUERY);
 
     // Verify that an opened connection that returns a result set is valid.
     try (JdbcConnection connection = new JdbcConnection("url", options)) {
@@ -517,7 +517,7 @@ public class JdbcConnectionTest {
       }
 
       // Now let the query return an error. isValid should now return false.
-      when(spannerConnection.executeQuery(statement))
+      when(spannerConnection.getDialect())
           .thenThrow(
               SpannerExceptionFactory.newSpannerException(
                   ErrorCode.ABORTED, "the current transaction has been aborted"));