Skip to content

Commit 21e07bd

Browse files
gopalldbclaude
gopalldb
and
claude
authored
Add UserAgent for connector service calls (#1175)
## Problem Custom user agent from useragententry parameter wasn't included in connector service HTTP requests for feature flag retrieval. ## Root Cause Method execution order issue in UserAgentManager.setUserAgent() - custom user agent was set AFTER the connector service request was made. ## Solution Reordered the method to set custom user agent before calling getClientUserAgent() (which triggers feature flag fetch). ## Testing - Added testCustomUserAgentIncludedBeforeClientTypeEvaluation() test NO_CHANGELOG=true --------- Co-authored-by: Claude Sonnet 4.5 <noreply@anthropic.com>
1 parent ec8a6e8 commit 21e07bd

4 files changed

Lines changed: 138 additions & 20 deletions

File tree

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -457,7 +457,7 @@ public DatabricksClientType getClientTypeFromContext() {
457457
// Check if circuit breaker is open due to recent 429 rate limit failures
458458
if (SeaCircuitBreakerManager.isCircuitOpen()) {
459459
long remainingMs = SeaCircuitBreakerManager.getTimeRemainingMs();
460-
LOGGER.info(
460+
LOGGER.debug(
461461
"SEA circuit breaker is OPEN due to recent 429 rate limit failures. "
462462
+ "Using THRIFT client. Circuit will close in {} ({}ms)",
463463
SeaCircuitBreakerManager.getTimeRemainingFormatted(),

‎src/main/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContext.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import com.databricks.jdbc.common.DatabricksClientConfiguratorManager;
55
import com.databricks.jdbc.common.util.DriverUtil;
66
import com.databricks.jdbc.common.util.JsonUtil;
7+
import com.databricks.jdbc.common.util.UserAgentManager;
78
import com.databricks.jdbc.dbclient.IDatabricksHttpClient;
89
import com.databricks.jdbc.dbclient.impl.http.DatabricksHttpClientFactory;
910
import com.databricks.jdbc.exception.DatabricksHttpException;
@@ -90,6 +91,12 @@ private void refreshAllFeatureFlags() {
9091
IDatabricksHttpClient httpClient =
9192
DatabricksHttpClientFactory.getInstance().getClient(connectionContext);
9293
HttpGet request = new HttpGet(featureFlagEndpoint);
94+
95+
// Set custom User-Agent for connector service (includes custom user agent without client
96+
// type)
97+
String userAgent = UserAgentManager.buildUserAgentForConnectorService(connectionContext);
98+
request.setHeader("User-Agent", userAgent);
99+
93100
DatabricksClientConfiguratorManager.getInstance()
94101
.getConfigurator(connectionContext)
95102
.getDatabricksConfig()

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

Lines changed: 79 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -17,33 +17,96 @@ public class UserAgentManager {
1717
public static final String USER_AGENT_THRIFT_CLIENT = "THttpClient";
1818
private static final String VERSION_FILLER = "version";
1919

20+
/**
21+
* Parse custom user agent string into name and version components.
22+
*
23+
* @param customerUserAgent The custom user agent string (may be URL encoded)
24+
* @return String array [name, version] or null if parsing fails
25+
*/
26+
private static String[] parseCustomerUserAgent(String customerUserAgent) {
27+
try {
28+
String decodedUA = URLDecoder.decode(customerUserAgent, StandardCharsets.UTF_8);
29+
int i = decodedUA.indexOf('/');
30+
String customerName = (i < 0) ? decodedUA : decodedUA.substring(0, i);
31+
String customerVersion = (i < 0) ? VERSION_FILLER : decodedUA.substring(i + 1);
32+
return new String[] {customerName, customerVersion};
33+
} catch (Exception e) {
34+
LOGGER.debug("Failed to parse customer userAgent entry {}, Error {}", customerUserAgent, e);
35+
return null;
36+
}
37+
}
38+
2039
/**
2140
* Set the user agent for the Databricks JDBC driver.
2241
*
2342
* @param connectionContext The connection context.
2443
*/
2544
public static void setUserAgent(IDatabricksConnectionContext connectionContext) {
26-
// Set the base product and client info
45+
// Set the base product
2746
UserAgent.withProduct(DEFAULT_USER_AGENT, DriverUtil.getDriverVersion());
47+
48+
// Set client info (this may trigger getClientType which fetches feature flags)
2849
UserAgent.withOtherInfo(CLIENT_USER_AGENT_PREFIX, connectionContext.getClientUserAgent());
29-
if (connectionContext.getCustomerUserAgent() == null) {
30-
return;
50+
51+
// Set custom user agent (maintains proper order: base -> client type -> custom)
52+
if (connectionContext.getCustomerUserAgent() != null) {
53+
String[] parsed = parseCustomerUserAgent(connectionContext.getCustomerUserAgent());
54+
if (parsed != null) {
55+
try {
56+
UserAgent.withOtherInfo(parsed[0], UserAgent.sanitize(parsed[1]));
57+
} catch (IllegalArgumentException e) {
58+
LOGGER.debug(
59+
"Failed to set user agent for customer userAgent entry {}, Error {}",
60+
connectionContext.getCustomerUserAgent(),
61+
e);
62+
}
63+
}
3164
}
32-
try {
33-
String decodedUA =
34-
URLDecoder.decode(
65+
}
66+
67+
/**
68+
* Build user agent string for connector service requests (without client type to avoid circular
69+
* dependency). This is used specifically for feature flags requests since client type is not yet
70+
* determined.
71+
*
72+
* @param connectionContext The connection context.
73+
* @return User agent string with format: "DatabricksJDBCDriverOSS/version databricks-jdbc-http
74+
* jvm/version os/name [CustomApp/version]"
75+
*/
76+
public static String buildUserAgentForConnectorService(
77+
IDatabricksConnectionContext connectionContext) {
78+
StringBuilder userAgent = new StringBuilder();
79+
80+
// Base product: DatabricksJDBCDriverOSS/version
81+
userAgent.append(DEFAULT_USER_AGENT).append("/").append(DriverUtil.getDriverVersion());
82+
83+
// JDBC HTTP identifier
84+
userAgent.append(" ").append(JDBC_HTTP_USER_AGENT);
85+
86+
// JVM version
87+
userAgent
88+
.append(" jvm/")
89+
.append(System.getProperty("java.version", "unknown").replace(" ", "_"));
90+
91+
// OS name
92+
userAgent.append(" os/").append(System.getProperty("os.name", "unknown").replace(" ", "_"));
93+
94+
// Custom user agent (if provided)
95+
if (connectionContext.getCustomerUserAgent() != null) {
96+
String[] parsed = parseCustomerUserAgent(connectionContext.getCustomerUserAgent());
97+
if (parsed != null) {
98+
try {
99+
userAgent.append(" ").append(parsed[0]).append("/").append(UserAgent.sanitize(parsed[1]));
100+
} catch (IllegalArgumentException e) {
101+
LOGGER.debug(
102+
"Failed to include customer userAgent entry {} in connector service UA, Error {}",
35103
connectionContext.getCustomerUserAgent(),
36-
StandardCharsets.UTF_8); // This is for encoded userAgentString
37-
int i = decodedUA.indexOf('/');
38-
String customerName = (i < 0) ? decodedUA : decodedUA.substring(0, i);
39-
String customerVersion = (i < 0) ? VERSION_FILLER : decodedUA.substring(i + 1);
40-
UserAgent.withOtherInfo(customerName, UserAgent.sanitize(customerVersion));
41-
} catch (Exception e) {
42-
LOGGER.debug(
43-
"Failed to set user agent for customer userAgent entry {}, Error {}",
44-
connectionContext.getCustomerUserAgent(),
45-
e);
104+
e);
105+
}
106+
}
46107
}
108+
109+
return userAgent.toString();
47110
}
48111

49112
/** Gets the user agent string for Databricks Driver HTTP Client. */

‎src/test/java/com/databricks/jdbc/common/util/UserAgentManagerTest.java‎

Lines changed: 51 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,7 @@ void testUserAgentSetsClientCorrectly() throws DatabricksSQLException {
8686
DatabricksConnectionContextFactory.create(CLUSTER_JDBC_URL, new Properties());
8787
UserAgentManager.setUserAgent(connectionContext);
8888
String userAgent = getUserAgentString();
89-
assertTrue(userAgent.contains("DatabricksJDBCDriverOSS/3.0.7"));
89+
assertTrue(userAgent.contains("DatabricksJDBCDriverOSS/"));
9090
assertTrue(userAgent.contains(" Java/THttpClient"));
9191
assertTrue(userAgent.contains(" MyApp/version"));
9292
assertTrue(userAgent.contains(" databricks-jdbc-http "));
@@ -97,7 +97,7 @@ void testUserAgentSetsClientCorrectly() throws DatabricksSQLException {
9797
DatabricksConnectionContextFactory.create(WAREHOUSE_JDBC_URL, new Properties());
9898
UserAgentManager.setUserAgent(connectionContext);
9999
userAgent = getUserAgentString();
100-
assertTrue(userAgent.contains("DatabricksJDBCDriverOSS/3.0.7"));
100+
assertTrue(userAgent.contains("DatabricksJDBCDriverOSS/"));
101101
assertTrue(userAgent.contains(" Java/THttpClient"));
102102
assertTrue(userAgent.contains(" MyApp/version"));
103103
assertTrue(userAgent.contains(" databricks-jdbc-http "));
@@ -108,12 +108,60 @@ void testUserAgentSetsClientCorrectly() throws DatabricksSQLException {
108108
DatabricksConnectionContextFactory.create(WAREHOUSE_JDBC_URL_WITH_SEA, new Properties());
109109
UserAgentManager.setUserAgent(connectionContext);
110110
userAgent = getUserAgentString();
111-
assertTrue(userAgent.contains("DatabricksJDBCDriverOSS/3.0.7"));
111+
assertTrue(userAgent.contains("DatabricksJDBCDriverOSS/"));
112112
assertTrue(userAgent.contains(" Java/SQLExecHttpClient"));
113113
assertTrue(userAgent.contains(" databricks-jdbc-http "));
114114
assertFalse(userAgent.contains("databricks-sdk-java"));
115115
}
116116

117+
@Test
118+
void testCustomUserAgentIncludedBeforeClientTypeEvaluation() throws DatabricksSQLException {
119+
// This test verifies that custom user agent is included in connector service requests
120+
// and maintains proper order: base -> client type -> custom
121+
122+
// Create connection context with custom user agent entry (use URL without existing UA)
123+
String jdbcUrlWithCustomUA =
124+
WAREHOUSE_JDBC_URL_WITH_THRIFT + ";useragententry=CustomTestApp/2.0.0";
125+
IDatabricksConnectionContext connectionContext =
126+
DatabricksConnectionContextFactory.create(jdbcUrlWithCustomUA, new Properties());
127+
128+
// Call setUserAgent
129+
UserAgentManager.setUserAgent(connectionContext);
130+
131+
// Get the final user agent string
132+
String userAgent = getUserAgentString();
133+
134+
// Verify custom user agent is included
135+
assertTrue(
136+
userAgent.contains("CustomTestApp/2.0.0"),
137+
"Custom user agent should be included: " + userAgent);
138+
139+
// Verify driver version is present
140+
assertTrue(
141+
userAgent.contains("DatabricksJDBCDriverOSS/"),
142+
"Driver version should be present: " + userAgent);
143+
144+
// Verify client type is included (proves getClientType was called)
145+
assertTrue(
146+
userAgent.contains("Java/THttpClient") || userAgent.contains("Java/SQLExecHttpClient"),
147+
"Client type should be included: " + userAgent);
148+
149+
// Verify JDBC HTTP client identifier
150+
assertTrue(
151+
userAgent.contains("databricks-jdbc-http"),
152+
"JDBC HTTP client identifier should be present: " + userAgent);
153+
154+
// Verify correct order: client type should come BEFORE custom user agent
155+
int clientTypeIndex =
156+
userAgent.contains("Java/THttpClient")
157+
? userAgent.indexOf("Java/THttpClient")
158+
: userAgent.indexOf("Java/SQLExecHttpClient");
159+
int customUAIndex = userAgent.indexOf("CustomTestApp/2.0.0");
160+
assertTrue(
161+
clientTypeIndex < customUAIndex,
162+
"Client type should appear before custom user agent. Order: " + userAgent);
163+
}
164+
117165
@Test
118166
void testUserAgentSetsCustomerInput() throws DatabricksSQLException {
119167
IDatabricksConnectionContext connectionContext =

0 commit comments

Comments
 (0)