Skip to content

Commit 3957211

Browse files
shivam2680shivam2680
authored
[PECOBLR-435] added application name to track BiTool (#837)
* added application name to track BiTool * add tests * some changes * refactored code; update user agent and telemetry client app * update user agent thrift mode as well * changes fn to package-private. * make parameterized test * spotless * merge UserAgentHelper and UserAgentManager * addressed reviews * update DatabricksHttpClient
1 parent cd24356 commit 3957211

12 files changed

Lines changed: 249 additions & 12 deletions

File tree

‎src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,11 +5,11 @@
55
import static com.databricks.jdbc.common.EnvironmentVariables.DEFAULT_ROW_LIMIT_PER_BLOCK;
66
import static com.databricks.jdbc.common.util.UserAgentManager.USER_AGENT_SEA_CLIENT;
77
import static com.databricks.jdbc.common.util.UserAgentManager.USER_AGENT_THRIFT_CLIENT;
8+
import static com.databricks.jdbc.common.util.WildcardUtil.isNullOrEmpty;
89

910
import com.databricks.jdbc.api.internal.IDatabricksConnectionContext;
1011
import com.databricks.jdbc.common.*;
1112
import com.databricks.jdbc.common.util.ValidationUtil;
12-
import com.databricks.jdbc.common.util.WildcardUtil;
1313
import com.databricks.jdbc.exception.DatabricksDriverException;
1414
import com.databricks.jdbc.exception.DatabricksParsingException;
1515
import com.databricks.jdbc.exception.DatabricksSQLException;
@@ -83,7 +83,7 @@ public static ImmutableMap<String, String> buildPropertiesMap(
8383
String connectionParamString, Properties properties) {
8484
ImmutableMap.Builder<String, String> parametersBuilder = ImmutableMap.builder();
8585
// check if connectionParamString is empty or null
86-
if (!WildcardUtil.isNullOrEmpty(connectionParamString)) {
86+
if (!isNullOrEmpty(connectionParamString)) {
8787
String[] urlParts = connectionParamString.split(DatabricksJdbcConstants.URL_DELIMITER);
8888
for (String urlPart : urlParts) {
8989
String[] pair = urlPart.split(DatabricksJdbcConstants.PAIR_DELIMITER);
@@ -379,6 +379,7 @@ public String getClientUserAgent() {
379379
: USER_AGENT_THRIFT_CLIENT;
380380
}
381381

382+
@Override
382383
public String getCustomerUserAgent() {
383384
return getParameter(DatabricksJdbcUrlParams.USER_AGENT_ENTRY);
384385
}
@@ -887,6 +888,11 @@ public boolean isTokenCacheEnabled() {
887888
return getParameter(DatabricksJdbcUrlParams.ENABLE_TOKEN_CACHE).equals("1");
888889
}
889890

891+
@Override
892+
public String getApplicationName() {
893+
return getParameter(DatabricksJdbcUrlParams.APPLICATION_NAME);
894+
}
895+
890896
private static boolean nullOrEmptyString(String s) {
891897
return s == null || s.isEmpty();
892898
}

‎src/main/java/com/databricks/jdbc/api/impl/DatabricksSession.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import com.databricks.jdbc.common.DatabricksClientType;
99
import com.databricks.jdbc.common.DatabricksJdbcUrlParams;
1010
import com.databricks.jdbc.common.IDatabricksComputeResource;
11+
import com.databricks.jdbc.common.util.UserAgentManager;
1112
import com.databricks.jdbc.dbclient.IDatabricksClient;
1213
import com.databricks.jdbc.dbclient.IDatabricksMetadataClient;
1314
import com.databricks.jdbc.dbclient.impl.sqlexec.DatabricksEmptyMetadataClient;
@@ -247,6 +248,12 @@ public void setClientInfoProperty(String name, String value) {
247248
this.databricksClient.resetAccessToken(value);
248249
value = REDACTED_TOKEN; // mask access token
249250
}
251+
252+
// If application name is being set, update both telemetry and user agent
253+
if (name.equalsIgnoreCase(DatabricksJdbcUrlParams.APPLICATION_NAME.getParamName())) {
254+
UserAgentManager.updateUserAgentAndTelemetry(connectionContext, value);
255+
}
256+
250257
clientInfoProperties.put(name, value);
251258
}
252259

‎src/main/java/com/databricks/jdbc/api/internal/IDatabricksConnectionContext.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -317,4 +317,7 @@ public interface IDatabricksConnectionContext {
317317

318318
/** Returns whether token caching is enabled for OAuth authentication */
319319
boolean isTokenCacheEnabled();
320+
321+
/** Returns the application name using JDBC Connection */
322+
String getApplicationName();
320323
}

‎src/main/java/com/databricks/jdbc/common/DatabricksJdbcConstants.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,8 @@ public final class DatabricksJdbcConstants {
8181
Set.of(
8282
ALLOWED_VOLUME_INGESTION_PATHS,
8383
ALLOWED_STAGING_INGESTION_PATHS,
84-
DatabricksJdbcUrlParams.AUTH_ACCESS_TOKEN.getParamName());
84+
DatabricksJdbcUrlParams.AUTH_ACCESS_TOKEN.getParamName(),
85+
DatabricksJdbcUrlParams.APPLICATION_NAME.getParamName());
8586
public static final Map<String, String> JSON_HTTP_HEADERS =
8687
Map.of(
8788
"Accept", "application/json",

‎src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -127,7 +127,9 @@ public enum DatabricksJdbcUrlParams {
127127
"255"),
128128
SOCKET_TIMEOUT("socketTimeout", "Socket timeout in seconds", "900"),
129129
TOKEN_CACHE_PASS_PHRASE("TokenCachePassPhrase", "Pass phrase to use for OAuth U2M Token Cache"),
130-
ENABLE_TOKEN_CACHE("EnableTokenCache", "Enable caching OAuth tokens", "1");
130+
ENABLE_TOKEN_CACHE("EnableTokenCache", "Enable caching OAuth tokens", "1"),
131+
APPLICATION_NAME("ApplicationName", "Name of application using the driver", ""),
132+
;
131133

132134
private final String paramName;
133135
private final String defaultValue;

‎src/main/java/com/databricks/jdbc/common/util/UserAgentManager.java‎

Lines changed: 58 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -3,31 +3,82 @@
33
import static com.databricks.jdbc.common.util.WildcardUtil.isNullOrEmpty;
44

55
import com.databricks.jdbc.api.internal.IDatabricksConnectionContext;
6+
import com.databricks.jdbc.telemetry.TelemetryHelper;
67
import com.databricks.sdk.core.UserAgent;
78

89
public class UserAgentManager {
910
private static final String SDK_USER_AGENT = "databricks-sdk-java";
1011
private static final String JDBC_HTTP_USER_AGENT = "databricks-jdbc-http";
11-
private static final String VERSION_FILLER_FOR_CUSTOMER_USER_AGENT = "version";
1212
private static final String DEFAULT_USER_AGENT = "DatabricksJDBCDriverOSS";
1313
private static final String CLIENT_USER_AGENT_PREFIX = "Java";
1414
public static final String USER_AGENT_SEA_CLIENT = "SQLExecHttpClient";
1515
public static final String USER_AGENT_THRIFT_CLIENT = "THttpClient";
16+
private static final String APP_NAME_SYSTEM_PROPERTY = "app.name";
17+
private static final String VERSION_FILLER = "version";
1618

1719
/**
1820
* Set the user agent for the Databricks JDBC driver.
1921
*
2022
* @param connectionContext The connection context.
2123
*/
2224
public static void setUserAgent(IDatabricksConnectionContext connectionContext) {
25+
// Set the base product and client info
2326
UserAgent.withProduct(DEFAULT_USER_AGENT, DriverUtil.getDriverVersion());
2427
UserAgent.withOtherInfo(CLIENT_USER_AGENT_PREFIX, connectionContext.getClientUserAgent());
25-
String customerUA = connectionContext.getCustomerUserAgent();
26-
if (!isNullOrEmpty(customerUA)) {
27-
int i = customerUA.indexOf('/');
28-
String customerName = (i < 0) ? customerUA : customerUA.substring(0, i);
29-
String customerVersion =
30-
(i < 0) ? VERSION_FILLER_FOR_CUSTOMER_USER_AGENT : customerUA.substring(i + 1);
28+
29+
// Update application name for both telemetry and user agent
30+
updateUserAgentAndTelemetry(connectionContext, null);
31+
}
32+
33+
/**
34+
* Determines the application name using a fallback mechanism: 1. useragententry url param 2.
35+
* applicationname url param 3. client info property "applicationname" 4. System property app.name
36+
*
37+
* @param connectionContext The connection context
38+
* @param clientInfoAppName The application name from client info properties, can be null
39+
* @return The determined application name or null if none is found
40+
*/
41+
static String determineApplicationName(
42+
IDatabricksConnectionContext connectionContext, String clientInfoAppName) {
43+
// First check URL params
44+
String appName = connectionContext.getCustomerUserAgent();
45+
if (!isNullOrEmpty(appName)) {
46+
return appName;
47+
}
48+
49+
// Then check applicationname URL param
50+
appName = connectionContext.getApplicationName();
51+
if (!isNullOrEmpty(appName)) {
52+
return appName;
53+
}
54+
55+
// Then check client info property
56+
if (!isNullOrEmpty(clientInfoAppName)) {
57+
return clientInfoAppName;
58+
}
59+
60+
// Finally check system property
61+
return System.getProperty(APP_NAME_SYSTEM_PROPERTY);
62+
}
63+
64+
/**
65+
* Updates both the telemetry client app name and HTTP user agent headers. To be called during
66+
* connection initialization and when app name changes.
67+
*
68+
* @param connectionContext The connection context
69+
* @param clientInfoAppName Optional client info app name, can be null
70+
*/
71+
public static void updateUserAgentAndTelemetry(
72+
IDatabricksConnectionContext connectionContext, String clientInfoAppName) {
73+
String appName = determineApplicationName(connectionContext, clientInfoAppName);
74+
if (!isNullOrEmpty(appName)) {
75+
// Update telemetry
76+
TelemetryHelper.updateClientAppName(appName);
77+
78+
// Update HTTP user agent
79+
int i = appName.indexOf('/');
80+
String customerName = (i < 0) ? appName : appName.substring(0, i);
81+
String customerVersion = (i < 0) ? VERSION_FILLER : appName.substring(i + 1);
3182
UserAgent.withOtherInfo(customerName, UserAgent.sanitize(customerVersion));
3283
}
3384
}

‎src/main/java/com/databricks/jdbc/dbclient/impl/http/DatabricksHttpClient.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package com.databricks.jdbc.dbclient.impl.http;
22

33
import static com.databricks.jdbc.common.DatabricksJdbcConstants.*;
4+
import static com.databricks.jdbc.common.util.WildcardUtil.isNullOrEmpty;
45
import static com.databricks.jdbc.dbclient.impl.common.ClientConfigurator.convertNonProxyHostConfigToBeSystemPropertyCompliant;
56
import static io.netty.util.NetUtil.LOCALHOST;
67

@@ -82,6 +83,10 @@ public CloseableHttpResponse execute(HttpUriRequest request, boolean supportGzip
8283
request.setHeader("Content-Encoding", "gzip");
8384
}
8485
try {
86+
String userAgentString = UserAgentManager.getUserAgentString();
87+
if (!isNullOrEmpty(userAgentString) && !request.containsHeader("User-Agent")) {
88+
request.setHeader("User-Agent", userAgentString);
89+
}
8590
return httpClient.execute(request);
8691
} catch (IOException e) {
8792
throwHttpException(e, request);

‎src/main/java/com/databricks/jdbc/model/telemetry/DriverSystemConfiguration.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import com.databricks.sdk.support.ToStringer;
44
import com.fasterxml.jackson.annotation.JsonProperty;
5+
import com.google.common.annotations.VisibleForTesting;
56

67
public class DriverSystemConfiguration {
78
// TODO : add json properties when proto is implemented completely
@@ -38,6 +39,11 @@ public class DriverSystemConfiguration {
3839
@JsonProperty("char_set_encoding")
3940
private String charSetEncoding;
4041

42+
@VisibleForTesting
43+
public String getClientAppName() {
44+
return clientAppName;
45+
}
46+
4147
public DriverSystemConfiguration setDriverName(String driverName) {
4248
this.driverName = driverName;
4349
return this;

‎src/main/java/com/databricks/jdbc/telemetry/TelemetryHelper.java‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
package com.databricks.jdbc.telemetry;
22

3+
import static com.databricks.jdbc.common.util.WildcardUtil.isNullOrEmpty;
4+
35
import com.databricks.jdbc.api.internal.IDatabricksConnectionContext;
46
import com.databricks.jdbc.common.DatabricksClientConfiguratorManager;
57
import com.databricks.jdbc.common.safe.DatabricksDriverFeatureFlagsContextFactory;
@@ -32,7 +34,6 @@ public class TelemetryHelper {
3234

3335
private static final DriverSystemConfiguration DRIVER_SYSTEM_CONFIGURATION =
3436
new DriverSystemConfiguration()
35-
.setClientAppName(null)
3637
.setCharSetEncoding(Charset.defaultCharset().displayName())
3738
.setDriverName(DriverUtil.getDriverName())
3839
.setDriverVersion(DriverUtil.getDriverVersion())
@@ -50,6 +51,12 @@ public static DriverSystemConfiguration getDriverSystemConfiguration() {
5051
return DRIVER_SYSTEM_CONFIGURATION;
5152
}
5253

54+
public static void updateClientAppName(String clientAppName) {
55+
if (!isNullOrEmpty(clientAppName)) {
56+
DRIVER_SYSTEM_CONFIGURATION.setClientAppName(clientAppName);
57+
}
58+
}
59+
5360
public static boolean isTelemetryAllowedForConnection(IDatabricksConnectionContext context) {
5461
return context != null
5562
&& context.isTelemetryEnabled()

‎src/test/java/com/databricks/jdbc/api/impl/DatabricksSessionTest.java‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import static com.databricks.jdbc.TestConstants.*;
44
import static com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode.TEMPORARY_REDIRECT_EXCEPTION;
5+
import static com.databricks.jdbc.telemetry.TelemetryHelper.getDriverSystemConfiguration;
56
import static org.junit.jupiter.api.Assertions.*;
67
import static org.mockito.ArgumentMatchers.any;
78
import static org.mockito.ArgumentMatchers.eq;
@@ -201,6 +202,16 @@ public void testSetClientInfoProperty() throws DatabricksSQLException {
201202
assertEquals("value", session.getClientInfoProperties().get("key"));
202203
}
203204

205+
@Test
206+
public void testGetClientInfoProperty_ApplicationName() throws DatabricksSQLException {
207+
DatabricksSession session =
208+
new DatabricksSession(
209+
DatabricksConnectionContext.parse(VALID_CLUSTER_URL, new Properties()), sdkClient);
210+
session.setClientInfoProperty("applicationName", "testApp");
211+
assertEquals("testApp", session.getClientInfoProperties().get("applicationName"));
212+
assertEquals("testApp", getDriverSystemConfiguration().getClientAppName());
213+
}
214+
204215
@Test
205216
public void testSetClientInfoProperty_AuthAccessToken() throws DatabricksSQLException {
206217
DatabricksSession session =

0 commit comments

Comments
 (0)