From 4af6d6b54eff5f0cdb44d7d44cade903bf35681f Mon Sep 17 00:00:00 2001 From: Shivam Raj Date: Thu, 22 May 2025 22:06:18 +0530 Subject: [PATCH 01/11] added application name to track BiTool --- src/main/java/com/databricks/client/jdbc/Driver.java | 1 + .../jdbc/api/impl/DatabricksConnectionContext.java | 5 +++++ .../com/databricks/jdbc/api/impl/DatabricksSession.java | 4 ++++ .../jdbc/api/internal/IDatabricksConnectionContext.java | 3 +++ .../databricks/jdbc/common/DatabricksJdbcConstants.java | 3 ++- .../databricks/jdbc/common/DatabricksJdbcUrlParams.java | 4 +++- .../com/databricks/jdbc/telemetry/TelemetryHelper.java | 9 ++++++++- 7 files changed, 26 insertions(+), 3 deletions(-) diff --git a/src/main/java/com/databricks/client/jdbc/Driver.java b/src/main/java/com/databricks/client/jdbc/Driver.java index e6907ec60d..bb0ac82874 100644 --- a/src/main/java/com/databricks/client/jdbc/Driver.java +++ b/src/main/java/com/databricks/client/jdbc/Driver.java @@ -58,6 +58,7 @@ public Connection connect(String url, Properties info) throws DatabricksSQLExcep DatabricksConnectionContextFactory.create(url, info); DriverUtil.setUpLogging(connectionContext); UserAgentManager.setUserAgent(connectionContext); + updateClientAppName(connectionContext.getApplicationName()); LOGGER.info(getDriverSystemConfiguration().toString()); DatabricksConnection connection = new DatabricksConnection(connectionContext); boolean isConnectionOpen = false; 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 7cc3a87aa7..f9c045a9bd 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java @@ -887,6 +887,11 @@ public boolean isTokenCacheEnabled() { return getParameter(DatabricksJdbcUrlParams.ENABLE_TOKEN_CACHE).equals("1"); } + @Override + public String getApplicationName() { + return getParameter(DatabricksJdbcUrlParams.APPLICATION_NAME); + } + private static boolean nullOrEmptyString(String s) { return s == null || s.isEmpty(); } diff --git a/src/main/java/com/databricks/jdbc/api/impl/DatabricksSession.java b/src/main/java/com/databricks/jdbc/api/impl/DatabricksSession.java index d0169a32c5..e71ede7ce9 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksSession.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksSession.java @@ -1,6 +1,7 @@ package com.databricks.jdbc.api.impl; import static com.databricks.jdbc.common.DatabricksJdbcConstants.REDACTED_TOKEN; +import static com.databricks.jdbc.telemetry.TelemetryHelper.updateClientAppName; import com.databricks.jdbc.api.internal.IDatabricksConnectionContext; import com.databricks.jdbc.api.internal.IDatabricksSession; @@ -247,6 +248,9 @@ public void setClientInfoProperty(String name, String value) { this.databricksClient.resetAccessToken(value); value = REDACTED_TOKEN; // mask access token } + if (name.equalsIgnoreCase(DatabricksJdbcUrlParams.APPLICATION_NAME.getParamName())) { + updateClientAppName(value); + } clientInfoProperties.put(name, value); } 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 1b61c02bbd..fa7638cc53 100644 --- a/src/main/java/com/databricks/jdbc/api/internal/IDatabricksConnectionContext.java +++ b/src/main/java/com/databricks/jdbc/api/internal/IDatabricksConnectionContext.java @@ -317,4 +317,7 @@ public interface IDatabricksConnectionContext { /** Returns whether token caching is enabled for OAuth authentication */ boolean isTokenCacheEnabled(); + + /** Returns the application name using JDBC Connection */ + String getApplicationName(); } diff --git a/src/main/java/com/databricks/jdbc/common/DatabricksJdbcConstants.java b/src/main/java/com/databricks/jdbc/common/DatabricksJdbcConstants.java index 46dcbb0feb..6aaef36a85 100644 --- a/src/main/java/com/databricks/jdbc/common/DatabricksJdbcConstants.java +++ b/src/main/java/com/databricks/jdbc/common/DatabricksJdbcConstants.java @@ -81,7 +81,8 @@ public final class DatabricksJdbcConstants { Set.of( ALLOWED_VOLUME_INGESTION_PATHS, ALLOWED_STAGING_INGESTION_PATHS, - DatabricksJdbcUrlParams.AUTH_ACCESS_TOKEN.getParamName()); + DatabricksJdbcUrlParams.AUTH_ACCESS_TOKEN.getParamName(), + DatabricksJdbcUrlParams.APPLICATION_NAME.getParamName()); public static final Map JSON_HTTP_HEADERS = Map.of( "Accept", "application/json", diff --git a/src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java b/src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java index fa445c3451..3e7cac2c9c 100644 --- a/src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java +++ b/src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java @@ -127,7 +127,9 @@ public enum DatabricksJdbcUrlParams { "255"), SOCKET_TIMEOUT("socketTimeout", "Socket timeout in seconds", "900"), TOKEN_CACHE_PASS_PHRASE("TokenCachePassPhrase", "Pass phrase to use for OAuth U2M Token Cache"), - ENABLE_TOKEN_CACHE("EnableTokenCache", "Enable caching OAuth tokens", "1"); + ENABLE_TOKEN_CACHE("EnableTokenCache", "Enable caching OAuth tokens", "1"), + APPLICATION_NAME("ApplicationName", "Name of application using the driver", ""), + ; private final String paramName; private final String defaultValue; diff --git a/src/main/java/com/databricks/jdbc/telemetry/TelemetryHelper.java b/src/main/java/com/databricks/jdbc/telemetry/TelemetryHelper.java index 8e6407eac8..4896396eeb 100644 --- a/src/main/java/com/databricks/jdbc/telemetry/TelemetryHelper.java +++ b/src/main/java/com/databricks/jdbc/telemetry/TelemetryHelper.java @@ -1,5 +1,7 @@ package com.databricks.jdbc.telemetry; +import static com.databricks.jdbc.common.util.WildcardUtil.isNullOrEmpty; + import com.databricks.jdbc.api.internal.IDatabricksConnectionContext; import com.databricks.jdbc.common.util.DatabricksThreadContextHolder; import com.databricks.jdbc.common.util.DriverUtil; @@ -22,7 +24,6 @@ public class TelemetryHelper { private static final DriverSystemConfiguration DRIVER_SYSTEM_CONFIGURATION = new DriverSystemConfiguration() - .setClientAppName(null) .setCharSetEncoding(Charset.defaultCharset().displayName()) .setDriverName(DriverUtil.getDriverName()) .setDriverVersion(DriverUtil.getDriverVersion()) @@ -40,6 +41,12 @@ public static DriverSystemConfiguration getDriverSystemConfiguration() { return DRIVER_SYSTEM_CONFIGURATION; } + public static void updateClientAppName(String clientAppName) { + if (!isNullOrEmpty(clientAppName)) { + DRIVER_SYSTEM_CONFIGURATION.setClientAppName(clientAppName); + } + } + // TODO : add an export even before connection context is built public static void exportInitialTelemetryLog(IDatabricksConnectionContext connectionContext) { if (connectionContext == null) { From bc5260831a0eabe5b17ca302e1b0f3429f5783a1 Mon Sep 17 00:00:00 2001 From: Shivam Raj Date: Thu, 22 May 2025 22:37:23 +0530 Subject: [PATCH 02/11] add tests --- .../telemetry/DriverSystemConfiguration.java | 6 +++++ .../jdbc/api/impl/DatabricksSessionTest.java | 11 ++++++++++ .../jdbc/telemetry/TelemetryHelperTest.java | 22 +++++++++++++++++++ 3 files changed, 39 insertions(+) diff --git a/src/main/java/com/databricks/jdbc/model/telemetry/DriverSystemConfiguration.java b/src/main/java/com/databricks/jdbc/model/telemetry/DriverSystemConfiguration.java index 606850a7c0..2e6bfb55d2 100644 --- a/src/main/java/com/databricks/jdbc/model/telemetry/DriverSystemConfiguration.java +++ b/src/main/java/com/databricks/jdbc/model/telemetry/DriverSystemConfiguration.java @@ -2,6 +2,7 @@ import com.databricks.sdk.support.ToStringer; import com.fasterxml.jackson.annotation.JsonProperty; +import com.google.common.annotations.VisibleForTesting; public class DriverSystemConfiguration { // TODO : add json properties when proto is implemented completely @@ -38,6 +39,11 @@ public class DriverSystemConfiguration { @JsonProperty("char_set_encoding") private String charSetEncoding; + @VisibleForTesting + public String getClientAppName() { + return clientAppName; + } + public DriverSystemConfiguration setDriverName(String driverName) { this.driverName = driverName; return this; diff --git a/src/test/java/com/databricks/jdbc/api/impl/DatabricksSessionTest.java b/src/test/java/com/databricks/jdbc/api/impl/DatabricksSessionTest.java index f4b2ae02bc..c4933208fd 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/DatabricksSessionTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/DatabricksSessionTest.java @@ -2,6 +2,7 @@ import static com.databricks.jdbc.TestConstants.*; import static com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode.TEMPORARY_REDIRECT_EXCEPTION; +import static com.databricks.jdbc.telemetry.TelemetryHelper.getDriverSystemConfiguration; import static org.junit.jupiter.api.Assertions.*; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; @@ -201,6 +202,16 @@ public void testSetClientInfoProperty() throws DatabricksSQLException { assertEquals("value", session.getClientInfoProperties().get("key")); } + @Test + public void testGetClientInfoProperty_ApplicationName() throws DatabricksSQLException { + DatabricksSession session = + new DatabricksSession( + DatabricksConnectionContext.parse(VALID_CLUSTER_URL, new Properties()), sdkClient); + session.setClientInfoProperty("applicationName", "testApp"); + assertEquals("testApp", session.getClientInfoProperties().get("applicationName")); + assertEquals("testApp", getDriverSystemConfiguration().getClientAppName()); + } + @Test public void testSetClientInfoProperty_AuthAccessToken() throws DatabricksSQLException { DatabricksSession session = diff --git a/src/test/java/com/databricks/jdbc/telemetry/TelemetryHelperTest.java b/src/test/java/com/databricks/jdbc/telemetry/TelemetryHelperTest.java index d5d106ca43..a8ad2dc1c2 100644 --- a/src/test/java/com/databricks/jdbc/telemetry/TelemetryHelperTest.java +++ b/src/test/java/com/databricks/jdbc/telemetry/TelemetryHelperTest.java @@ -2,6 +2,7 @@ import static com.databricks.jdbc.TestConstants.TEST_STRING; import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertEquals; import static org.mockito.Mockito.*; import com.databricks.jdbc.api.internal.IDatabricksConnectionContext; @@ -68,4 +69,25 @@ void testErrorTelemetryLogDoesNotThrowError() { void testGetDriverSystemConfigurationDoesNotThrowError() { assertDoesNotThrow(TelemetryHelper::getDriverSystemConfiguration); } + + @Test + void testUpdateClientAppName() { + // Set the client app name to null first to ensure a clean state + TelemetryHelper.updateClientAppName(null); + + // Test valid app name + TelemetryHelper.updateClientAppName("TestApplicationName"); + assertEquals( + "TestApplicationName", TelemetryHelper.getDriverSystemConfiguration().getClientAppName()); + + // Test empty app name - should not change the existing value + TelemetryHelper.updateClientAppName(""); + assertEquals( + "TestApplicationName", TelemetryHelper.getDriverSystemConfiguration().getClientAppName()); + + // Test null app name - should not change the existing value + TelemetryHelper.updateClientAppName(null); + assertEquals( + "TestApplicationName", TelemetryHelper.getDriverSystemConfiguration().getClientAppName()); + } } From c28d24ce68f1949fc692f11e034444e67ad0222b Mon Sep 17 00:00:00 2001 From: Shivam Raj Date: Fri, 23 May 2025 14:46:29 +0530 Subject: [PATCH 03/11] some changes --- .../jdbc/api/impl/DatabricksConnectionContext.java | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) 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 f9c045a9bd..da7093f7e6 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java @@ -5,11 +5,11 @@ import static com.databricks.jdbc.common.EnvironmentVariables.DEFAULT_ROW_LIMIT_PER_BLOCK; import static com.databricks.jdbc.common.util.UserAgentManager.USER_AGENT_SEA_CLIENT; import static com.databricks.jdbc.common.util.UserAgentManager.USER_AGENT_THRIFT_CLIENT; +import static com.databricks.jdbc.common.util.WildcardUtil.isNullOrEmpty; import com.databricks.jdbc.api.internal.IDatabricksConnectionContext; import com.databricks.jdbc.common.*; import com.databricks.jdbc.common.util.ValidationUtil; -import com.databricks.jdbc.common.util.WildcardUtil; import com.databricks.jdbc.exception.DatabricksDriverException; import com.databricks.jdbc.exception.DatabricksParsingException; import com.databricks.jdbc.exception.DatabricksSQLException; @@ -83,7 +83,7 @@ public static ImmutableMap buildPropertiesMap( String connectionParamString, Properties properties) { ImmutableMap.Builder parametersBuilder = ImmutableMap.builder(); // check if connectionParamString is empty or null - if (!WildcardUtil.isNullOrEmpty(connectionParamString)) { + if (!isNullOrEmpty(connectionParamString)) { String[] urlParts = connectionParamString.split(DatabricksJdbcConstants.URL_DELIMITER); for (String urlPart : urlParts) { String[] pair = urlPart.split(DatabricksJdbcConstants.PAIR_DELIMITER); @@ -380,7 +380,11 @@ public String getClientUserAgent() { } public String getCustomerUserAgent() { - return getParameter(DatabricksJdbcUrlParams.USER_AGENT_ENTRY); + String customerUA = getParameter(DatabricksJdbcUrlParams.USER_AGENT_ENTRY); + if (isNullOrEmpty(customerUA)) { + customerUA = getParameter(DatabricksJdbcUrlParams.APPLICATION_NAME); + } + return customerUA; } @Override From 764ede9f0b985cde23ed10c2aa8c2464121dd112 Mon Sep 17 00:00:00 2001 From: Shivam Raj Date: Fri, 23 May 2025 18:32:33 +0530 Subject: [PATCH 04/11] refactored code; update user agent and telemetry client app --- .../com/databricks/client/jdbc/Driver.java | 2 +- .../api/impl/DatabricksConnectionContext.java | 7 +- .../jdbc/api/impl/DatabricksSession.java | 7 +- .../jdbc/common/util/UserAgentHelper.java | 66 +++++++++++ .../jdbc/common/util/UserAgentManager.java | 15 +-- .../jdbc/common/util/UserAgentHelperTest.java | 109 ++++++++++++++++++ 6 files changed, 187 insertions(+), 19 deletions(-) create mode 100644 src/main/java/com/databricks/jdbc/common/util/UserAgentHelper.java create mode 100644 src/test/java/com/databricks/jdbc/common/util/UserAgentHelperTest.java diff --git a/src/main/java/com/databricks/client/jdbc/Driver.java b/src/main/java/com/databricks/client/jdbc/Driver.java index bb0ac82874..89cfe2049c 100644 --- a/src/main/java/com/databricks/client/jdbc/Driver.java +++ b/src/main/java/com/databricks/client/jdbc/Driver.java @@ -57,8 +57,8 @@ public Connection connect(String url, Properties info) throws DatabricksSQLExcep IDatabricksConnectionContext connectionContext = DatabricksConnectionContextFactory.create(url, info); DriverUtil.setUpLogging(connectionContext); + // This initializes both user agent headers and telemetry UserAgentManager.setUserAgent(connectionContext); - updateClientAppName(connectionContext.getApplicationName()); LOGGER.info(getDriverSystemConfiguration().toString()); DatabricksConnection connection = new DatabricksConnection(connectionContext); boolean isConnectionOpen = false; 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 da7093f7e6..4e24a371c4 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java @@ -379,12 +379,9 @@ public String getClientUserAgent() { : USER_AGENT_THRIFT_CLIENT; } + @Override public String getCustomerUserAgent() { - String customerUA = getParameter(DatabricksJdbcUrlParams.USER_AGENT_ENTRY); - if (isNullOrEmpty(customerUA)) { - customerUA = getParameter(DatabricksJdbcUrlParams.APPLICATION_NAME); - } - return customerUA; + return getParameter(DatabricksJdbcUrlParams.USER_AGENT_ENTRY); } @Override diff --git a/src/main/java/com/databricks/jdbc/api/impl/DatabricksSession.java b/src/main/java/com/databricks/jdbc/api/impl/DatabricksSession.java index e71ede7ce9..5226238b5e 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksSession.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksSession.java @@ -1,7 +1,6 @@ package com.databricks.jdbc.api.impl; import static com.databricks.jdbc.common.DatabricksJdbcConstants.REDACTED_TOKEN; -import static com.databricks.jdbc.telemetry.TelemetryHelper.updateClientAppName; import com.databricks.jdbc.api.internal.IDatabricksConnectionContext; import com.databricks.jdbc.api.internal.IDatabricksSession; @@ -9,6 +8,7 @@ import com.databricks.jdbc.common.DatabricksClientType; import com.databricks.jdbc.common.DatabricksJdbcUrlParams; import com.databricks.jdbc.common.IDatabricksComputeResource; +import com.databricks.jdbc.common.util.UserAgentHelper; import com.databricks.jdbc.dbclient.IDatabricksClient; import com.databricks.jdbc.dbclient.IDatabricksMetadataClient; import com.databricks.jdbc.dbclient.impl.sqlexec.DatabricksEmptyMetadataClient; @@ -248,9 +248,12 @@ public void setClientInfoProperty(String name, String value) { this.databricksClient.resetAccessToken(value); value = REDACTED_TOKEN; // mask access token } + + // If application name is being set, update both telemetry and user agent if (name.equalsIgnoreCase(DatabricksJdbcUrlParams.APPLICATION_NAME.getParamName())) { - updateClientAppName(value); + UserAgentHelper.updateUserAgentAndTelemetry(connectionContext, value); } + clientInfoProperties.put(name, value); } diff --git a/src/main/java/com/databricks/jdbc/common/util/UserAgentHelper.java b/src/main/java/com/databricks/jdbc/common/util/UserAgentHelper.java new file mode 100644 index 0000000000..f61b5f1a1e --- /dev/null +++ b/src/main/java/com/databricks/jdbc/common/util/UserAgentHelper.java @@ -0,0 +1,66 @@ +package com.databricks.jdbc.common.util; + +import static com.databricks.jdbc.common.util.WildcardUtil.isNullOrEmpty; + +import com.databricks.jdbc.api.internal.IDatabricksConnectionContext; +import com.databricks.jdbc.telemetry.TelemetryHelper; +import com.databricks.sdk.core.UserAgent; + +/** Helper class for determining and setting user agent and client app name. */ +public class UserAgentHelper { + private static final String APP_NAME_SYSTEM_PROPERTY = "app.name"; + private static final String VERSION_FILLER = "version"; + + /** + * Determines the application name using a fallback mechanism: 1. useragententry url param 2. + * applicationname url param 3. client info property "applicationname" 4. System property app.name + * + * @param connectionContext The connection context + * @param clientInfoAppName The application name from client info properties, can be null + * @return The determined application name or null if none is found + */ + public static String determineApplicationName( + IDatabricksConnectionContext connectionContext, String clientInfoAppName) { + // First check URL params + String appName = connectionContext.getCustomerUserAgent(); + if (!isNullOrEmpty(appName)) { + return appName; + } + + // Then check applicationname URL param + appName = connectionContext.getApplicationName(); + if (!isNullOrEmpty(appName)) { + return appName; + } + + // Then check client info property + if (!isNullOrEmpty(clientInfoAppName)) { + return clientInfoAppName; + } + + // Finally check system property + return System.getProperty(APP_NAME_SYSTEM_PROPERTY); + } + + /** + * Updates both the telemetry client app name and HTTP user agent headers. To be called during + * connection initialization and when app name changes. + * + * @param connectionContext The connection context + * @param clientInfoAppName Optional client info app name, can be null + */ + public static void updateUserAgentAndTelemetry( + IDatabricksConnectionContext connectionContext, String clientInfoAppName) { + String appName = determineApplicationName(connectionContext, clientInfoAppName); + if (!isNullOrEmpty(appName)) { + // Update telemetry + TelemetryHelper.updateClientAppName(appName); + + // Update HTTP user agent + int i = appName.indexOf('/'); + String customerName = (i < 0) ? appName : appName.substring(0, i); + String customerVersion = (i < 0) ? VERSION_FILLER : appName.substring(i + 1); + UserAgent.withOtherInfo(customerName, UserAgent.sanitize(customerVersion)); + } + } +} diff --git a/src/main/java/com/databricks/jdbc/common/util/UserAgentManager.java b/src/main/java/com/databricks/jdbc/common/util/UserAgentManager.java index 17ee19088d..433d61c1db 100644 --- a/src/main/java/com/databricks/jdbc/common/util/UserAgentManager.java +++ b/src/main/java/com/databricks/jdbc/common/util/UserAgentManager.java @@ -1,14 +1,11 @@ package com.databricks.jdbc.common.util; -import static com.databricks.jdbc.common.util.WildcardUtil.isNullOrEmpty; - import com.databricks.jdbc.api.internal.IDatabricksConnectionContext; import com.databricks.sdk.core.UserAgent; public class UserAgentManager { private static final String SDK_USER_AGENT = "databricks-sdk-java"; private static final String JDBC_HTTP_USER_AGENT = "databricks-jdbc-http"; - private static final String VERSION_FILLER_FOR_CUSTOMER_USER_AGENT = "version"; private static final String DEFAULT_USER_AGENT = "DatabricksJDBCDriverOSS"; private static final String CLIENT_USER_AGENT_PREFIX = "Java"; public static final String USER_AGENT_SEA_CLIENT = "SQLExecHttpClient"; @@ -20,16 +17,12 @@ public class UserAgentManager { * @param connectionContext The connection context. */ public static void setUserAgent(IDatabricksConnectionContext connectionContext) { + // Set the base product and client info UserAgent.withProduct(DEFAULT_USER_AGENT, DriverUtil.getDriverVersion()); UserAgent.withOtherInfo(CLIENT_USER_AGENT_PREFIX, connectionContext.getClientUserAgent()); - String customerUA = connectionContext.getCustomerUserAgent(); - if (!isNullOrEmpty(customerUA)) { - int i = customerUA.indexOf('/'); - String customerName = (i < 0) ? customerUA : customerUA.substring(0, i); - String customerVersion = - (i < 0) ? VERSION_FILLER_FOR_CUSTOMER_USER_AGENT : customerUA.substring(i + 1); - UserAgent.withOtherInfo(customerName, UserAgent.sanitize(customerVersion)); - } + + // Let UserAgentHelper handle the application name for both telemetry and user agent + UserAgentHelper.updateUserAgentAndTelemetry(connectionContext, null); } /** Gets the user agent string for Databricks Driver HTTP Client. */ diff --git a/src/test/java/com/databricks/jdbc/common/util/UserAgentHelperTest.java b/src/test/java/com/databricks/jdbc/common/util/UserAgentHelperTest.java new file mode 100644 index 0000000000..280a944631 --- /dev/null +++ b/src/test/java/com/databricks/jdbc/common/util/UserAgentHelperTest.java @@ -0,0 +1,109 @@ +package com.databricks.jdbc.common.util; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.mockito.Mockito.when; + +import com.databricks.jdbc.api.internal.IDatabricksConnectionContext; +import com.databricks.jdbc.telemetry.TelemetryHelper; +import com.databricks.sdk.core.UserAgent; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.MockedStatic; +import org.mockito.Mockito; +import org.mockito.junit.jupiter.MockitoExtension; + +@ExtendWith(MockitoExtension.class) +public class UserAgentHelperTest { + + @Mock private IDatabricksConnectionContext connectionContext; + + private MockedStatic telemetryHelperMock; + private MockedStatic userAgentMock; + + @BeforeEach + public void setup() { + telemetryHelperMock = Mockito.mockStatic(TelemetryHelper.class); + userAgentMock = Mockito.mockStatic(UserAgent.class); + } + + @AfterEach + public void tearDown() { + telemetryHelperMock.close(); + userAgentMock.close(); + } + + @Test + public void testDetermineApplicationName_WithUserAgentEntry() { + // When useragententry is set + when(connectionContext.getCustomerUserAgent()).thenReturn("MyUserAgent"); + + String result = + UserAgentHelper.determineApplicationName(connectionContext, "AlsoShouldNotUseThis"); + + assertEquals("MyUserAgent", result); + } + + @Test + public void testDetermineApplicationName_WithApplicationName() { + // When useragententry is not set but applicationname is + when(connectionContext.getCustomerUserAgent()).thenReturn(null); + when(connectionContext.getApplicationName()).thenReturn("AppNameValue"); + + String result = UserAgentHelper.determineApplicationName(connectionContext, "ShouldNotUseThis"); + + assertEquals("AppNameValue", result); + } + + @Test + public void testDetermineApplicationName_WithClientInfo() { + // When URL params are not set but client info is provided + when(connectionContext.getCustomerUserAgent()).thenReturn(null); + when(connectionContext.getApplicationName()).thenReturn(null); + + String result = UserAgentHelper.determineApplicationName(connectionContext, "ClientInfoApp"); + + assertEquals("ClientInfoApp", result); + } + + @Test + public void testDetermineApplicationName_WithSystemProperty() { + // When falling back to system property + when(connectionContext.getCustomerUserAgent()).thenReturn(null); + when(connectionContext.getApplicationName()).thenReturn(null); + + System.setProperty("app.name", "SystemPropApp"); + try { + String result = UserAgentHelper.determineApplicationName(connectionContext, null); + assertEquals("SystemPropApp", result); + } finally { + System.clearProperty("app.name"); + } + } + + @Test + public void testUpdateUserAgentAndTelemetry() { + // Test that both telemetry and user agent are updated + when(connectionContext.getCustomerUserAgent()).thenReturn("TestApp"); + when(UserAgent.sanitize("version")).thenReturn("version"); + + UserAgentHelper.updateUserAgentAndTelemetry(connectionContext, null); + + telemetryHelperMock.verify(() -> TelemetryHelper.updateClientAppName("TestApp")); + userAgentMock.verify(() -> UserAgent.withOtherInfo("TestApp", "version")); + } + + @Test + public void testUpdateUserAgentAndTelemetry_WithVersion() { + // Test with app name containing version + when(connectionContext.getCustomerUserAgent()).thenReturn("MyApp/1.2.3"); + when(UserAgent.sanitize("1.2.3")).thenReturn("1.2.3"); + + UserAgentHelper.updateUserAgentAndTelemetry(connectionContext, null); + + telemetryHelperMock.verify(() -> TelemetryHelper.updateClientAppName("MyApp/1.2.3")); + userAgentMock.verify(() -> UserAgent.withOtherInfo("MyApp", "1.2.3")); + } +} From e737325cbb54dedf35c118aff67bfdd151454d0b Mon Sep 17 00:00:00 2001 From: Shivam Raj Date: Mon, 26 May 2025 14:57:57 +0530 Subject: [PATCH 05/11] update user agent thrift mode as well --- .../jdbc/api/impl/DatabricksSession.java | 2 ++ .../jdbc/dbclient/IDatabricksClient.java | 3 +++ .../impl/sqlexec/DatabricksSdkClient.java | 6 +++++ .../impl/thrift/DatabricksHttpTTransport.java | 18 +++++++++++++-- .../thrift/DatabricksThriftServiceClient.java | 6 +++++ .../jdbc/api/impl/DatabricksSessionTest.java | 22 +++++++++++++++++++ 6 files changed, 55 insertions(+), 2 deletions(-) diff --git a/src/main/java/com/databricks/jdbc/api/impl/DatabricksSession.java b/src/main/java/com/databricks/jdbc/api/impl/DatabricksSession.java index 5226238b5e..1f8d9b33c8 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksSession.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksSession.java @@ -252,6 +252,8 @@ public void setClientInfoProperty(String name, String value) { // If application name is being set, update both telemetry and user agent if (name.equalsIgnoreCase(DatabricksJdbcUrlParams.APPLICATION_NAME.getParamName())) { UserAgentHelper.updateUserAgentAndTelemetry(connectionContext, value); + // Update the client's user agent headers for existing ThriftClient + databricksClient.updateUserAgent(); } clientInfoProperties.put(name, value); diff --git a/src/main/java/com/databricks/jdbc/dbclient/IDatabricksClient.java b/src/main/java/com/databricks/jdbc/dbclient/IDatabricksClient.java index 14c1890bf9..c7965e5ad7 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/IDatabricksClient.java +++ b/src/main/java/com/databricks/jdbc/dbclient/IDatabricksClient.java @@ -126,6 +126,9 @@ Collection getResultChunks(StatementId statementId, long chunkInde */ void resetAccessToken(String newAccessToken); + /** Update the user agent headers when application name changes */ + void updateUserAgent(); + TFetchResultsResp getMoreResults(IDatabricksStatementInternal parentStatement) throws DatabricksSQLException; diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClient.java b/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClient.java index 242d1bf162..54417b89fa 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClient.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClient.java @@ -412,6 +412,12 @@ public synchronized void resetAccessToken(String newAccessToken) { this.apiClient = workspaceClient.apiClient(); } + @Override + public void updateUserAgent() { + // For SDK client, the user agent is handled by the global UserAgent state + // which is automatically picked up by the SDK apiClient for new requests + } + @Override public TFetchResultsResp getMoreResults(IDatabricksStatementInternal parentStatement) throws DatabricksSQLException { diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksHttpTTransport.java b/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksHttpTTransport.java index e4c239bba3..b6ecb758ee 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksHttpTTransport.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksHttpTTransport.java @@ -3,6 +3,7 @@ import static com.databricks.jdbc.common.util.DatabricksAuthUtil.initializeConfigWithToken; import com.databricks.jdbc.api.internal.IDatabricksConnectionContext; +import com.databricks.jdbc.common.util.UserAgentManager; import com.databricks.jdbc.common.util.ValidationUtil; import com.databricks.jdbc.dbclient.IDatabricksHttpClient; import com.databricks.jdbc.dbclient.impl.common.TracingUtil; @@ -145,8 +146,12 @@ public void checkReadBytesAvailable(long numBytes) throws TTransportException {} /** Refreshes the custom headers by re-authenticating if necessary. */ private void refreshHeadersIfRequired() { Map refreshedHeaders = databricksConfig.authenticate(); - customHeaders = - refreshedHeaders != null ? new HashMap<>(refreshedHeaders) : Collections.emptyMap(); + // Preserve existing custom headers (like User-Agent) and add/update auth headers + Map newHeaders = new HashMap<>(customHeaders); + if (refreshedHeaders != null) { + newHeaders.putAll(refreshedHeaders); + } + customHeaders = newHeaders; } void resetAccessToken(String newAccessToken) { @@ -154,6 +159,15 @@ void resetAccessToken(String newAccessToken) { this.databricksConfig.resolve(); } + /** Updates the user agent for thrift requests by adding it as a header */ + void updateUserAgent() { + // Add user agent as a header since the HTTP client's user agent is set during construction + // and cannot be changed. This ensures the updated user agent is sent with thrift requests. + Map headers = new HashMap<>(customHeaders); + headers.put("User-Agent", UserAgentManager.getUserAgentString()); + customHeaders = headers; + } + @VisibleForTesting void setResponseBuffer(ByteArrayInputStream responseBuffer) { this.responseBuffer = responseBuffer; diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftServiceClient.java b/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftServiceClient.java index 76ce3cd3bd..ca912b0efe 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftServiceClient.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftServiceClient.java @@ -76,6 +76,12 @@ public void resetAccessToken(String newAccessToken) { .resetAccessToken(newAccessToken); } + @Override + public void updateUserAgent() { + ((DatabricksHttpTTransport) thriftAccessor.getThriftClient().getInputProtocol().getTransport()) + .updateUserAgent(); + } + @Override public ImmutableSessionInfo createSession( IDatabricksComputeResource cluster, diff --git a/src/test/java/com/databricks/jdbc/api/impl/DatabricksSessionTest.java b/src/test/java/com/databricks/jdbc/api/impl/DatabricksSessionTest.java index c4933208fd..241e6ec0e7 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/DatabricksSessionTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/DatabricksSessionTest.java @@ -231,4 +231,26 @@ public void testSetClientInfoProperty_AuthAccessTokenThrift() throws DatabricksS DatabricksJdbcUrlParams.AUTH_ACCESS_TOKEN.getParamName(), "token"); verify(thriftClient).resetAccessToken("token"); } + + @Test + public void testSetClientInfoProperty_ApplicationNameUpdatesUserAgent() + throws DatabricksSQLException { + DatabricksSession session = + new DatabricksSession( + DatabricksConnectionContext.parse(VALID_CLUSTER_URL, new Properties()), sdkClient); + session.setClientInfoProperty( + DatabricksJdbcUrlParams.APPLICATION_NAME.getParamName(), "testApp"); + verify(sdkClient).updateUserAgent(); + } + + @Test + public void testSetClientInfoProperty_ApplicationNameUpdatesUserAgentThrift() + throws DatabricksSQLException { + DatabricksSession session = + new DatabricksSession( + DatabricksConnectionContext.parse(VALID_CLUSTER_URL, new Properties()), thriftClient); + session.setClientInfoProperty( + DatabricksJdbcUrlParams.APPLICATION_NAME.getParamName(), "testApp"); + verify(thriftClient).updateUserAgent(); + } } From 8b10ecf2bd31e3fdec9ce215e6413ffcf3a03f3c Mon Sep 17 00:00:00 2001 From: Shivam Raj Date: Mon, 2 Jun 2025 10:55:37 +0530 Subject: [PATCH 06/11] changes fn to package-private. --- .../java/com/databricks/jdbc/common/util/UserAgentHelper.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/main/java/com/databricks/jdbc/common/util/UserAgentHelper.java b/src/main/java/com/databricks/jdbc/common/util/UserAgentHelper.java index f61b5f1a1e..5531d7cb13 100644 --- a/src/main/java/com/databricks/jdbc/common/util/UserAgentHelper.java +++ b/src/main/java/com/databricks/jdbc/common/util/UserAgentHelper.java @@ -19,7 +19,7 @@ public class UserAgentHelper { * @param clientInfoAppName The application name from client info properties, can be null * @return The determined application name or null if none is found */ - public static String determineApplicationName( + static String determineApplicationName( IDatabricksConnectionContext connectionContext, String clientInfoAppName) { // First check URL params String appName = connectionContext.getCustomerUserAgent(); From 1f57dbc690743bb4756ea91229843cdc08d7e870 Mon Sep 17 00:00:00 2001 From: Shivam Raj Date: Mon, 2 Jun 2025 12:44:13 +0530 Subject: [PATCH 07/11] make parameterized test --- .../jdbc/common/util/UserAgentHelperTest.java | 73 +++++++++++-------- 1 file changed, 44 insertions(+), 29 deletions(-) diff --git a/src/test/java/com/databricks/jdbc/common/util/UserAgentHelperTest.java b/src/test/java/com/databricks/jdbc/common/util/UserAgentHelperTest.java index 280a944631..910ef74ca8 100644 --- a/src/test/java/com/databricks/jdbc/common/util/UserAgentHelperTest.java +++ b/src/test/java/com/databricks/jdbc/common/util/UserAgentHelperTest.java @@ -6,10 +6,14 @@ import com.databricks.jdbc.api.internal.IDatabricksConnectionContext; import com.databricks.jdbc.telemetry.TelemetryHelper; import com.databricks.sdk.core.UserAgent; +import java.util.stream.Stream; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; import org.mockito.Mock; import org.mockito.MockedStatic; import org.mockito.Mockito; @@ -35,37 +39,48 @@ public void tearDown() { userAgentMock.close(); } - @Test - public void testDetermineApplicationName_WithUserAgentEntry() { - // When useragententry is set - when(connectionContext.getCustomerUserAgent()).thenReturn("MyUserAgent"); - - String result = - UserAgentHelper.determineApplicationName(connectionContext, "AlsoShouldNotUseThis"); - - assertEquals("MyUserAgent", result); - } - - @Test - public void testDetermineApplicationName_WithApplicationName() { - // When useragententry is not set but applicationname is - when(connectionContext.getCustomerUserAgent()).thenReturn(null); - when(connectionContext.getApplicationName()).thenReturn("AppNameValue"); - - String result = UserAgentHelper.determineApplicationName(connectionContext, "ShouldNotUseThis"); - - assertEquals("AppNameValue", result); + private static Stream provideApplicationNameTestCases() { + return Stream.of( + // Test case 1: UserAgentEntry takes precedence + Arguments.of( + "MyUserAgent", + "AppNameValue", + "ClientInfoApp", + "MyUserAgent", + "When useragententry is set"), + // Test case 2: ApplicationName is used when UserAgentEntry is null + Arguments.of( + null, + "AppNameValue", + "ClientInfoApp", + "AppNameValue", + "When useragententry is not set but applicationname is"), + // Test case 3: ClientInfo is used when both UserAgentEntry and ApplicationName are null + Arguments.of( + null, + null, // applicationName + "ClientInfoApp", + "ClientInfoApp", + "When URL params are not set but client info is provided")); } - @Test - public void testDetermineApplicationName_WithClientInfo() { - // When URL params are not set but client info is provided - when(connectionContext.getCustomerUserAgent()).thenReturn(null); - when(connectionContext.getApplicationName()).thenReturn(null); - - String result = UserAgentHelper.determineApplicationName(connectionContext, "ClientInfoApp"); - - assertEquals("ClientInfoApp", result); + @ParameterizedTest(name = "{4}") + @MethodSource("provideApplicationNameTestCases") + public void testDetermineApplicationName( + String customerUserAgent, + String applicationName, + String clientInfoApp, + String expectedResult, + String testDescription) { + // Setup only necessary stubs + Mockito.lenient().when(connectionContext.getCustomerUserAgent()).thenReturn(customerUserAgent); + Mockito.lenient().when(connectionContext.getApplicationName()).thenReturn(applicationName); + + // Execute + String result = UserAgentHelper.determineApplicationName(connectionContext, clientInfoApp); + + // Verify + assertEquals(expectedResult, result); } @Test From 0408c4a3bd4d61e74542589160104b2509cb999e Mon Sep 17 00:00:00 2001 From: Shivam Raj Date: Mon, 2 Jun 2025 13:49:09 +0530 Subject: [PATCH 08/11] spotless --- .../com/databricks/jdbc/telemetry/TelemetryHelperTest.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/test/java/com/databricks/jdbc/telemetry/TelemetryHelperTest.java b/src/test/java/com/databricks/jdbc/telemetry/TelemetryHelperTest.java index 0c0ad5f930..a58a17d60e 100644 --- a/src/test/java/com/databricks/jdbc/telemetry/TelemetryHelperTest.java +++ b/src/test/java/com/databricks/jdbc/telemetry/TelemetryHelperTest.java @@ -1,12 +1,12 @@ package com.databricks.jdbc.telemetry; import static com.databricks.jdbc.TestConstants.TEST_STRING; -import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; -import static org.junit.jupiter.api.Assertions.assertEquals; import static com.databricks.jdbc.TestConstants.WAREHOUSE_COMPUTE; import static com.databricks.jdbc.common.safe.FeatureFlagTestUtil.enableFeatureFlagForTesting; import static com.databricks.jdbc.telemetry.TelemetryHelper.isTelemetryAllowedForConnection; import static org.junit.jupiter.api.Assertions.*; +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertEquals; import static org.mockito.Mockito.*; import com.databricks.jdbc.api.internal.IDatabricksConnectionContext; From 8043483eb13fc42181c563be875d06d7e260ce0f Mon Sep 17 00:00:00 2001 From: Shivam Raj Date: Mon, 2 Jun 2025 16:06:47 +0530 Subject: [PATCH 09/11] merge UserAgentHelper and UserAgentManager --- .../com/databricks/client/jdbc/Driver.java | 1 - .../jdbc/api/impl/DatabricksSession.java | 4 +- .../jdbc/common/util/UserAgentHelper.java | 66 ---------- .../jdbc/common/util/UserAgentManager.java | 62 ++++++++- .../jdbc/common/util/UserAgentHelperTest.java | 124 ------------------ .../common/util/UserAgentManagerTest.java | 115 ++++++++++++++++ 6 files changed, 177 insertions(+), 195 deletions(-) delete mode 100644 src/main/java/com/databricks/jdbc/common/util/UserAgentHelper.java delete mode 100644 src/test/java/com/databricks/jdbc/common/util/UserAgentHelperTest.java diff --git a/src/main/java/com/databricks/client/jdbc/Driver.java b/src/main/java/com/databricks/client/jdbc/Driver.java index 89cfe2049c..e6907ec60d 100644 --- a/src/main/java/com/databricks/client/jdbc/Driver.java +++ b/src/main/java/com/databricks/client/jdbc/Driver.java @@ -57,7 +57,6 @@ public Connection connect(String url, Properties info) throws DatabricksSQLExcep IDatabricksConnectionContext connectionContext = DatabricksConnectionContextFactory.create(url, info); DriverUtil.setUpLogging(connectionContext); - // This initializes both user agent headers and telemetry UserAgentManager.setUserAgent(connectionContext); LOGGER.info(getDriverSystemConfiguration().toString()); DatabricksConnection connection = new DatabricksConnection(connectionContext); diff --git a/src/main/java/com/databricks/jdbc/api/impl/DatabricksSession.java b/src/main/java/com/databricks/jdbc/api/impl/DatabricksSession.java index 1f8d9b33c8..6a7f97f978 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksSession.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksSession.java @@ -8,7 +8,7 @@ import com.databricks.jdbc.common.DatabricksClientType; import com.databricks.jdbc.common.DatabricksJdbcUrlParams; import com.databricks.jdbc.common.IDatabricksComputeResource; -import com.databricks.jdbc.common.util.UserAgentHelper; +import com.databricks.jdbc.common.util.UserAgentManager; import com.databricks.jdbc.dbclient.IDatabricksClient; import com.databricks.jdbc.dbclient.IDatabricksMetadataClient; import com.databricks.jdbc.dbclient.impl.sqlexec.DatabricksEmptyMetadataClient; @@ -251,7 +251,7 @@ public void setClientInfoProperty(String name, String value) { // If application name is being set, update both telemetry and user agent if (name.equalsIgnoreCase(DatabricksJdbcUrlParams.APPLICATION_NAME.getParamName())) { - UserAgentHelper.updateUserAgentAndTelemetry(connectionContext, value); + UserAgentManager.updateUserAgentAndTelemetry(connectionContext, value); // Update the client's user agent headers for existing ThriftClient databricksClient.updateUserAgent(); } diff --git a/src/main/java/com/databricks/jdbc/common/util/UserAgentHelper.java b/src/main/java/com/databricks/jdbc/common/util/UserAgentHelper.java deleted file mode 100644 index 5531d7cb13..0000000000 --- a/src/main/java/com/databricks/jdbc/common/util/UserAgentHelper.java +++ /dev/null @@ -1,66 +0,0 @@ -package com.databricks.jdbc.common.util; - -import static com.databricks.jdbc.common.util.WildcardUtil.isNullOrEmpty; - -import com.databricks.jdbc.api.internal.IDatabricksConnectionContext; -import com.databricks.jdbc.telemetry.TelemetryHelper; -import com.databricks.sdk.core.UserAgent; - -/** Helper class for determining and setting user agent and client app name. */ -public class UserAgentHelper { - private static final String APP_NAME_SYSTEM_PROPERTY = "app.name"; - private static final String VERSION_FILLER = "version"; - - /** - * Determines the application name using a fallback mechanism: 1. useragententry url param 2. - * applicationname url param 3. client info property "applicationname" 4. System property app.name - * - * @param connectionContext The connection context - * @param clientInfoAppName The application name from client info properties, can be null - * @return The determined application name or null if none is found - */ - static String determineApplicationName( - IDatabricksConnectionContext connectionContext, String clientInfoAppName) { - // First check URL params - String appName = connectionContext.getCustomerUserAgent(); - if (!isNullOrEmpty(appName)) { - return appName; - } - - // Then check applicationname URL param - appName = connectionContext.getApplicationName(); - if (!isNullOrEmpty(appName)) { - return appName; - } - - // Then check client info property - if (!isNullOrEmpty(clientInfoAppName)) { - return clientInfoAppName; - } - - // Finally check system property - return System.getProperty(APP_NAME_SYSTEM_PROPERTY); - } - - /** - * Updates both the telemetry client app name and HTTP user agent headers. To be called during - * connection initialization and when app name changes. - * - * @param connectionContext The connection context - * @param clientInfoAppName Optional client info app name, can be null - */ - public static void updateUserAgentAndTelemetry( - IDatabricksConnectionContext connectionContext, String clientInfoAppName) { - String appName = determineApplicationName(connectionContext, clientInfoAppName); - if (!isNullOrEmpty(appName)) { - // Update telemetry - TelemetryHelper.updateClientAppName(appName); - - // Update HTTP user agent - int i = appName.indexOf('/'); - String customerName = (i < 0) ? appName : appName.substring(0, i); - String customerVersion = (i < 0) ? VERSION_FILLER : appName.substring(i + 1); - UserAgent.withOtherInfo(customerName, UserAgent.sanitize(customerVersion)); - } - } -} diff --git a/src/main/java/com/databricks/jdbc/common/util/UserAgentManager.java b/src/main/java/com/databricks/jdbc/common/util/UserAgentManager.java index 433d61c1db..2f6e56c5a9 100644 --- a/src/main/java/com/databricks/jdbc/common/util/UserAgentManager.java +++ b/src/main/java/com/databricks/jdbc/common/util/UserAgentManager.java @@ -1,6 +1,9 @@ package com.databricks.jdbc.common.util; +import static com.databricks.jdbc.common.util.WildcardUtil.isNullOrEmpty; + import com.databricks.jdbc.api.internal.IDatabricksConnectionContext; +import com.databricks.jdbc.telemetry.TelemetryHelper; import com.databricks.sdk.core.UserAgent; public class UserAgentManager { @@ -10,6 +13,8 @@ public class UserAgentManager { private static final String CLIENT_USER_AGENT_PREFIX = "Java"; public static final String USER_AGENT_SEA_CLIENT = "SQLExecHttpClient"; public static final String USER_AGENT_THRIFT_CLIENT = "THttpClient"; + private static final String APP_NAME_SYSTEM_PROPERTY = "app.name"; + private static final String VERSION_FILLER = "version"; /** * Set the user agent for the Databricks JDBC driver. @@ -21,8 +26,61 @@ public static void setUserAgent(IDatabricksConnectionContext connectionContext) UserAgent.withProduct(DEFAULT_USER_AGENT, DriverUtil.getDriverVersion()); UserAgent.withOtherInfo(CLIENT_USER_AGENT_PREFIX, connectionContext.getClientUserAgent()); - // Let UserAgentHelper handle the application name for both telemetry and user agent - UserAgentHelper.updateUserAgentAndTelemetry(connectionContext, null); + // Update application name for both telemetry and user agent + updateUserAgentAndTelemetry(connectionContext, null); + } + + /** + * Determines the application name using a fallback mechanism: 1. useragententry url param 2. + * applicationname url param 3. client info property "applicationname" 4. System property app.name + * + * @param connectionContext The connection context + * @param clientInfoAppName The application name from client info properties, can be null + * @return The determined application name or null if none is found + */ + static String determineApplicationName( + IDatabricksConnectionContext connectionContext, String clientInfoAppName) { + // First check URL params + String appName = connectionContext.getCustomerUserAgent(); + if (!isNullOrEmpty(appName)) { + return appName; + } + + // Then check applicationname URL param + appName = connectionContext.getApplicationName(); + if (!isNullOrEmpty(appName)) { + return appName; + } + + // Then check client info property + if (!isNullOrEmpty(clientInfoAppName)) { + return clientInfoAppName; + } + + // Finally check system property + return System.getProperty(APP_NAME_SYSTEM_PROPERTY); + } + + /** + * Updates both the telemetry client app name and HTTP user agent headers. To be called during + * connection initialization and when app name changes. + * + * @param connectionContext The connection context + * @param clientInfoAppName Optional client info app name, can be null + */ + public static void updateUserAgentAndTelemetry( + IDatabricksConnectionContext connectionContext, String clientInfoAppName) { + String appName = determineApplicationName(connectionContext, clientInfoAppName); + if (!isNullOrEmpty(appName)) { + // Update telemetry + TelemetryHelper.updateClientAppName(appName); + + // Update HTTP user agent + int i = appName.indexOf('/'); + String customerName = (i < 0) ? appName : appName.substring(0, i); + String customerVersion = (i < 0) ? VERSION_FILLER : appName.substring(i + 1); + UserAgent.withOtherInfo(customerName, UserAgent.sanitize(customerVersion)); + } } /** Gets the user agent string for Databricks Driver HTTP Client. */ diff --git a/src/test/java/com/databricks/jdbc/common/util/UserAgentHelperTest.java b/src/test/java/com/databricks/jdbc/common/util/UserAgentHelperTest.java deleted file mode 100644 index 910ef74ca8..0000000000 --- a/src/test/java/com/databricks/jdbc/common/util/UserAgentHelperTest.java +++ /dev/null @@ -1,124 +0,0 @@ -package com.databricks.jdbc.common.util; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.mockito.Mockito.when; - -import com.databricks.jdbc.api.internal.IDatabricksConnectionContext; -import com.databricks.jdbc.telemetry.TelemetryHelper; -import com.databricks.sdk.core.UserAgent; -import java.util.stream.Stream; -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.extension.ExtendWith; -import org.junit.jupiter.params.ParameterizedTest; -import org.junit.jupiter.params.provider.Arguments; -import org.junit.jupiter.params.provider.MethodSource; -import org.mockito.Mock; -import org.mockito.MockedStatic; -import org.mockito.Mockito; -import org.mockito.junit.jupiter.MockitoExtension; - -@ExtendWith(MockitoExtension.class) -public class UserAgentHelperTest { - - @Mock private IDatabricksConnectionContext connectionContext; - - private MockedStatic telemetryHelperMock; - private MockedStatic userAgentMock; - - @BeforeEach - public void setup() { - telemetryHelperMock = Mockito.mockStatic(TelemetryHelper.class); - userAgentMock = Mockito.mockStatic(UserAgent.class); - } - - @AfterEach - public void tearDown() { - telemetryHelperMock.close(); - userAgentMock.close(); - } - - private static Stream provideApplicationNameTestCases() { - return Stream.of( - // Test case 1: UserAgentEntry takes precedence - Arguments.of( - "MyUserAgent", - "AppNameValue", - "ClientInfoApp", - "MyUserAgent", - "When useragententry is set"), - // Test case 2: ApplicationName is used when UserAgentEntry is null - Arguments.of( - null, - "AppNameValue", - "ClientInfoApp", - "AppNameValue", - "When useragententry is not set but applicationname is"), - // Test case 3: ClientInfo is used when both UserAgentEntry and ApplicationName are null - Arguments.of( - null, - null, // applicationName - "ClientInfoApp", - "ClientInfoApp", - "When URL params are not set but client info is provided")); - } - - @ParameterizedTest(name = "{4}") - @MethodSource("provideApplicationNameTestCases") - public void testDetermineApplicationName( - String customerUserAgent, - String applicationName, - String clientInfoApp, - String expectedResult, - String testDescription) { - // Setup only necessary stubs - Mockito.lenient().when(connectionContext.getCustomerUserAgent()).thenReturn(customerUserAgent); - Mockito.lenient().when(connectionContext.getApplicationName()).thenReturn(applicationName); - - // Execute - String result = UserAgentHelper.determineApplicationName(connectionContext, clientInfoApp); - - // Verify - assertEquals(expectedResult, result); - } - - @Test - public void testDetermineApplicationName_WithSystemProperty() { - // When falling back to system property - when(connectionContext.getCustomerUserAgent()).thenReturn(null); - when(connectionContext.getApplicationName()).thenReturn(null); - - System.setProperty("app.name", "SystemPropApp"); - try { - String result = UserAgentHelper.determineApplicationName(connectionContext, null); - assertEquals("SystemPropApp", result); - } finally { - System.clearProperty("app.name"); - } - } - - @Test - public void testUpdateUserAgentAndTelemetry() { - // Test that both telemetry and user agent are updated - when(connectionContext.getCustomerUserAgent()).thenReturn("TestApp"); - when(UserAgent.sanitize("version")).thenReturn("version"); - - UserAgentHelper.updateUserAgentAndTelemetry(connectionContext, null); - - telemetryHelperMock.verify(() -> TelemetryHelper.updateClientAppName("TestApp")); - userAgentMock.verify(() -> UserAgent.withOtherInfo("TestApp", "version")); - } - - @Test - public void testUpdateUserAgentAndTelemetry_WithVersion() { - // Test with app name containing version - when(connectionContext.getCustomerUserAgent()).thenReturn("MyApp/1.2.3"); - when(UserAgent.sanitize("1.2.3")).thenReturn("1.2.3"); - - UserAgentHelper.updateUserAgentAndTelemetry(connectionContext, null); - - telemetryHelperMock.verify(() -> TelemetryHelper.updateClientAppName("MyApp/1.2.3")); - userAgentMock.verify(() -> UserAgent.withOtherInfo("MyApp", "1.2.3")); - } -} diff --git a/src/test/java/com/databricks/jdbc/common/util/UserAgentManagerTest.java b/src/test/java/com/databricks/jdbc/common/util/UserAgentManagerTest.java index 6756db8616..2df34ab0a3 100644 --- a/src/test/java/com/databricks/jdbc/common/util/UserAgentManagerTest.java +++ b/src/test/java/com/databricks/jdbc/common/util/UserAgentManagerTest.java @@ -1,16 +1,131 @@ package com.databricks.jdbc.common.util; import static com.databricks.jdbc.TestConstants.*; +import static com.databricks.jdbc.common.util.UserAgentManager.determineApplicationName; import static com.databricks.jdbc.common.util.UserAgentManager.getUserAgentString; +import static com.databricks.jdbc.common.util.UserAgentManager.updateUserAgentAndTelemetry; import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.Mockito.when; import com.databricks.jdbc.api.impl.DatabricksConnectionContextFactory; import com.databricks.jdbc.api.internal.IDatabricksConnectionContext; import com.databricks.jdbc.exception.DatabricksSQLException; +import com.databricks.jdbc.telemetry.TelemetryHelper; import java.util.Properties; +import java.util.stream.Stream; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; +import org.mockito.Mock; +import org.mockito.MockedStatic; +import org.mockito.Mockito; +import org.mockito.junit.jupiter.MockitoExtension; +@ExtendWith(MockitoExtension.class) public class UserAgentManagerTest { + + @Mock private IDatabricksConnectionContext connectionContext; + + private MockedStatic telemetryHelperMock; + + @BeforeEach + public void setup() { + telemetryHelperMock = Mockito.mockStatic(TelemetryHelper.class); + } + + @AfterEach + public void tearDown() { + telemetryHelperMock.close(); + } + + private static Stream provideApplicationNameTestCases() { + return Stream.of( + // Test case 1: UserAgentEntry takes precedence + Arguments.of( + "MyUserAgent", + "AppNameValue", + "ClientInfoApp", + "MyUserAgent", + "When useragententry is set"), + // Test case 2: ApplicationName is used when UserAgentEntry is null + Arguments.of( + null, + "AppNameValue", + "ClientInfoApp", + "AppNameValue", + "When useragententry is not set but applicationname is"), + // Test case 3: ClientInfo is used when both UserAgentEntry and ApplicationName are null + Arguments.of( + null, + null, // applicationName + "ClientInfoApp", + "ClientInfoApp", + "When URL params are not set but client info is provided")); + } + + @ParameterizedTest + @MethodSource("provideApplicationNameTestCases") + public void testDetermineApplicationName( + String customerUserAgent, + String applicationName, + String clientInfoApp, + String expectedResult) { + // Setup only necessary stubs + Mockito.lenient().when(connectionContext.getCustomerUserAgent()).thenReturn(customerUserAgent); + Mockito.lenient().when(connectionContext.getApplicationName()).thenReturn(applicationName); + + // Execute + String result = determineApplicationName(connectionContext, clientInfoApp); + + // Verify + assertEquals(expectedResult, result); + } + + @Test + public void testDetermineApplicationName_WithSystemProperty() { + // When falling back to system property + when(connectionContext.getCustomerUserAgent()).thenReturn(null); + when(connectionContext.getApplicationName()).thenReturn(null); + + System.setProperty("app.name", "SystemPropApp"); + try { + String result = determineApplicationName(connectionContext, null); + assertEquals("SystemPropApp", result); + } finally { + System.clearProperty("app.name"); + } + } + + @Test + public void testUpdateUserAgentAndTelemetry() { + // Test that both telemetry and user agent are updated + when(connectionContext.getCustomerUserAgent()).thenReturn("TestApp"); + // when(UserAgent.sanitize("version")).thenReturn("version"); + + updateUserAgentAndTelemetry(connectionContext, null); + + telemetryHelperMock.verify(() -> TelemetryHelper.updateClientAppName("TestApp")); + String userAgent = getUserAgentString(); + assertTrue(userAgent.contains("TestApp/version")); + } + + @Test + public void testUpdateUserAgentAndTelemetry_WithVersion() { + // Test with app name containing version + when(connectionContext.getCustomerUserAgent()).thenReturn("MyApp/1.2.3"); + // when(UserAgent.sanitize("1.2.3")).thenReturn("1.2.3"); + + updateUserAgentAndTelemetry(connectionContext, null); + + telemetryHelperMock.verify(() -> TelemetryHelper.updateClientAppName("MyApp/1.2.3")); + String userAgent = getUserAgentString(); + assertTrue(userAgent.contains("MyApp/1.2.3")); + } + @Test void testUserAgentSetsClientCorrectly() throws DatabricksSQLException { // Thrift with all-purpose cluster From 0e1b2f3e494709be89e2978d81b291a362e8df5d Mon Sep 17 00:00:00 2001 From: Shivam Raj Date: Mon, 2 Jun 2025 16:23:23 +0530 Subject: [PATCH 10/11] addressed reviews --- .../dbclient/impl/thrift/DatabricksHttpTTransport.java | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksHttpTTransport.java b/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksHttpTTransport.java index b6ecb758ee..2298d5cb9a 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksHttpTTransport.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksHttpTTransport.java @@ -147,11 +147,10 @@ public void checkReadBytesAvailable(long numBytes) throws TTransportException {} private void refreshHeadersIfRequired() { Map refreshedHeaders = databricksConfig.authenticate(); // Preserve existing custom headers (like User-Agent) and add/update auth headers - Map newHeaders = new HashMap<>(customHeaders); if (refreshedHeaders != null) { - newHeaders.putAll(refreshedHeaders); + customHeaders = new HashMap<>(customHeaders); + customHeaders.putAll(refreshedHeaders); } - customHeaders = newHeaders; } void resetAccessToken(String newAccessToken) { @@ -163,9 +162,8 @@ void resetAccessToken(String newAccessToken) { void updateUserAgent() { // Add user agent as a header since the HTTP client's user agent is set during construction // and cannot be changed. This ensures the updated user agent is sent with thrift requests. - Map headers = new HashMap<>(customHeaders); - headers.put("User-Agent", UserAgentManager.getUserAgentString()); - customHeaders = headers; + customHeaders = new HashMap<>(customHeaders); + customHeaders.put("User-Agent", UserAgentManager.getUserAgentString()); } @VisibleForTesting From 83fd62c3c7cb58a05a1a443ba67e2863685aaf50 Mon Sep 17 00:00:00 2001 From: Shivam Raj Date: Mon, 2 Jun 2025 19:40:11 +0530 Subject: [PATCH 11/11] update DatabricksHttpClient --- .../jdbc/api/impl/DatabricksSession.java | 2 -- .../jdbc/dbclient/IDatabricksClient.java | 3 --- .../impl/http/DatabricksHttpClient.java | 5 +++++ .../impl/sqlexec/DatabricksSdkClient.java | 6 ----- .../impl/thrift/DatabricksHttpTTransport.java | 16 ++------------ .../thrift/DatabricksThriftServiceClient.java | 6 ----- .../jdbc/api/impl/DatabricksSessionTest.java | 22 ------------------- 7 files changed, 7 insertions(+), 53 deletions(-) diff --git a/src/main/java/com/databricks/jdbc/api/impl/DatabricksSession.java b/src/main/java/com/databricks/jdbc/api/impl/DatabricksSession.java index 6a7f97f978..508f2def79 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksSession.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksSession.java @@ -252,8 +252,6 @@ public void setClientInfoProperty(String name, String value) { // If application name is being set, update both telemetry and user agent if (name.equalsIgnoreCase(DatabricksJdbcUrlParams.APPLICATION_NAME.getParamName())) { UserAgentManager.updateUserAgentAndTelemetry(connectionContext, value); - // Update the client's user agent headers for existing ThriftClient - databricksClient.updateUserAgent(); } clientInfoProperties.put(name, value); diff --git a/src/main/java/com/databricks/jdbc/dbclient/IDatabricksClient.java b/src/main/java/com/databricks/jdbc/dbclient/IDatabricksClient.java index c7965e5ad7..14c1890bf9 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/IDatabricksClient.java +++ b/src/main/java/com/databricks/jdbc/dbclient/IDatabricksClient.java @@ -126,9 +126,6 @@ Collection getResultChunks(StatementId statementId, long chunkInde */ void resetAccessToken(String newAccessToken); - /** Update the user agent headers when application name changes */ - void updateUserAgent(); - TFetchResultsResp getMoreResults(IDatabricksStatementInternal parentStatement) throws DatabricksSQLException; diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/http/DatabricksHttpClient.java b/src/main/java/com/databricks/jdbc/dbclient/impl/http/DatabricksHttpClient.java index ffd599e373..da854eb1ad 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/http/DatabricksHttpClient.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/http/DatabricksHttpClient.java @@ -1,6 +1,7 @@ package com.databricks.jdbc.dbclient.impl.http; import static com.databricks.jdbc.common.DatabricksJdbcConstants.*; +import static com.databricks.jdbc.common.util.WildcardUtil.isNullOrEmpty; import static com.databricks.jdbc.dbclient.impl.common.ClientConfigurator.convertNonProxyHostConfigToBeSystemPropertyCompliant; import static io.netty.util.NetUtil.LOCALHOST; @@ -82,6 +83,10 @@ public CloseableHttpResponse execute(HttpUriRequest request, boolean supportGzip request.setHeader("Content-Encoding", "gzip"); } try { + String userAgentString = UserAgentManager.getUserAgentString(); + if (!isNullOrEmpty(userAgentString) && !request.containsHeader("User-Agent")) { + request.setHeader("User-Agent", userAgentString); + } return httpClient.execute(request); } catch (IOException e) { throwHttpException(e, request); diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClient.java b/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClient.java index 2d234cf0be..210d19861e 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClient.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/sqlexec/DatabricksSdkClient.java @@ -415,12 +415,6 @@ public synchronized void resetAccessToken(String newAccessToken) { this.apiClient = workspaceClient.apiClient(); } - @Override - public void updateUserAgent() { - // For SDK client, the user agent is handled by the global UserAgent state - // which is automatically picked up by the SDK apiClient for new requests - } - @Override public TFetchResultsResp getMoreResults(IDatabricksStatementInternal parentStatement) throws DatabricksSQLException { diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksHttpTTransport.java b/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksHttpTTransport.java index 2298d5cb9a..e4c239bba3 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksHttpTTransport.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksHttpTTransport.java @@ -3,7 +3,6 @@ import static com.databricks.jdbc.common.util.DatabricksAuthUtil.initializeConfigWithToken; import com.databricks.jdbc.api.internal.IDatabricksConnectionContext; -import com.databricks.jdbc.common.util.UserAgentManager; import com.databricks.jdbc.common.util.ValidationUtil; import com.databricks.jdbc.dbclient.IDatabricksHttpClient; import com.databricks.jdbc.dbclient.impl.common.TracingUtil; @@ -146,11 +145,8 @@ public void checkReadBytesAvailable(long numBytes) throws TTransportException {} /** Refreshes the custom headers by re-authenticating if necessary. */ private void refreshHeadersIfRequired() { Map refreshedHeaders = databricksConfig.authenticate(); - // Preserve existing custom headers (like User-Agent) and add/update auth headers - if (refreshedHeaders != null) { - customHeaders = new HashMap<>(customHeaders); - customHeaders.putAll(refreshedHeaders); - } + customHeaders = + refreshedHeaders != null ? new HashMap<>(refreshedHeaders) : Collections.emptyMap(); } void resetAccessToken(String newAccessToken) { @@ -158,14 +154,6 @@ void resetAccessToken(String newAccessToken) { this.databricksConfig.resolve(); } - /** Updates the user agent for thrift requests by adding it as a header */ - void updateUserAgent() { - // Add user agent as a header since the HTTP client's user agent is set during construction - // and cannot be changed. This ensures the updated user agent is sent with thrift requests. - customHeaders = new HashMap<>(customHeaders); - customHeaders.put("User-Agent", UserAgentManager.getUserAgentString()); - } - @VisibleForTesting void setResponseBuffer(ByteArrayInputStream responseBuffer) { this.responseBuffer = responseBuffer; diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftServiceClient.java b/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftServiceClient.java index ca912b0efe..76ce3cd3bd 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftServiceClient.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftServiceClient.java @@ -76,12 +76,6 @@ public void resetAccessToken(String newAccessToken) { .resetAccessToken(newAccessToken); } - @Override - public void updateUserAgent() { - ((DatabricksHttpTTransport) thriftAccessor.getThriftClient().getInputProtocol().getTransport()) - .updateUserAgent(); - } - @Override public ImmutableSessionInfo createSession( IDatabricksComputeResource cluster, diff --git a/src/test/java/com/databricks/jdbc/api/impl/DatabricksSessionTest.java b/src/test/java/com/databricks/jdbc/api/impl/DatabricksSessionTest.java index 241e6ec0e7..c4933208fd 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/DatabricksSessionTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/DatabricksSessionTest.java @@ -231,26 +231,4 @@ public void testSetClientInfoProperty_AuthAccessTokenThrift() throws DatabricksS DatabricksJdbcUrlParams.AUTH_ACCESS_TOKEN.getParamName(), "token"); verify(thriftClient).resetAccessToken("token"); } - - @Test - public void testSetClientInfoProperty_ApplicationNameUpdatesUserAgent() - throws DatabricksSQLException { - DatabricksSession session = - new DatabricksSession( - DatabricksConnectionContext.parse(VALID_CLUSTER_URL, new Properties()), sdkClient); - session.setClientInfoProperty( - DatabricksJdbcUrlParams.APPLICATION_NAME.getParamName(), "testApp"); - verify(sdkClient).updateUserAgent(); - } - - @Test - public void testSetClientInfoProperty_ApplicationNameUpdatesUserAgentThrift() - throws DatabricksSQLException { - DatabricksSession session = - new DatabricksSession( - DatabricksConnectionContext.parse(VALID_CLUSTER_URL, new Properties()), thriftClient); - session.setClientInfoProperty( - DatabricksJdbcUrlParams.APPLICATION_NAME.getParamName(), "testApp"); - verify(thriftClient).updateUserAgent(); - } }