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..4e24a371c4 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); @@ -379,6 +379,7 @@ public String getClientUserAgent() { : USER_AGENT_THRIFT_CLIENT; } + @Override public String getCustomerUserAgent() { return getParameter(DatabricksJdbcUrlParams.USER_AGENT_ENTRY); } @@ -887,6 +888,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..508f2def79 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksSession.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksSession.java @@ -8,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.UserAgentManager; import com.databricks.jdbc.dbclient.IDatabricksClient; import com.databricks.jdbc.dbclient.IDatabricksMetadataClient; import com.databricks.jdbc.dbclient.impl.sqlexec.DatabricksEmptyMetadataClient; @@ -247,6 +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())) { + UserAgentManager.updateUserAgentAndTelemetry(connectionContext, 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/common/util/UserAgentManager.java b/src/main/java/com/databricks/jdbc/common/util/UserAgentManager.java index 17ee19088d..2f6e56c5a9 100644 --- a/src/main/java/com/databricks/jdbc/common/util/UserAgentManager.java +++ b/src/main/java/com/databricks/jdbc/common/util/UserAgentManager.java @@ -3,16 +3,18 @@ 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 { 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"; 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. @@ -20,14 +22,63 @@ 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); + + // 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)); } } 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/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/main/java/com/databricks/jdbc/telemetry/TelemetryHelper.java b/src/main/java/com/databricks/jdbc/telemetry/TelemetryHelper.java index c0870b8d6a..67800a6853 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.DatabricksClientConfiguratorManager; import com.databricks.jdbc.common.safe.DatabricksDriverFeatureFlagsContextFactory; @@ -32,7 +34,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()) @@ -50,6 +51,12 @@ public static DriverSystemConfiguration getDriverSystemConfiguration() { return DRIVER_SYSTEM_CONFIGURATION; } + public static void updateClientAppName(String clientAppName) { + if (!isNullOrEmpty(clientAppName)) { + DRIVER_SYSTEM_CONFIGURATION.setClientAppName(clientAppName); + } + } + public static boolean isTelemetryAllowedForConnection(IDatabricksConnectionContext context) { return context != null && context.isTelemetryEnabled() 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/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 diff --git a/src/test/java/com/databricks/jdbc/telemetry/TelemetryHelperTest.java b/src/test/java/com/databricks/jdbc/telemetry/TelemetryHelperTest.java index 6c694e9755..a58a17d60e 100644 --- a/src/test/java/com/databricks/jdbc/telemetry/TelemetryHelperTest.java +++ b/src/test/java/com/databricks/jdbc/telemetry/TelemetryHelperTest.java @@ -5,6 +5,8 @@ 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; @@ -82,6 +84,27 @@ 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()); + } + @Test public void testGetDatabricksConfigSafely_ReturnsNullOnError() { try (MockedStatic mockedFactory =