Skip to content

Commit 924ac30

Browse files
samikshya-dbsamikshya-db
authored
Revert "[PECOBLR-1190] Make both telemetryClient and factory on a hos… (#1080)
…t level (#1075)" This reverts commit 409b69d. This is because we want to release this in the upcoming version ## Description <!-- Provide a brief summary of the changes made and the issue they aim to address.--> ## Testing <!-- Describe how the changes have been tested--> ## Additional Notes to the Reviewer <!-- Share any additional context or insights that may help the reviewer understand the changes better. This could include challenges faced, limitations, or compromises made during the development process. Also, mention any areas of the code that you would like the reviewer to focus on specifically. -->
1 parent 409b69d commit 924ac30

7 files changed

Lines changed: 32 additions & 109 deletions

File tree

‎NEXT_CHANGELOG.md‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,6 @@
99
- Added a new config attribute `DisableOauthRefreshToken` to control whether refresh tokens are requested in OAuth exchanges. By default, the driver does not include the `offline_access` scope. If `offline_access` is explicitly provided by the user, it is preserved and not removed.
1010

1111
### Updated
12-
- Minimized OAuth requests by reducing calls in feature flags and telemetry.
1312
- Updated sdk version from 0.65.0 to 0.69.0
1413

1514
### Fixed

‎development/.release-freeze.json‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
11
{
2-
"freeze": false,
3-
"reason": ""
2+
"freeze": true,
3+
"reason": "release 3.0.4"
44
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,6 @@ static void setFeatureFlagsContext(
7373
}
7474

7575
private static String keyOf(IDatabricksConnectionContext context) {
76-
return context.getHostForOAuth();
76+
return context.getComputeResource().getUniqueIdentifier();
7777
}
7878
}

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

Lines changed: 20 additions & 71 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,7 @@
88
import com.databricks.jdbc.log.JdbcLoggerFactory;
99
import com.databricks.sdk.core.DatabricksConfig;
1010
import com.google.common.annotations.VisibleForTesting;
11-
import java.util.Map;
12-
import java.util.concurrent.ConcurrentHashMap;
11+
import java.util.LinkedHashMap;
1312
import java.util.concurrent.ExecutorService;
1413
import java.util.concurrent.Executors;
1514
import java.util.concurrent.ThreadFactory;
@@ -24,10 +23,10 @@ public class TelemetryClientFactory {
2423
private static final TelemetryClientFactory INSTANCE = new TelemetryClientFactory();
2524

2625
@VisibleForTesting
27-
final Map<String, TelemetryClientHolder> telemetryClientHolders = new ConcurrentHashMap<>();
26+
final LinkedHashMap<String, TelemetryClient> telemetryClients = new LinkedHashMap<>();
2827

2928
@VisibleForTesting
30-
final Map<String, TelemetryClientHolder> noauthTelemetryClientHolders = new ConcurrentHashMap<>();
29+
final LinkedHashMap<String, TelemetryClient> noauthTelemetryClients = new LinkedHashMap<>();
3130

3231
private final ExecutorService telemetryExecutorService;
3332

@@ -60,60 +59,24 @@ public ITelemetryClient getTelemetryClient(IDatabricksConnectionContext connecti
6059
DatabricksConfig databricksConfig =
6160
TelemetryHelper.getDatabricksConfigSafely(connectionContext);
6261
if (databricksConfig != null) {
63-
String key = keyOf(connectionContext);
64-
TelemetryClientHolder holder =
65-
telemetryClientHolders.compute(
66-
key,
67-
(k, existing) -> {
68-
if (existing == null) {
69-
return new TelemetryClientHolder(
70-
new TelemetryClient(
71-
connectionContext, getTelemetryExecutorService(), databricksConfig),
72-
1);
73-
}
74-
existing.refCount.incrementAndGet();
75-
return existing;
76-
});
77-
return holder.client;
62+
return telemetryClients.computeIfAbsent(
63+
connectionContext.getConnectionUuid(),
64+
k ->
65+
new TelemetryClient(
66+
connectionContext, getTelemetryExecutorService(), databricksConfig));
7867
}
7968
// Use no-auth telemetry client if connection creation failed.
80-
String key = keyOf(connectionContext);
81-
TelemetryClientHolder holder =
82-
noauthTelemetryClientHolders.compute(
83-
key,
84-
(k, existing) -> {
85-
if (existing == null) {
86-
return new TelemetryClientHolder(
87-
new TelemetryClient(connectionContext, getTelemetryExecutorService()), 1);
88-
}
89-
existing.refCount.incrementAndGet();
90-
return existing;
91-
});
92-
return holder.client;
69+
return noauthTelemetryClients.computeIfAbsent(
70+
connectionContext.getConnectionUuid(),
71+
k -> new TelemetryClient(connectionContext, getTelemetryExecutorService()));
9372
}
9473

9574
public void closeTelemetryClient(IDatabricksConnectionContext connectionContext) {
96-
String key = keyOf(connectionContext);
97-
telemetryClientHolders.computeIfPresent(
98-
key,
99-
(k, holder) -> {
100-
if (holder.refCount.get() <= 1) {
101-
closeTelemetryClient(holder.client, "telemetry client");
102-
return null;
103-
}
104-
holder.refCount.decrementAndGet();
105-
return holder;
106-
});
107-
noauthTelemetryClientHolders.computeIfPresent(
108-
key,
109-
(k, holder) -> {
110-
if (holder.refCount.get() <= 1) {
111-
closeTelemetryClient(holder.client, "unauthenticated telemetry client");
112-
return null;
113-
}
114-
holder.refCount.decrementAndGet();
115-
return holder;
116-
});
75+
closeTelemetryClient(
76+
telemetryClients.remove(connectionContext.getConnectionUuid()), "telemetry client");
77+
closeTelemetryClient(
78+
noauthTelemetryClients.remove(connectionContext.getConnectionUuid()),
79+
"unauthenticated telemetry client");
11780
}
11881

11982
public ExecutorService getTelemetryExecutorService() {
@@ -146,12 +109,12 @@ static ITelemetryPushClient getTelemetryPushClient(
146109
@VisibleForTesting
147110
public void reset() {
148111
// Close all existing clients
149-
telemetryClientHolders.values().forEach(holder -> holder.client.close());
150-
noauthTelemetryClientHolders.values().forEach(holder -> holder.client.close());
112+
telemetryClients.values().forEach(TelemetryClient::close);
113+
noauthTelemetryClients.values().forEach(TelemetryClient::close);
151114

152115
// Clear the maps
153-
telemetryClientHolders.clear();
154-
noauthTelemetryClientHolders.clear();
116+
telemetryClients.clear();
117+
noauthTelemetryClients.clear();
155118
}
156119

157120
private void closeTelemetryClient(ITelemetryClient client, String clientType) {
@@ -163,18 +126,4 @@ private void closeTelemetryClient(ITelemetryClient client, String clientType) {
163126
}
164127
}
165128
}
166-
167-
private static final class TelemetryClientHolder {
168-
final TelemetryClient client;
169-
final AtomicInteger refCount;
170-
171-
TelemetryClientHolder(TelemetryClient client, int initialCount) {
172-
this.client = client;
173-
this.refCount = new AtomicInteger(initialCount);
174-
}
175-
}
176-
177-
private static String keyOf(IDatabricksConnectionContext context) {
178-
return context.getHostForOAuth();
179-
}
180129
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -88,7 +88,7 @@ private static void exportTelemetryEvent(
8888
TelemetryLogLevel logLevel) {
8989
if (connectionContext == null
9090
|| telemetryDetails == null
91-
|| logLevel.toInt() < connectionContext.getTelemetryLogLevel().toInt()) {
91+
|| logLevel.toInt() <= connectionContext.getTelemetryLogLevel().toInt()) {
9292
// We don't export telemetry logs in the following three scenarios:
9393
// 1. When the context is not set.
9494
// 2. When telemetry details are not set.

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

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -44,8 +44,8 @@ public void testGetNoOpTelemetryClient() throws Exception {
4444
ITelemetryClient telemetryClient =
4545
TelemetryClientFactory.getInstance().getTelemetryClient(context);
4646
assertInstanceOf(NoopTelemetryClient.class, telemetryClient);
47-
assertEquals(0, TelemetryClientFactory.getInstance().telemetryClientHolders.size());
48-
assertEquals(0, TelemetryClientFactory.getInstance().noauthTelemetryClientHolders.size());
47+
assertEquals(0, TelemetryClientFactory.getInstance().telemetryClients.size());
48+
assertEquals(0, TelemetryClientFactory.getInstance().noauthTelemetryClients.size());
4949
}
5050

5151
@Test
@@ -58,11 +58,11 @@ public void testGetAuthenticatedTelemetryClient() throws Exception {
5858
ITelemetryClient telemetryClient =
5959
TelemetryClientFactory.getInstance().getTelemetryClient(context);
6060
assertInstanceOf(TelemetryClient.class, telemetryClient);
61-
assertEquals(1, TelemetryClientFactory.getInstance().telemetryClientHolders.size());
62-
assertEquals(0, TelemetryClientFactory.getInstance().noauthTelemetryClientHolders.size());
61+
assertEquals(1, TelemetryClientFactory.getInstance().telemetryClients.size());
62+
assertEquals(0, TelemetryClientFactory.getInstance().noauthTelemetryClients.size());
6363
TelemetryClientFactory.getInstance().closeTelemetryClient(context);
64-
assertEquals(0, TelemetryClientFactory.getInstance().telemetryClientHolders.size());
65-
assertEquals(0, TelemetryClientFactory.getInstance().noauthTelemetryClientHolders.size());
64+
assertEquals(0, TelemetryClientFactory.getInstance().telemetryClients.size());
65+
assertEquals(0, TelemetryClientFactory.getInstance().noauthTelemetryClients.size());
6666
TelemetryClientFactory.getInstance().closeTelemetryClient(context);
6767
}
6868

@@ -80,8 +80,8 @@ public void testGetNoOpTelemetryClientWhenDatabricksConfigIsNull() throws Except
8080
TelemetryClientFactory.getInstance().getTelemetryClient(context);
8181

8282
assertInstanceOf(NoopTelemetryClient.class, telemetryClient);
83-
assertEquals(0, TelemetryClientFactory.getInstance().telemetryClientHolders.size());
84-
assertEquals(0, TelemetryClientFactory.getInstance().noauthTelemetryClientHolders.size());
83+
assertEquals(0, TelemetryClientFactory.getInstance().telemetryClients.size());
84+
assertEquals(0, TelemetryClientFactory.getInstance().noauthTelemetryClients.size());
8585
TelemetryClientFactory.getInstance().closeTelemetryClient(context);
8686
}
8787
}

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

Lines changed: 0 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -205,7 +205,6 @@ public void testGetDatabricksConfigSafely_HandlesNullContext() {
205205
public void testTelemetryNotAllowedUsecase() {
206206
// Clear thread context to ensure telemetry is not allowed
207207
when(connectionContext.forceEnableTelemetry()).thenReturn(false);
208-
when(connectionContext.getHostForOAuth()).thenReturn("test-host");
209208
when(connectionContext.isTelemetryEnabled()).thenReturn(false);
210209
assertFalse(isTelemetryAllowedForConnection(connectionContext));
211210
when(connectionContext.getComputeResource()).thenReturn(WAREHOUSE_COMPUTE);
@@ -217,7 +216,6 @@ public void testTelemetryNotAllowedUsecase() {
217216
public void testTelemetryAllowedWithForceTelemetryFlag() {
218217
when(connectionContext.getTelemetryLogLevel()).thenReturn(TelemetryLogLevel.DEBUG);
219218
when(connectionContext.getComputeResource()).thenReturn(WAREHOUSE_COMPUTE);
220-
when(connectionContext.getHostForOAuth()).thenReturn("test-host");
221219
enableFeatureFlagForTesting(connectionContext, Collections.emptyMap());
222220
assertTrue(() -> isTelemetryAllowedForConnection(connectionContext));
223221
}
@@ -262,29 +260,6 @@ void testExportTelemetryLog_EmitsWhenEventLevelHigherThanConfigured() {
262260
}
263261
}
264262

265-
@Test
266-
void testExportTelemetryLog_EmitsWhenEventLevelEqualToConfigured() {
267-
// Configured level: INFO (4); Event level: INFO (4) -> should export (4 >= 4)
268-
when(connectionContext.getTelemetryLogLevel()).thenReturn(TelemetryLogLevel.INFO);
269-
StatementTelemetryDetails details =
270-
new StatementTelemetryDetails("stmt-eq").setOperationLatencyMillis(15L);
271-
272-
ITelemetryClient clientMock = mock(ITelemetryClient.class);
273-
TelemetryClientFactory factoryMock = mock(TelemetryClientFactory.class);
274-
275-
try (MockedStatic<TelemetryClientFactory> mocked =
276-
Mockito.mockStatic(TelemetryClientFactory.class)) {
277-
mocked.when(TelemetryClientFactory::getInstance).thenReturn(factoryMock);
278-
when(factoryMock.getTelemetryClient(connectionContext)).thenReturn(clientMock);
279-
280-
TelemetryHelper.exportTelemetryLog(details, TelemetryLogLevel.INFO);
281-
282-
mocked.verify(TelemetryClientFactory::getInstance, times(1));
283-
verify(factoryMock, times(1)).getTelemetryClient(connectionContext);
284-
verify(clientMock, times(1)).exportEvent(any());
285-
}
286-
}
287-
288263
static Stream<Object[]> failureLogParameters() {
289264
return Stream.of(
290265
new Object[] {"test-statement-id", null},

0 commit comments

Comments
 (0)