Skip to content

Commit 4c634e3

Browse files
samikshya-dbsamikshya-db
authored
Add client side telemetry fields inline with proto (#832)
1 parent e7d3641 commit 4c634e3

4 files changed

Lines changed: 121 additions & 6 deletions

File tree

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,8 @@
11
package com.databricks.jdbc.common.util;
22

3+
import java.util.Collections;
4+
import java.util.List;
5+
36
public class StringUtil {
47
public static String convertJdbcEscapeSequences(String sql) {
58
// Replace JDBC escape sequences;
@@ -40,6 +43,13 @@ public static boolean checkPrefixMatch(String prefix, String fileName, boolean c
4043
/* lengthToMatch= */ prefix.length());
4144
}
4245

46+
public static List<String> split(String value) {
47+
if (value == null) {
48+
return Collections.emptyList();
49+
}
50+
return List.of(value.split(","));
51+
}
52+
4353
/** Function to return the folder name from the path */
4454
public static String getFolderNameFromPath(String path) {
4555
if (path == null) return "";

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

Lines changed: 95 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,16 +2,16 @@
22

33
import com.databricks.jdbc.common.AuthFlow;
44
import com.databricks.jdbc.common.AuthMech;
5-
import com.databricks.jdbc.common.DatabricksClientType;
65
import com.databricks.sdk.support.ToStringer;
76
import com.fasterxml.jackson.annotation.JsonProperty;
7+
import java.util.List;
88

99
public class DriverConnectionParameters {
1010
@JsonProperty("http_path")
1111
String httpPath;
1212

1313
@JsonProperty("mode")
14-
DatabricksClientType driverMode;
14+
String driverMode;
1515

1616
@JsonProperty("host_info")
1717
HostDetails hostDetails;
@@ -106,13 +106,43 @@ public class DriverConnectionParameters {
106106
@JsonProperty("enable_token_cache")
107107
boolean enableTokenCache;
108108

109+
@JsonProperty("auth_endpoint")
110+
String authEndpoint;
111+
112+
@JsonProperty("token_endpoint")
113+
String tokenEndpoint;
114+
115+
@JsonProperty("non_proxy_hosts")
116+
List<String> nonProxyHosts;
117+
118+
@JsonProperty("http_connection_pool_size")
119+
int httpConnectionPoolSize;
120+
121+
@JsonProperty("enable_sea_hybrid_results")
122+
boolean enableSeaHybridResults;
123+
124+
@JsonProperty("enable_complex_datatype_support")
125+
boolean enableComplexSupport;
126+
127+
@JsonProperty("allow_self_signed_support")
128+
boolean allowSelfSignedSupport;
129+
130+
@JsonProperty("use_system_trust_store")
131+
boolean useSystemTrustStore;
132+
133+
@JsonProperty("rows_fetched_per_block")
134+
int rowsFetchedPerBlock;
135+
136+
@JsonProperty("async_poll_interval_millis")
137+
int asyncPollIntervalMillis;
138+
109139
public DriverConnectionParameters setHttpPath(String httpPath) {
110140
this.httpPath = httpPath;
111141
return this;
112142
}
113143

114-
public DriverConnectionParameters setDriverMode(DatabricksClientType clientType) {
115-
this.driverMode = clientType;
144+
public DriverConnectionParameters setDriverMode(String clientType) {
145+
this.driverMode = clientType.toString();
116146
return this;
117147
}
118148

@@ -276,6 +306,56 @@ public DriverConnectionParameters setEnableTokenCache(boolean enableTokenCache)
276306
return this;
277307
}
278308

309+
public DriverConnectionParameters setAuthEndpoint(String authEndpoint) {
310+
this.authEndpoint = authEndpoint;
311+
return this;
312+
}
313+
314+
public DriverConnectionParameters setTokenEndpoint(String tokenEndpoint) {
315+
this.tokenEndpoint = tokenEndpoint;
316+
return this;
317+
}
318+
319+
public DriverConnectionParameters setNonProxyHosts(List<String> nonProxyHosts) {
320+
this.nonProxyHosts = nonProxyHosts;
321+
return this;
322+
}
323+
324+
public DriverConnectionParameters setHttpConnectionPoolSize(int httpConnectionPoolSize) {
325+
this.httpConnectionPoolSize = httpConnectionPoolSize;
326+
return this;
327+
}
328+
329+
public DriverConnectionParameters setEnableSeaHybridResults(boolean enableSeaHybridResults) {
330+
this.enableSeaHybridResults = enableSeaHybridResults;
331+
return this;
332+
}
333+
334+
public DriverConnectionParameters setEnableComplexSupport(boolean enableComplexSupport) {
335+
this.enableComplexSupport = enableComplexSupport;
336+
return this;
337+
}
338+
339+
public DriverConnectionParameters setAllowSelfSignedSupport(boolean allowSelfSignedSupport) {
340+
this.allowSelfSignedSupport = allowSelfSignedSupport;
341+
return this;
342+
}
343+
344+
public DriverConnectionParameters setUseSystemTrustStore(boolean useSystemTrustStore) {
345+
this.useSystemTrustStore = useSystemTrustStore;
346+
return this;
347+
}
348+
349+
public DriverConnectionParameters setRowsFetchedPerBlock(int rowsFetchedPerBlock) {
350+
this.rowsFetchedPerBlock = rowsFetchedPerBlock;
351+
return this;
352+
}
353+
354+
public DriverConnectionParameters setAsyncPollIntervalMillis(int asyncPollIntervalMillis) {
355+
this.asyncPollIntervalMillis = asyncPollIntervalMillis;
356+
return this;
357+
}
358+
279359
@Override
280360
public String toString() {
281361
return new ToStringer(DriverConnectionParameters.class)
@@ -300,6 +380,7 @@ public String toString() {
300380
.add("acceptUndeterminedCertificateRevocation", acceptUndeterminedCertificateRevocation)
301381
.add("enableArrow", enableArrow)
302382
.add("enableDirectResults", enableDirectResults)
383+
.add("enableJwtAssertion", enableJwtAssertion)
303384
.add("jwtKeyFile", jwtKeyFile)
304385
.add("jwtAlgorithm", jwtAlgorithm)
305386
.add("googleServiceAccount", googleServiceAccount)
@@ -311,6 +392,16 @@ public String toString() {
311392
.add("stringColumnLength", stringColumnLength)
312393
.add("socketTimeout", socketTimeout)
313394
.add("enableTokenCache", enableTokenCache)
395+
.add("authEndpoint", authEndpoint)
396+
.add("tokenEndpoint", tokenEndpoint)
397+
.add("nonProxyHosts", nonProxyHosts)
398+
.add("httpConnectionPoolSize", httpConnectionPoolSize)
399+
.add("enableSeaHybridResults", enableSeaHybridResults)
400+
.add("enableComplexSupport", enableComplexSupport)
401+
.add("allowSelfSignedSupport", allowSelfSignedSupport)
402+
.add("useSystemTrustStore", useSystemTrustStore)
403+
.add("rowsFetchedPerBlock", rowsFetchedPerBlock)
404+
.add("asyncPollIntervalMillis", asyncPollIntervalMillis)
314405
.toString();
315406
}
316407
}

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

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import com.databricks.jdbc.api.internal.IDatabricksConnectionContext;
44
import com.databricks.jdbc.common.util.DatabricksThreadContextHolder;
55
import com.databricks.jdbc.common.util.DriverUtil;
6+
import com.databricks.jdbc.common.util.StringUtil;
67
import com.databricks.jdbc.dbclient.impl.common.StatementId;
78
import com.databricks.jdbc.exception.DatabricksParsingException;
89
import com.databricks.jdbc.model.telemetry.*;
@@ -209,7 +210,17 @@ private static DriverConnectionParameters buildDriverConnectionParameters(
209210
.setCheckCertificateRevocation(connectionContext.checkCertificateRevocation())
210211
.setAcceptUndeterminedCertificateRevocation(
211212
connectionContext.acceptUndeterminedCertificateRevocation())
212-
.setDriverMode(connectionContext.getClientType())
213+
.setDriverMode(connectionContext.getClientType().toString())
214+
.setAuthEndpoint(connectionContext.getAuthEndpoint())
215+
.setTokenEndpoint(connectionContext.getTokenEndpoint())
216+
.setNonProxyHosts(StringUtil.split(connectionContext.getNonProxyHosts()))
217+
.setHttpConnectionPoolSize(connectionContext.getHttpConnectionPoolSize())
218+
.setEnableSeaHybridResults(connectionContext.isSqlExecHybridResultsEnabled())
219+
.setEnableComplexSupport(connectionContext.isComplexDatatypeSupportEnabled())
220+
.setAllowSelfSignedSupport(connectionContext.allowSelfSignedCerts())
221+
.setUseSystemTrustStore(connectionContext.useSystemTrustStore())
222+
.setRowsFetchedPerBlock(connectionContext.getRowsFetchedPerBlock())
223+
.setAsyncPollIntervalMillis(connectionContext.getAsyncExecPollInterval())
213224
.setEnableTokenCache(connectionContext.isTokenCacheEnabled())
214225
.setHttpPath(connectionContext.getHttpPath());
215226
if (connectionContext.useJWTAssertion()) {

‎src/test/java/com/databricks/jdbc/telemetry/TelemetryHelperTest.java‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
import static org.mockito.Mockito.*;
66

77
import com.databricks.jdbc.api.internal.IDatabricksConnectionContext;
8+
import com.databricks.jdbc.common.DatabricksClientType;
89
import com.databricks.jdbc.common.StatementType;
910
import com.databricks.jdbc.exception.DatabricksParsingException;
1011
import com.databricks.jdbc.model.telemetry.SqlExecutionEvent;
@@ -27,6 +28,7 @@ void testInitialTelemetryLogDoesNotThrowError() {
2728
when(connectionContext.getProxyAuthType()).thenReturn(ProxyConfig.ProxyAuthType.BASIC);
2829
when(connectionContext.getProxyPort()).thenReturn(443);
2930
when(connectionContext.getProxyHost()).thenReturn(TEST_STRING);
31+
when(connectionContext.getClientType()).thenReturn(DatabricksClientType.SEA);
3032
when(connectionContext.getUseCloudFetchProxy()).thenReturn(true);
3133
when(connectionContext.getCloudFetchProxyAuthType())
3234
.thenReturn(ProxyConfig.ProxyAuthType.BASIC);
@@ -43,6 +45,7 @@ void testInitialTelemetryLogWithNullContextDoesNotThrowError() {
4345
@Test
4446
void testHostFetchThrowsErrorInTelemetryLog() throws DatabricksParsingException {
4547
when(connectionContext.getConnectionUuid()).thenReturn(UUID.randomUUID().toString());
48+
when(connectionContext.getClientType()).thenReturn(DatabricksClientType.SEA);
4649
when(connectionContext.getHostUrl())
4750
.thenThrow(
4851
new DatabricksParsingException(TEST_STRING, DatabricksDriverErrorCode.INVALID_STATE));
@@ -52,14 +55,14 @@ void testHostFetchThrowsErrorInTelemetryLog() throws DatabricksParsingException
5255
@Test
5356
void testLatencyTelemetryLogDoesNotThrowError() {
5457
when(connectionContext.getConnectionUuid()).thenReturn(TEST_STRING);
58+
when(connectionContext.getClientType()).thenReturn(DatabricksClientType.SEA);
5559
SqlExecutionEvent event = new SqlExecutionEvent().setDriverStatementType(StatementType.QUERY);
5660
assertDoesNotThrow(() -> TelemetryHelper.exportLatencyLog(connectionContext, 150, event, null));
5761
}
5862

5963
@Test
6064
void testErrorTelemetryLogDoesNotThrowError() {
6165
when(connectionContext.getConnectionUuid()).thenReturn(TEST_STRING);
62-
SqlExecutionEvent event = new SqlExecutionEvent().setDriverStatementType(StatementType.QUERY);
6366
assertDoesNotThrow(
6467
() -> TelemetryHelper.exportFailureLog(connectionContext, TEST_STRING, TEST_STRING));
6568
}

0 commit comments

Comments
 (0)