diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index d0e1d5b441..f9e9376421 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -8,6 +8,7 @@ - **SQL Scripting support**: Added support for [SQL Scripting](https://docs.databricks.com/aws/en/sql/language-manual/sql-ref-scripting) - Added a client property `enableVolumeOperations` to enable GET/PUT/REMOVE volume operations on a stream. For backward compatibility, allowedVolumeIngestionPaths can also be used for REMOVE operation. - Support for fetching schemas across all catalogs (when catalog is specified as null or a wildcard) in `DatabaseMetaData#getSchemas` API in SQL Execution mode. +- **Configurable SQL validation in isValid()**: Added `EnableSQLValidationForIsValid` connection property to control whether `isValid()` method executes an actual SQL query for server-side validation. Default value is 0. ### Updated - Databricks SDK dependency upgraded to latest version 0.60.0 diff --git a/README.md b/README.md index 41f930e294..db37290510 100644 --- a/README.md +++ b/README.md @@ -91,6 +91,7 @@ Optional parameters: - `OAuth2RedirectUrlPort` - Ports for redirect URL (default: 8020) - `EnableOIDCDiscovery` - Enable OIDC discovery (default: 1) - `OAuthDiscoveryURL` - OIDC discovery endpoint (default: /oidc/.well-known/oauth-authorization-server) +- `EnableSQLValidationForIsValid` - Enable SQL query based validation in `isValid()` connection checks (default: 0) ### Logging diff --git a/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnection.java b/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnection.java index 4c1555a20d..5775002da0 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnection.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnection.java @@ -414,7 +414,21 @@ public SQLXML createSQLXML() throws SQLException { @Override public boolean isValid(int timeout) throws SQLException { ValidationUtil.checkIfNonNegative(timeout, "timeout"); - return !isClosed(); + if (isClosed()) { + return false; + } + if (connectionContext.getEnableSQLValidationForIsValid()) { + try (Statement stmt = createStatement()) { + stmt.setQueryTimeout(timeout); + // This is a lightweight query to check if the connection is valid + stmt.execute("SELECT VERSION()"); + return true; + } catch (Exception e) { + LOGGER.debug("Validation failed for isValid(): {}", e.getMessage()); + return false; + } + } + return true; } /** diff --git a/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java b/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java index 9b26b4f4aa..e2fe52c771 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java @@ -225,6 +225,12 @@ public String getHttpPath() { return getParameter(DatabricksJdbcUrlParams.HTTP_PATH); } + public boolean getEnableSQLValidationForIsValid() { + LOGGER.debug("String getEnableSQLValidationForIsValid()"); + return getParameter(DatabricksJdbcUrlParams.ENABLE_SQL_VALIDATION_FOR_IS_VALID, "0") + .equals("1"); + } + @Override public String getHostForOAuth() { return this.host; diff --git a/src/main/java/com/databricks/jdbc/api/internal/IDatabricksConnectionContext.java b/src/main/java/com/databricks/jdbc/api/internal/IDatabricksConnectionContext.java index 2da74837d4..1b22e1bdd4 100644 --- a/src/main/java/com/databricks/jdbc/api/internal/IDatabricksConnectionContext.java +++ b/src/main/java/com/databricks/jdbc/api/internal/IDatabricksConnectionContext.java @@ -89,6 +89,9 @@ public interface IDatabricksConnectionContext { String getHttpPath(); + /** Returns the value of the EnableSQLValidationForIsValid connection property. */ + boolean getEnableSQLValidationForIsValid(); + String getProxyHost(); int getProxyPort(); diff --git a/src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java b/src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java index 433447a162..d0b04dee37 100644 --- a/src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java +++ b/src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java @@ -153,7 +153,11 @@ public enum DatabricksJdbcUrlParams { HTTP_CONNECTION_REQUEST_TIMEOUT( "HttpConnectionRequestTimeout", "HTTP connection request timeout in seconds"), CLOUD_FETCH_SPEED_THRESHOLD( - "CloudFetchSpeedThreshold", "Minimum expected download speed in MB/s", "0.1"); + "CloudFetchSpeedThreshold", "Minimum expected download speed in MB/s", "0.1"), + ENABLE_SQL_VALIDATION_FOR_IS_VALID( + "EnableSQLValidationForIsValid", + "Enable SQL query execution for connection validation in isValid() method", + "0"); private final String paramName; private final String defaultValue; diff --git a/src/main/java/com/databricks/jdbc/common/util/DatabricksDriverPropertyUtil.java b/src/main/java/com/databricks/jdbc/common/util/DatabricksDriverPropertyUtil.java index 05935c1ee5..8ec0766c3a 100644 --- a/src/main/java/com/databricks/jdbc/common/util/DatabricksDriverPropertyUtil.java +++ b/src/main/java/com/databricks/jdbc/common/util/DatabricksDriverPropertyUtil.java @@ -37,7 +37,8 @@ public class DatabricksDriverPropertyUtil { DatabricksJdbcUrlParams.ROWS_FETCHED_PER_BLOCK, DatabricksJdbcUrlParams.DEFAULT_STRING_COLUMN_LENGTH, DatabricksJdbcUrlParams.SOCKET_TIMEOUT, - DatabricksJdbcUrlParams.ENABLE_TOKEN_CACHE); + DatabricksJdbcUrlParams.ENABLE_TOKEN_CACHE, + DatabricksJdbcUrlParams.ENABLE_SQL_VALIDATION_FOR_IS_VALID); public static List getMissingProperties(String url, Properties info) throws DatabricksParsingException { diff --git a/src/test/java/com/databricks/jdbc/api/impl/DatabricksConnectionTest.java b/src/test/java/com/databricks/jdbc/api/impl/DatabricksConnectionTest.java index 25ccc392e3..afc0905117 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/DatabricksConnectionTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/DatabricksConnectionTest.java @@ -483,4 +483,38 @@ public void testQueryTagsInSessionConfigs() throws SQLException { assertTrue(sessionConfigs.containsKey("query_tags")); assertEquals("team:marketing,dashboard:abc123", sessionConfigs.get("query_tags")); } + + @Test + public void testIsValidWithSQLValidationEnabled() throws SQLException { + String jdbcUrlWithValidation = CATALOG_SCHEMA_JDBC_URL + ";EnableSQLValidationForIsValid=1"; + IDatabricksConnectionContext connectionContextWithValidation = + DatabricksConnectionContext.parse(jdbcUrlWithValidation, new Properties()); + when(databricksClient.createSession( + new Warehouse(WAREHOUSE_ID), CATALOG, SCHEMA, new HashMap<>())) + .thenReturn(IMMUTABLE_SESSION_INFO); + connection = new DatabricksConnection(connectionContextWithValidation, databricksClient); + connection.open(); + DatabricksConnection spyConnection = spy(connection); + DatabricksStatement mockStatement = mock(DatabricksStatement.class); + doReturn(mockStatement).when(spyConnection).createStatement(); + doNothing().when(mockStatement).setQueryTimeout(anyInt()); + when(mockStatement.execute("SELECT VERSION()")).thenReturn(true); + + assertTrue(spyConnection.isValid(5)); + verify(spyConnection).createStatement(); + verify(mockStatement).setQueryTimeout(5); + verify(mockStatement).execute("SELECT VERSION()"); + + DatabricksStatement mockStatementFail = mock(DatabricksStatement.class); + doReturn(mockStatementFail).when(spyConnection).createStatement(); + doNothing().when(mockStatementFail).setQueryTimeout(anyInt()); + when(mockStatementFail.execute("SELECT VERSION()")) + .thenThrow(new SQLException("Connection lost")); + + assertFalse(spyConnection.isValid(5)); + verify(mockStatementFail).setQueryTimeout(5); + verify(mockStatementFail).execute("SELECT VERSION()"); + + connection.close(); + } }