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 a72870ec90018d40be49f02f1c60dc1f067095e3 Mon Sep 17 00:00:00 2001 From: Gopal Lal Date: Fri, 25 Apr 2025 12:25:27 +0530 Subject: [PATCH 3/5] Adding extra null check for thread local variables in Telemetry --- .../jdbc/telemetry/TelemetryHelper.java | 70 +++++++++++-------- 1 file changed, 39 insertions(+), 31 deletions(-) diff --git a/src/main/java/com/databricks/jdbc/telemetry/TelemetryHelper.java b/src/main/java/com/databricks/jdbc/telemetry/TelemetryHelper.java index 4673fdd2d9..ce837157fa 100644 --- a/src/main/java/com/databricks/jdbc/telemetry/TelemetryHelper.java +++ b/src/main/java/com/databricks/jdbc/telemetry/TelemetryHelper.java @@ -113,44 +113,52 @@ static void exportLatencyLog( long latencyMilliseconds, SqlExecutionEvent executionEvent, StatementId statementId) { - TelemetryEvent telemetryEvent = - new TelemetryEvent() - .setLatency(latencyMilliseconds) - .setSqlOperation(executionEvent) - .setDriverConnectionParameters(getDriverConnectionParameter(connectionContext)); - if (statementId != null) { - telemetryEvent.setSqlStatementId(statementId.toString()); + // Though we already handle null connectionContext in the downstream implementation, + // we are adding this check for extra sanity + if (connectionContext != null) { + TelemetryEvent telemetryEvent = + new TelemetryEvent() + .setLatency(latencyMilliseconds) + .setSqlOperation(executionEvent) + .setDriverConnectionParameters(getDriverConnectionParameter(connectionContext)); + if (statementId != null) { + telemetryEvent.setSqlStatementId(statementId.toString()); + } + TelemetryFrontendLog telemetryFrontendLog = + new TelemetryFrontendLog() + .setFrontendLogEventId(getEventUUID()) + .setContext(getLogContext()) + .setEntry(new FrontendLogEntry().setSqlDriverLog(telemetryEvent)); + TelemetryClientFactory.getInstance() + .getTelemetryClient(connectionContext, DatabricksThreadContextHolder.getDatabricksConfig()) + .exportEvent(telemetryFrontendLog); } - TelemetryFrontendLog telemetryFrontendLog = - new TelemetryFrontendLog() - .setFrontendLogEventId(getEventUUID()) - .setContext(getLogContext()) - .setEntry(new FrontendLogEntry().setSqlDriverLog(telemetryEvent)); - TelemetryClientFactory.getInstance() - .getTelemetryClient(connectionContext, DatabricksThreadContextHolder.getDatabricksConfig()) - .exportEvent(telemetryFrontendLog); } public static void exportLatencyLog( IDatabricksConnectionContext connectionContext, long latencyMilliseconds, DriverVolumeOperation volumeOperationEvent) { - TelemetryFrontendLog telemetryFrontendLog = - new TelemetryFrontendLog() - .setFrontendLogEventId(getEventUUID()) - .setContext(getLogContext()) - .setEntry( - new FrontendLogEntry() - .setSqlDriverLog( - new TelemetryEvent() - .setLatency(latencyMilliseconds) - .setVolumeOperation(volumeOperationEvent) - .setDriverConnectionParameters( - getDriverConnectionParameter(connectionContext)))); - - TelemetryClientFactory.getInstance() - .getTelemetryClient(connectionContext, DatabricksThreadContextHolder.getDatabricksConfig()) - .exportEvent(telemetryFrontendLog); + // Though we already handle null connectionContext in the downstream implementation, + // we are adding this check for extra sanity + if (connectionContext != null) { + TelemetryFrontendLog telemetryFrontendLog = + new TelemetryFrontendLog() + .setFrontendLogEventId(getEventUUID()) + .setContext(getLogContext()) + .setEntry( + new FrontendLogEntry() + .setSqlDriverLog( + new TelemetryEvent() + .setLatency(latencyMilliseconds) + .setVolumeOperation(volumeOperationEvent) + .setDriverConnectionParameters( + getDriverConnectionParameter(connectionContext)))); + + TelemetryClientFactory.getInstance() + .getTelemetryClient(connectionContext, DatabricksThreadContextHolder.getDatabricksConfig()) + .exportEvent(telemetryFrontendLog); + } } private static DriverConnectionParameters getDriverConnectionParameter( From 99c83f841ad73a8c5b97aa46af11848dc00b859c Mon Sep 17 00:00:00 2001 From: Gopal Lal Date: Fri, 25 Apr 2025 12:59:57 +0530 Subject: [PATCH 4/5] . --- .../jdbc/telemetry/TelemetryHelper.java | 48 ++++++++++--------- 1 file changed, 25 insertions(+), 23 deletions(-) diff --git a/src/main/java/com/databricks/jdbc/telemetry/TelemetryHelper.java b/src/main/java/com/databricks/jdbc/telemetry/TelemetryHelper.java index ce837157fa..0640ac04ac 100644 --- a/src/main/java/com/databricks/jdbc/telemetry/TelemetryHelper.java +++ b/src/main/java/com/databricks/jdbc/telemetry/TelemetryHelper.java @@ -117,21 +117,22 @@ static void exportLatencyLog( // we are adding this check for extra sanity if (connectionContext != null) { TelemetryEvent telemetryEvent = - new TelemetryEvent() - .setLatency(latencyMilliseconds) - .setSqlOperation(executionEvent) - .setDriverConnectionParameters(getDriverConnectionParameter(connectionContext)); + new TelemetryEvent() + .setLatency(latencyMilliseconds) + .setSqlOperation(executionEvent) + .setDriverConnectionParameters(getDriverConnectionParameter(connectionContext)); if (statementId != null) { telemetryEvent.setSqlStatementId(statementId.toString()); } TelemetryFrontendLog telemetryFrontendLog = - new TelemetryFrontendLog() - .setFrontendLogEventId(getEventUUID()) - .setContext(getLogContext()) - .setEntry(new FrontendLogEntry().setSqlDriverLog(telemetryEvent)); + new TelemetryFrontendLog() + .setFrontendLogEventId(getEventUUID()) + .setContext(getLogContext()) + .setEntry(new FrontendLogEntry().setSqlDriverLog(telemetryEvent)); TelemetryClientFactory.getInstance() - .getTelemetryClient(connectionContext, DatabricksThreadContextHolder.getDatabricksConfig()) - .exportEvent(telemetryFrontendLog); + .getTelemetryClient( + connectionContext, DatabricksThreadContextHolder.getDatabricksConfig()) + .exportEvent(telemetryFrontendLog); } } @@ -143,21 +144,22 @@ public static void exportLatencyLog( // we are adding this check for extra sanity if (connectionContext != null) { TelemetryFrontendLog telemetryFrontendLog = - new TelemetryFrontendLog() - .setFrontendLogEventId(getEventUUID()) - .setContext(getLogContext()) - .setEntry( - new FrontendLogEntry() - .setSqlDriverLog( - new TelemetryEvent() - .setLatency(latencyMilliseconds) - .setVolumeOperation(volumeOperationEvent) - .setDriverConnectionParameters( - getDriverConnectionParameter(connectionContext)))); + new TelemetryFrontendLog() + .setFrontendLogEventId(getEventUUID()) + .setContext(getLogContext()) + .setEntry( + new FrontendLogEntry() + .setSqlDriverLog( + new TelemetryEvent() + .setLatency(latencyMilliseconds) + .setVolumeOperation(volumeOperationEvent) + .setDriverConnectionParameters( + getDriverConnectionParameter(connectionContext)))); TelemetryClientFactory.getInstance() - .getTelemetryClient(connectionContext, DatabricksThreadContextHolder.getDatabricksConfig()) - .exportEvent(telemetryFrontendLog); + .getTelemetryClient( + connectionContext, DatabricksThreadContextHolder.getDatabricksConfig()) + .exportEvent(telemetryFrontendLog); } } From 3d0389329e9962cde0f5aa5fefa6b378ca6da7aa Mon Sep 17 00:00:00 2001 From: Gopal Lal Date: Fri, 25 Apr 2025 14:28:08 +0530 Subject: [PATCH 5/5] Add extra null check --- .../com/databricks/jdbc/telemetry/TelemetryClientFactory.java | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/main/java/com/databricks/jdbc/telemetry/TelemetryClientFactory.java b/src/main/java/com/databricks/jdbc/telemetry/TelemetryClientFactory.java index 482faa2b8d..f61b5140f8 100644 --- a/src/main/java/com/databricks/jdbc/telemetry/TelemetryClientFactory.java +++ b/src/main/java/com/databricks/jdbc/telemetry/TelemetryClientFactory.java @@ -34,7 +34,9 @@ public static TelemetryClientFactory getInstance() { public ITelemetryClient getTelemetryClient( IDatabricksConnectionContext connectionContext, DatabricksConfig databricksConfig) { - if (connectionContext != null && connectionContext.isTelemetryEnabled()) { + if (connectionContext != null + && connectionContext.isTelemetryEnabled() + && databricksConfig != null) { return telemetryClients.computeIfAbsent( connectionContext.getConnectionUuid(), k ->