Skip to content

Commit 6f3ac5a

Browse files
keshavdandevakeshavdandeva
committed
chore: use global handler
1 parent 93b15fb commit 6f3ac5a

4 files changed

Lines changed: 128 additions & 93 deletions

File tree

‎java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,6 @@
6464
import java.util.concurrent.ConcurrentHashMap;
6565
import java.util.concurrent.Executor;
6666
import java.util.concurrent.TimeUnit;
67-
import java.util.logging.Logger;
6867

6968
/**
7069
* An implementation of {@link java.sql.Connection} for establishing a connection with BigQuery and
@@ -955,6 +954,7 @@ private void closeImpl() throws SQLException {
955954
} finally {
956955
BigQueryJdbcMdc.removeInstance(this);
957956
BigQueryJdbcRootLogger.closeConnectionHandler(this.connectionId);
957+
BigQueryJdbcOpenTelemetry.unregisterConnection(this.connectionId);
958958
}
959959
this.isClosed = true;
960960
}
@@ -1059,10 +1059,8 @@ private BigQuery getBigQueryConnection() {
10591059
this.enableGcpTraceExporter, this.enableGcpLogExporter, this.customOpenTelemetry);
10601060

10611061
if (this.enableGcpLogExporter || this.customOpenTelemetry != null) {
1062-
OpenTelemetryJulHandler otelHandler =
1063-
new OpenTelemetryJulHandler(
1064-
null, openTelemetry, this.enableGcpLogExporter, this.connectionId);
1065-
Logger.getLogger("com.google.cloud.bigquery").addHandler(otelHandler);
1062+
BigQueryJdbcOpenTelemetry.registerConnection(
1063+
this.connectionId, openTelemetry, null, this.enableGcpLogExporter);
10661064
}
10671065

10681066
if (this.enableGcpTraceExporter || this.customOpenTelemetry != null) {

‎java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcOpenTelemetry.java‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,15 +16,60 @@
1616

1717
package com.google.cloud.bigquery.jdbc;
1818

19+
import com.google.cloud.logging.Logging;
1920
import io.opentelemetry.api.OpenTelemetry;
2021
import io.opentelemetry.api.trace.Tracer;
22+
import java.util.Collection;
23+
import java.util.concurrent.ConcurrentHashMap;
24+
import java.util.logging.Logger;
2125

2226
public class BigQueryJdbcOpenTelemetry {
2327

2428
static final String INSTRUMENTATION_SCOPE_NAME = "com.google.cloud.bigquery.jdbc";
29+
static final String BIGQUERY_NAMESPACE = "com.google.cloud.bigquery";
30+
31+
static class TelemetryConfig {
32+
final OpenTelemetry openTelemetry;
33+
final Logging loggingClient;
34+
final boolean isGcpFallback;
35+
36+
TelemetryConfig(OpenTelemetry openTelemetry, Logging loggingClient, boolean isGcpFallback) {
37+
this.openTelemetry = openTelemetry;
38+
this.loggingClient = loggingClient;
39+
this.isGcpFallback = isGcpFallback;
40+
}
41+
}
42+
43+
private static final ConcurrentHashMap<String, TelemetryConfig> connectionConfigs =
44+
new ConcurrentHashMap<>();
2545

2646
private BigQueryJdbcOpenTelemetry() {}
2747

48+
static {
49+
Logger.getLogger(BIGQUERY_NAMESPACE).addHandler(new OpenTelemetryJulHandler());
50+
}
51+
52+
public static void registerConnection(
53+
String connectionId,
54+
OpenTelemetry openTelemetry,
55+
Logging loggingClient,
56+
boolean isGcpFallback) {
57+
connectionConfigs.put(
58+
connectionId, new TelemetryConfig(openTelemetry, loggingClient, isGcpFallback));
59+
}
60+
61+
public static void unregisterConnection(String connectionId) {
62+
connectionConfigs.remove(connectionId);
63+
}
64+
65+
public static TelemetryConfig getConnectionConfig(String connectionId) {
66+
return connectionConfigs.get(connectionId);
67+
}
68+
69+
public static Collection<TelemetryConfig> getRegisteredConfigs() {
70+
return connectionConfigs.values();
71+
}
72+
2873
/**
2974
* Initializes or returns the OpenTelemetry instance based on hybrid logic. Prefer
3075
* customOpenTelemetry if provided; fallback to an auto-configured GCP exporter if requested.

‎java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/OpenTelemetryJulHandler.java‎

Lines changed: 45 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import io.opentelemetry.api.OpenTelemetry;
2323
import io.opentelemetry.api.baggage.Baggage;
2424
import io.opentelemetry.api.common.AttributeKey;
25+
import io.opentelemetry.api.logs.LogRecordBuilder;
2526
import io.opentelemetry.api.logs.Logger;
2627
import io.opentelemetry.api.logs.Severity;
2728
import io.opentelemetry.api.trace.Span;
@@ -39,54 +40,46 @@
3940
*/
4041
public class OpenTelemetryJulHandler extends Handler {
4142

42-
private final Logging loggingClient;
43-
private final OpenTelemetry openTelemetry;
44-
private final boolean isGcpFallback;
45-
private final String expectedConnectionId;
46-
47-
public OpenTelemetryJulHandler(
48-
Logging loggingClient,
49-
OpenTelemetry openTelemetry,
50-
boolean isGcpFallback,
51-
String expectedConnectionId) {
52-
this.loggingClient = loggingClient;
53-
this.openTelemetry = openTelemetry;
54-
this.isGcpFallback = isGcpFallback;
55-
this.expectedConnectionId = expectedConnectionId;
56-
}
43+
public OpenTelemetryJulHandler() {}
5744

5845
@Override
5946
public void publish(LogRecord record) {
6047
if (!isLoggable(record)) {
6148
return;
6249
}
6350

64-
Context context = Context.current();
65-
SpanContext spanContext = Span.fromContext(context).getSpanContext();
66-
67-
String traceId = spanContext.isValid() ? spanContext.getTraceId() : null;
68-
String spanId = spanContext.isValid() ? spanContext.getSpanId() : null;
69-
7051
// Extract connection ID from baggage
71-
String connectionId = Baggage.fromContext(context).getEntryValue("jdbc.connection_id");
52+
String connectionId =
53+
Baggage.fromContext(Context.current()).getEntryValue("jdbc.connection_id");
7254

7355
// Fallback to MDC if not in baggage (if MDC is available and used)
7456
if (connectionId == null) {
7557
connectionId = BigQueryJdbcMdc.getConnectionId();
7658
}
7759

78-
if (expectedConnectionId != null && !expectedConnectionId.equals(connectionId)) {
60+
if (connectionId == null) {
61+
return;
62+
}
63+
64+
BigQueryJdbcOpenTelemetry.TelemetryConfig config =
65+
BigQueryJdbcOpenTelemetry.getConnectionConfig(connectionId);
66+
if (config == null) {
7967
return;
8068
}
8169

82-
if (isGcpFallback && loggingClient != null) {
83-
publishToGcp(record, traceId, spanId, connectionId);
84-
} else if (openTelemetry != null) {
85-
publishToOTel(record, traceId, spanId, connectionId);
70+
if (config.isGcpFallback && config.loggingClient != null) {
71+
publishToGcp(record, connectionId, config.loggingClient);
72+
} else if (config.openTelemetry != null) {
73+
publishToOTel(record, connectionId, config.openTelemetry);
8674
}
8775
}
8876

89-
private void publishToGcp(LogRecord record, String traceId, String spanId, String connectionId) {
77+
private void publishToGcp(LogRecord record, String connectionId, Logging loggingClient) {
78+
Context context = Context.current();
79+
SpanContext spanContext = Span.fromContext(context).getSpanContext();
80+
String traceId = spanContext.isValid() ? spanContext.getTraceId() : null;
81+
String spanId = spanContext.isValid() ? spanContext.getSpanId() : null;
82+
9083
// TODO(b/491238299): May require refinement for structured logging or error handling
9184
if (loggingClient == null) {
9285
return;
@@ -119,7 +112,7 @@ private com.google.cloud.logging.Severity mapGcpSeverity(Level level) {
119112
return com.google.cloud.logging.Severity.DEBUG;
120113
}
121114

122-
private void publishToOTel(LogRecord record, String traceId, String spanId, String connectionId) {
115+
private void publishToOTel(LogRecord record, String connectionId, OpenTelemetry openTelemetry) {
123116
if (openTelemetry == null) {
124117
return;
125118
}
@@ -128,17 +121,24 @@ private void publishToOTel(LogRecord record, String traceId, String spanId, Stri
128121
Logger logger =
129122
openTelemetry
130123
.getLogsBridge()
131-
.get(loggerName != null ? loggerName : "com.google.cloud.bigquery.jdbc");
132-
133-
logger
134-
.logRecordBuilder()
135-
.setBody(record.getMessage())
136-
.setSeverity(mapSeverity(record.getLevel()))
137-
.setTimestamp(Instant.ofEpochMilli(record.getMillis()))
138-
.setContext(Context.current())
139-
.setAttribute(
140-
AttributeKey.stringKey("jdbc.connection_id"), connectionId != null ? connectionId : "")
141-
.emit();
124+
.get(
125+
loggerName != null
126+
? loggerName
127+
: BigQueryJdbcOpenTelemetry.INSTRUMENTATION_SCOPE_NAME);
128+
129+
LogRecordBuilder builder =
130+
logger
131+
.logRecordBuilder()
132+
.setBody(record.getMessage())
133+
.setSeverity(mapSeverity(record.getLevel()))
134+
.setTimestamp(Instant.ofEpochMilli(record.getMillis()))
135+
.setContext(Context.current());
136+
137+
if (connectionId != null) {
138+
builder.setAttribute(AttributeKey.stringKey("jdbc.connection_id"), connectionId);
139+
}
140+
141+
builder.emit();
142142
}
143143

144144
private Severity mapSeverity(Level level) {
@@ -154,8 +154,11 @@ private Severity mapSeverity(Level level) {
154154

155155
@Override
156156
public void flush() {
157-
if (isGcpFallback && loggingClient != null) {
158-
loggingClient.flush();
157+
for (BigQueryJdbcOpenTelemetry.TelemetryConfig config :
158+
BigQueryJdbcOpenTelemetry.getRegisteredConfigs()) {
159+
if (config.isGcpFallback && config.loggingClient != null) {
160+
config.loggingClient.flush();
161+
}
159162
}
160163
}
161164

‎java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/OpenTelemetryJulHandlerTest.java‎

Lines changed: 35 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -31,8 +31,8 @@
3131
import io.opentelemetry.sdk.logs.data.LogRecordData;
3232
import io.opentelemetry.sdk.testing.junit5.OpenTelemetryExtension;
3333
import java.util.List;
34-
import java.util.logging.Level;
35-
import java.util.logging.LogRecord;
34+
import java.util.logging.Logger;
35+
import org.junit.jupiter.api.AfterEach;
3636
import org.junit.jupiter.api.Test;
3737
import org.junit.jupiter.api.extension.RegisterExtension;
3838
import org.mockito.ArgumentCaptor;
@@ -42,71 +42,60 @@ public class OpenTelemetryJulHandlerTest {
4242
@RegisterExtension
4343
static final OpenTelemetryExtension otelTesting = OpenTelemetryExtension.create();
4444

45+
private static final Logger logger = Logger.getLogger("com.google.cloud.bigquery");
46+
47+
@AfterEach
48+
public void tearDown() {
49+
BigQueryJdbcOpenTelemetry.unregisterConnection("test-uuid");
50+
BigQueryJdbcOpenTelemetry.unregisterConnection("wrong-uuid");
51+
BigQueryJdbcOpenTelemetry.unregisterConnection("gcp-uuid");
52+
}
53+
4554
@Test
4655
public void testPublishToOTel() {
47-
OpenTelemetryJulHandler handler =
48-
new OpenTelemetryJulHandler(null, otelTesting.getOpenTelemetry(), false, null);
56+
BigQueryJdbcOpenTelemetry.registerConnection(
57+
"test-uuid", otelTesting.getOpenTelemetry(), null, false);
4958

50-
LogRecord record = new LogRecord(Level.INFO, "Test message");
51-
record.setLoggerName("test.logger");
52-
record.setMillis(System.currentTimeMillis());
53-
54-
handler.publish(record);
59+
Baggage baggage = Baggage.builder().put("jdbc.connection_id", "test-uuid").build();
60+
try (Scope scope = baggage.makeCurrent()) {
61+
logger.info("Test message");
62+
}
5563

5664
List<LogRecordData> logs = otelTesting.getLogRecords();
5765
assertEquals(1, logs.size());
5866
LogRecordData log = logs.get(0);
5967
assertEquals("Test message", log.getBody().asString());
60-
assertEquals("test.logger", log.getInstrumentationScopeInfo().getName());
6168
assertEquals(Severity.INFO, log.getSeverity());
69+
assertEquals(
70+
"test-uuid", log.getAttributes().get(AttributeKey.stringKey("jdbc.connection_id")));
6271
}
6372

6473
@Test
65-
public void testPublishToOTelWithFiltering() {
66-
// Handler expects "correct-uuid"
67-
OpenTelemetryJulHandler handler =
68-
new OpenTelemetryJulHandler(null, otelTesting.getOpenTelemetry(), false, "correct-uuid");
69-
70-
// Simulate log with NO connection ID
71-
LogRecord record = new LogRecord(Level.INFO, "Test message");
72-
record.setLoggerName("test.logger");
73-
74-
handler.publish(record);
75-
76-
List<LogRecordData> logs = otelTesting.getLogRecords();
77-
assertTrue(logs.isEmpty()); // Should be filtered out
78-
}
79-
80-
@Test
81-
public void testPublishToOTelWithBaggage() {
82-
OpenTelemetryJulHandler handler =
83-
new OpenTelemetryJulHandler(null, otelTesting.getOpenTelemetry(), false, null);
74+
public void testPublishWithFiltering() {
75+
// Register for "test-uuid"
76+
BigQueryJdbcOpenTelemetry.registerConnection(
77+
"test-uuid", otelTesting.getOpenTelemetry(), null, false);
8478

85-
Baggage baggage = Baggage.builder().put("jdbc.connection_id", "test-uuid").build();
79+
// Log with WRONG connection ID
80+
Baggage baggage = Baggage.builder().put("jdbc.connection_id", "wrong-uuid").build();
8681
try (Scope scope = baggage.makeCurrent()) {
87-
LogRecord record = new LogRecord(Level.INFO, "Test message");
88-
record.setLoggerName("test.logger");
89-
handler.publish(record);
90-
91-
List<LogRecordData> logs = otelTesting.getLogRecords();
92-
assertEquals(1, logs.size());
93-
LogRecordData log = logs.get(0);
94-
assertEquals(
95-
"test-uuid", log.getAttributes().get(AttributeKey.stringKey("jdbc.connection_id")));
82+
logger.info("Test message");
9683
}
84+
85+
List<LogRecordData> logs = otelTesting.getLogRecords();
86+
assertTrue(logs.isEmpty()); // Should be filtered out because "wrong-uuid" has no config
9787
}
9888

9989
@Test
10090
public void testPublishToGcp() {
10191
Logging loggingClient = mock(Logging.class);
102-
OpenTelemetryJulHandler handler =
103-
new OpenTelemetryJulHandler(loggingClient, otelTesting.getOpenTelemetry(), true, null);
104-
105-
LogRecord record = new LogRecord(Level.INFO, "Test message");
106-
record.setLoggerName("test.logger");
107-
record.setMillis(System.currentTimeMillis());
92+
BigQueryJdbcOpenTelemetry.registerConnection(
93+
"gcp-uuid", otelTesting.getOpenTelemetry(), loggingClient, true);
10894

109-
handler.publish(record);
95+
Baggage baggage = Baggage.builder().put("jdbc.connection_id", "gcp-uuid").build();
96+
try (Scope scope = baggage.makeCurrent()) {
97+
logger.info("Test message");
98+
}
11099

111100
ArgumentCaptor<Iterable<LogEntry>> captor = ArgumentCaptor.forClass(Iterable.class);
112101
verify(loggingClient).write(captor.capture());

0 commit comments

Comments
 (0)