[PECOBLR-435] added application name to track BiTool

src/main/java/com/databricks/client/jdbc/Driver.java

@@ -57,6 +57,7 @@ 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);

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<String, String> buildPropertiesMap(
       String connectionParamString, Properties properties) {
     ImmutableMap.Builder<String, String> 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();
   }

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.UserAgentHelper;
 import com.databricks.jdbc.dbclient.IDatabricksClient;
 import com.databricks.jdbc.dbclient.IDatabricksMetadataClient;
 import com.databricks.jdbc.dbclient.impl.sqlexec.DatabricksEmptyMetadataClient;
@@ -247,6 +248,14 @@ 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())) {
+      UserAgentHelper.updateUserAgentAndTelemetry(connectionContext, value);
+      // Update the client's user agent headers for existing ThriftClient
+      databricksClient.updateUserAgent();
+    }
+
     clientInfoProperties.put(name, value);
   }
 

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();
 }

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<String, String> JSON_HTTP_HEADERS =
       Map.of(
           "Accept", "application/json",

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;

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
+   */
+  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));
+    }
+  }
+}

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. */

src/main/java/com/databricks/jdbc/dbclient/IDatabricksClient.java

@@ -126,6 +126,9 @@ Collection<ExternalLink> 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;
 

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 {

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,15 +146,28 @@ public void checkReadBytesAvailable(long numBytes) throws TTransportException {}
   /** Refreshes the custom headers by re-authenticating if necessary. */
   private void refreshHeadersIfRequired() {
     Map<String, String> refreshedHeaders = databricksConfig.authenticate();
-    customHeaders =
-        refreshedHeaders != null ? new HashMap<>(refreshedHeaders) : Collections.emptyMap();
+    // Preserve existing custom headers (like User-Agent) and add/update auth headers
+    Map<String, String> newHeaders = new HashMap<>(customHeaders);
+    if (refreshedHeaders != null) {
+      newHeaders.putAll(refreshedHeaders);
+    }
+    customHeaders = newHeaders;
   }
 
   void resetAccessToken(String newAccessToken) {
     this.databricksConfig = initializeConfigWithToken(newAccessToken, databricksConfig);
     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<String, String> headers = new HashMap<>(customHeaders);
+    headers.put("User-Agent", UserAgentManager.getUserAgentString());
+    customHeaders = headers;
+  }
+
   @VisibleForTesting
   void setResponseBuffer(ByteArrayInputStream responseBuffer) {
     this.responseBuffer = responseBuffer;

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,

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;

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) {

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 =
@@ -220,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();
+  }
 }