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 9248dace..4cef6b3f 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.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. + * + *

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 NO_GENERATED_KEY_COLUMNS = ImmutableList.of(); private Map> 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 @@ public void setTypeMap(Map> map) throws SQLException { 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; } 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 ff811f37..8d26bb31 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 void testIsValid() throws SQLException { 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 void testIsValid() throws SQLException { } // 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"));