From 688cf75d94387c327229bf1c3339d78f49d8f813 Mon Sep 17 00:00:00 2001 From: Gopal Lal Date: Wed, 12 Feb 2025 15:41:09 +0530 Subject: [PATCH 1/5] Add support for configuring HTTP connection pool size --- .../jdbc/api/IDatabricksConnectionContext.java | 3 +++ .../jdbc/api/impl/DatabricksConnectionContext.java | 5 +++++ .../databricks/jdbc/common/DatabricksJdbcUrlParams.java | 1 + .../jdbc/dbclient/impl/common/ClientConfigurator.java | 4 ++-- .../dbclient/impl/common/ClientConfiguratorTest.java | 9 +++++++++ .../dbclient/impl/sqlexec/DatabricksSdkClientTest.java | 1 + 6 files changed, 21 insertions(+), 2 deletions(-) diff --git a/src/main/java/com/databricks/jdbc/api/IDatabricksConnectionContext.java b/src/main/java/com/databricks/jdbc/api/IDatabricksConnectionContext.java index 53bbdf19d3..9995baeb00 100644 --- a/src/main/java/com/databricks/jdbc/api/IDatabricksConnectionContext.java +++ b/src/main/java/com/databricks/jdbc/api/IDatabricksConnectionContext.java @@ -223,4 +223,7 @@ public interface IDatabricksConnectionContext { /** Returns true if driver return complex data type java objects natively as opposed to string */ boolean isComplexDatatypeSupportEnabled(); + + /** Returns the size for HTTP connection pool */ + int getHttpConnectionPoolSize(); } 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 651c601ded..03ac732233 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java @@ -699,6 +699,11 @@ public boolean isRequestTracingEnabled() { return getParameter(DatabricksJdbcUrlParams.ENABLE_REQUEST_TRACING).equals("1"); } + @Override + public int getHttpConnectionPoolSize() { + return Integer.parseInt(getParameter(DatabricksJdbcUrlParams.HTTP_CONNECTION_POOL_SIZE, "100")); + } + private static boolean nullOrEmptyString(String s) { return s == null || s.isEmpty(); } diff --git a/src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java b/src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java index bd1f0a44c5..5a102d31ad 100644 --- a/src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java +++ b/src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java @@ -79,6 +79,7 @@ public enum DatabricksJdbcUrlParams { ALLOWED_VOLUME_INGESTION_PATHS("VolumeOperationAllowedLocalPaths", ""), ALLOWED_STAGING_INGESTION_PATHS("StagingAllowedLocalPaths", ""), ENABLE_REQUEST_TRACING("EnableRequestTracing", "flag to enable request tracing", "0"), + HTTP_CONNECTION_POOL_SIZE("HttpConnectionPoolSize", "Maximum HTTP connection pool size", "100"), ENABLE_SQL_EXEC_HYBRID_RESULTS( "EnableSQLExecHybridResults", "flag to enable hybrid results", "0"), ENABLE_COMPLEX_DATATYPE_SUPPORT( diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/common/ClientConfigurator.java b/src/main/java/com/databricks/jdbc/dbclient/impl/common/ClientConfigurator.java index 82f1a39315..cd05d5ced2 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/common/ClientConfigurator.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/common/ClientConfigurator.java @@ -53,8 +53,8 @@ public ClientConfigurator(IDatabricksConnectionContext connectionContext) { void setupConnectionManager(CommonsHttpClient.Builder httpClientBuilder) { PoolingHttpClientConnectionManager connManager = ConfiguratorUtils.getBaseConnectionManager(connectionContext); - // This is consistent with the value in the SDK - connManager.setMaxTotal(100); + // Default value is 100 which is consistent with the value in the SDK + connManager.setMaxTotal(connectionContext.getHttpConnectionPoolSize()); httpClientBuilder.withConnectionManager(connManager); } diff --git a/src/test/java/com/databricks/jdbc/dbclient/impl/common/ClientConfiguratorTest.java b/src/test/java/com/databricks/jdbc/dbclient/impl/common/ClientConfiguratorTest.java index 6acb52e5a0..b12067fbc0 100644 --- a/src/test/java/com/databricks/jdbc/dbclient/impl/common/ClientConfiguratorTest.java +++ b/src/test/java/com/databricks/jdbc/dbclient/impl/common/ClientConfiguratorTest.java @@ -40,6 +40,7 @@ void getWorkspaceClient_PAT_AuthenticatesWithAccessToken() throws DatabricksPars when(mockContext.getAuthMech()).thenReturn(AuthMech.PAT); when(mockContext.getHostUrl()).thenReturn("https://pat.databricks.com"); when(mockContext.getToken()).thenReturn("pat-token"); + when(mockContext.getHttpConnectionPoolSize()).thenReturn(100); configurator = new ClientConfigurator(mockContext); WorkspaceClient client = configurator.getWorkspaceClient(); @@ -58,6 +59,7 @@ void getWorkspaceClient_OAuthWithTokenPassthrough_AuthenticatesCorrectly() when(mockContext.getAuthFlow()).thenReturn(AuthFlow.TOKEN_PASSTHROUGH); when(mockContext.getHostUrl()).thenReturn("https://oauth-token.databricks.com"); when(mockContext.getPassThroughAccessToken()).thenReturn("oauth-token"); + when(mockContext.getHttpConnectionPoolSize()).thenReturn(100); configurator = new ClientConfigurator(mockContext); WorkspaceClient client = configurator.getWorkspaceClient(); @@ -77,6 +79,7 @@ void getWorkspaceClient_OAuthWithClientCredentials_AuthenticatesCorrectly() when(mockContext.getHostForOAuth()).thenReturn("https://oauth-client.databricks.com"); when(mockContext.getClientId()).thenReturn("client-id"); when(mockContext.getClientSecret()).thenReturn("client-secret"); + when(mockContext.getHttpConnectionPoolSize()).thenReturn(100); configurator = new ClientConfigurator(mockContext); WorkspaceClient client = configurator.getWorkspaceClient(); @@ -98,6 +101,7 @@ void getWorkspaceClient_OAuthWithClientCredentials_AuthenticatesCorrectlyGCP() when(mockContext.getCloud()).thenReturn(Cloud.GCP); when(mockContext.getGcpAuthType()).thenReturn(GCP_GOOGLE_CREDENTIALS_AUTH_TYPE); when(mockContext.getGoogleCredentials()).thenReturn("google-credentials"); + when(mockContext.getHttpConnectionPoolSize()).thenReturn(100); configurator = new ClientConfigurator(mockContext); WorkspaceClient client = configurator.getWorkspaceClient(); assertNotNull(client); @@ -119,6 +123,7 @@ void getWorkspaceClient_OAuthWithClientCredentials_AuthenticatesCorrectlyWithJWT when(mockContext.getClientSecret()).thenReturn("client-secret"); when(mockContext.useJWTAssertion()).thenReturn(true); when(mockContext.getTokenEndpoint()).thenReturn("token-endpoint"); + when(mockContext.getHttpConnectionPoolSize()).thenReturn(100); configurator = new ClientConfigurator(mockContext); WorkspaceClient client = configurator.getWorkspaceClient(); @@ -163,6 +168,7 @@ void getWorkspaceClient_OAuthWithBrowserBasedAuthentication_AuthenticatesCorrect when(mockContext.getClientId()).thenReturn("browser-client-id"); when(mockContext.getClientSecret()).thenReturn("browser-client-secret"); when(mockContext.getOAuthScopesForU2M()).thenReturn(List.of(new String[] {"scope1", "scope2"})); + when(mockContext.getHttpConnectionPoolSize()).thenReturn(100); configurator = new ClientConfigurator(mockContext); WorkspaceClient client = configurator.getWorkspaceClient(); assertNotNull(client); @@ -188,6 +194,7 @@ void getWorkspaceClient_OAuthWithBrowserBasedAuthentication_AuthenticatesCorrect when(mockContext.getOAuthScopesForU2M()).thenReturn(List.of(new String[] {"scope1", "scope2"})); when(mockContext.isOAuthDiscoveryModeEnabled()).thenReturn(true); when(mockContext.getOAuthDiscoveryURL()).thenReturn(TEST_DISCOVERY_URL); + when(mockContext.getHttpConnectionPoolSize()).thenReturn(100); configurator = new ClientConfigurator(mockContext); WorkspaceClient client = configurator.getWorkspaceClient(); assertNotNull(client); @@ -205,6 +212,7 @@ void getWorkspaceClient_OAuthWithBrowserBasedAuthentication_AuthenticatesCorrect @Test void testNonOauth() { when(mockContext.getAuthMech()).thenReturn(AuthMech.OTHER); + when(mockContext.getHttpConnectionPoolSize()).thenReturn(100); configurator = new ClientConfigurator(mockContext); DatabricksConfig config = configurator.getDatabricksConfig(); assertEquals(DatabricksJdbcConstants.ACCESS_TOKEN_AUTH_TYPE, config.getAuthType()); @@ -240,6 +248,7 @@ void testSetupProxyConfig() { when(mockContext.getProxyUser()).thenReturn("proxyUser"); when(mockContext.getProxyPassword()).thenReturn("proxyPass"); when(mockContext.getProxyAuthType()).thenReturn(ProxyConfig.ProxyAuthType.values()[0]); + when(mockContext.getHttpConnectionPoolSize()).thenReturn(100); // For non-proxy hosts conversion, an input of ".example.com,localhost" // is expected to be converted to "*.example.com|localhost" when(mockContext.getNonProxyHosts()).thenReturn(".example.com,localhost"); diff --git a/src/test/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClientTest.java b/src/test/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClientTest.java index b0af0e2d0d..58b3f6a2b0 100644 --- a/src/test/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClientTest.java +++ b/src/test/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClientTest.java @@ -40,6 +40,7 @@ void testHandleFailedExecution() throws SQLException { when(status.getSqlState()).thenReturn(DEFAULT_HTTP_EXCEPTION_SQLSTATE); when(errorInfo.getMessage()).thenReturn("Error message"); when(errorInfo.getErrorCode()).thenReturn(ServiceErrorCode.BAD_REQUEST); + when(connectionContext.getHttpConnectionPoolSize()).thenReturn(100); DatabricksSdkClient databricksSdkClient = new DatabricksSdkClient(connectionContext, statementExecutionService, apiClient); assertThrows(SQLException.class, () -> databricksSdkClient.getMoreResults(statementInternal)); From 2e6bfaa260d7585476f369f224c6e608dfe62a55 Mon Sep 17 00:00:00 2001 From: Gopal Lal Date: Wed, 12 Feb 2025 16:20:45 +0530 Subject: [PATCH 2/5] Changes for default value --- .../databricks/jdbc/api/impl/DatabricksConnectionContext.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 03ac732233..b3144ec140 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java @@ -701,7 +701,7 @@ public boolean isRequestTracingEnabled() { @Override public int getHttpConnectionPoolSize() { - return Integer.parseInt(getParameter(DatabricksJdbcUrlParams.HTTP_CONNECTION_POOL_SIZE, "100")); + return Integer.parseInt(getParameter(DatabricksJdbcUrlParams.HTTP_CONNECTION_POOL_SIZE)); } private static boolean nullOrEmptyString(String s) { From 82c186f87ab2d70965226389bd8f7ec8f3238f72 Mon Sep 17 00:00:00 2001 From: Gopal Lal Date: Fri, 30 May 2025 14:27:13 +0530 Subject: [PATCH 3/5] Update SDK version --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index ed476605c4..b98b32bd03 100644 --- a/pom.xml +++ b/pom.xml @@ -54,7 +54,7 @@ 4.5.14 2.10.1 2.14.0 - 0.46.0 + 0.52.0 3.1.2 0.3 1.8.0 From 30cdedb8ef17bda56a2e1e99a5f4846c0ab0e561 Mon Sep 17 00:00:00 2001 From: Gopal Lal Date: Fri, 30 May 2025 14:27:56 +0530 Subject: [PATCH 4/5] Update SDK version2 --- NEXT_CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 1314352a68..d45eb959c1 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -3,7 +3,7 @@ ## [Unreleased] ### Added -- +- Added support for DoD (.mil) domains ### Updated - From 3b8a8490f4351d2f7784cb75e9565cff12c8fc66 Mon Sep 17 00:00:00 2001 From: Gopal Lal Date: Mon, 2 Jun 2025 14:09:56 +0530 Subject: [PATCH 5/5] Fix test --- .../jdbc/dbclient/impl/common/ClientConfiguratorTest.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/test/java/com/databricks/jdbc/dbclient/impl/common/ClientConfiguratorTest.java b/src/test/java/com/databricks/jdbc/dbclient/impl/common/ClientConfiguratorTest.java index 94a37f918b..a091e37ca7 100644 --- a/src/test/java/com/databricks/jdbc/dbclient/impl/common/ClientConfiguratorTest.java +++ b/src/test/java/com/databricks/jdbc/dbclient/impl/common/ClientConfiguratorTest.java @@ -137,7 +137,7 @@ void getWorkspaceClient_OAuthWithClientCredentials_AuthenticatesCorrectlyWithJWT assertEquals("https://jwt-auth.databricks.com", config.getHost()); assertEquals("client-id", config.getClientId()); assertEquals("client-secret", config.getClientSecret()); - assertEquals(DatabricksJdbcConstants.M2M_AUTH_TYPE, config.getAuthType()); + assertEquals("custom-oauth-m2m", config.getAuthType()); DatabricksTokenFederationProvider provider = (DatabricksTokenFederationProvider) config.getCredentialsProvider(); assertNotNull(provider);