Skip to content

Commit 27ed8de

Browse files
keshavdandevakeshavdandeva
authored
feat(bqjdbc): Complete OpenTelemetry instrumentation and context propagation (#13187)
b/496720140 ## Changes ### Context Propagation & Session Tracking * **Baggage Injection**: Injected the generated Connection UUID into OpenTelemetry Baggage upon `BigQueryConnection` initialization to enable reliable log correlation. * **Log Handler Update**: Updated `OpenTelemetryJulHandler` to rely on Baggage for retrieving the connection ID, removing the legacy MDC fallback. * **Thread Pool Audit**: Wrapped tasks submitted to background executors in `BigQueryDatabaseMetaData` with `Context.current().wrap()`, ensuring trace context is not lost during parallel metadata fetching. ### Span Enrichment & Semantic Conventions * **Attributes**: Enriched JDBC spans with standard attributes: `db.system = "bigquery"`, `db.connection_id`, and `db.application` (derived from `partnerToken` or falling back to `"Google-BigQuery-JDBC-Driver"`). * **Scope Separation**: Implemented separate tracers for the JDBC driver (`com.google.cloud.bigquery.jdbc`) and the SDK (`com.google.cloud.bigquery`) to allow clean filtering in tracing UIs while maintaining correlation. ### Instrumentation * **PreparedStatement**: Added missing instrumentation for `BigQueryPreparedStatement` execution methods (`execute`, `executeQuery`, `executeLargeUpdate`) to generate spans. ### Refactoring & Cleanups * **Centralized Tracing**: Created a centralized `withTracing` helper in `BigQueryJdbcOpenTelemetry.java` to eliminate duplicated tracing logic in `BigQueryStatement` and `BigQueryDatabaseMetaData`. * **Constants**: Defined all semantic convention keys as constants in `BigQueryJdbcOpenTelemetry.java` to eliminate magic strings from method bodies. * **Simplifications**: Simplified redundant boolean checks in `BigQueryConnection.java`.
1 parent e5996c4 commit 27ed8de

8 files changed

Lines changed: 249 additions & 73 deletions

File tree

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

Lines changed: 22 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,9 @@
4343
import com.google.cloud.http.HttpTransportOptions;
4444
import com.google.cloud.logging.Logging;
4545
import io.opentelemetry.api.OpenTelemetry;
46+
import io.opentelemetry.api.baggage.Baggage;
4647
import io.opentelemetry.api.trace.Tracer;
48+
import io.opentelemetry.context.Context;
4749
import java.io.IOException;
4850
import java.io.InputStream;
4951
import java.sql.CallableStatement;
@@ -149,6 +151,7 @@ public class BigQueryConnection extends BigQueryNoOpsConnection {
149151
boolean enableGcpLogExporter;
150152
OpenTelemetry customOpenTelemetry;
151153
private OpenTelemetry openTelemetry;
154+
private Context otelContext;
152155
Tracer tracer =
153156
OpenTelemetry.noop().getTracer(BigQueryJdbcOpenTelemetry.INSTRUMENTATION_SCOPE_NAME);
154157
DatabaseMetaData databaseMetaData;
@@ -161,6 +164,11 @@ public class BigQueryConnection extends BigQueryNoOpsConnection {
161164

162165
BigQueryConnection(String url, DataSource ds) throws IOException {
163166
this.connectionId = UUID.randomUUID().toString();
167+
Baggage baggage =
168+
Baggage.builder()
169+
.put(BigQueryJdbcOpenTelemetry.CONNECTION_ID_BAGGAGE_KEY, this.connectionId)
170+
.build();
171+
this.otelContext = Context.current().with(baggage);
164172
try (BigQueryJdbcMdc.MdcCloseable mdc = BigQueryJdbcMdc.registerInstance(this.connectionId)) {
165173
LOG.finest("++enter++");
166174

@@ -1059,9 +1067,11 @@ private BigQuery getBigQueryConnection() {
10591067
if (this.httpTransportOptions != null) {
10601068
bigQueryOptions.setTransportOptions(this.httpTransportOptions);
10611069
}
1062-
if (Boolean.TRUE.equals(this.enableGcpTraceExporter) || this.customOpenTelemetry != null) {
1063-
this.tracer = BigQueryJdbcOpenTelemetry.getTracer(this.openTelemetry);
1064-
bigQueryOptions.setOpenTelemetryTracer(this.tracer);
1070+
if (this.enableGcpTraceExporter || this.customOpenTelemetry != null) {
1071+
Tracer sdkTracer = this.openTelemetry.getTracer(BigQueryJdbcOpenTelemetry.BIGQUERY_NAMESPACE);
1072+
bigQueryOptions.setOpenTelemetryTracer(sdkTracer);
1073+
this.tracer =
1074+
this.openTelemetry.getTracer(BigQueryJdbcOpenTelemetry.INSTRUMENTATION_SCOPE_NAME);
10651075
}
10661076

10671077
BigQueryOptions options = bigQueryOptions.setHeaderProvider(this.headerProvider).build();
@@ -1112,7 +1122,7 @@ private BigQueryReadClient getBigQueryReadClientConnection() throws IOException
11121122

11131123
bigQueryReadSettings.setTransportChannelProvider(activeProvider);
11141124

1115-
if (Boolean.TRUE.equals(this.enableGcpTraceExporter) || this.customOpenTelemetry != null) {
1125+
if (this.enableGcpTraceExporter || this.customOpenTelemetry != null) {
11161126
bigQueryReadSettings.setOpenTelemetryTracerProvider(this.openTelemetry.getTracerProvider());
11171127
}
11181128

@@ -1221,6 +1231,14 @@ public Tracer getTracer() {
12211231
return this.tracer;
12221232
}
12231233

1234+
public Context getOtelContext() {
1235+
return this.otelContext;
1236+
}
1237+
1238+
public String getPartnerToken() {
1239+
return this.partnerToken;
1240+
}
1241+
12241242
public boolean isReadOnlyTokenUsed() {
12251243
return this.isReadOnlyTokenUsed;
12261244
}

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

Lines changed: 24 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -45,8 +45,6 @@
4545
import com.google.common.annotations.VisibleForTesting;
4646
import io.opentelemetry.api.trace.Span;
4747
import io.opentelemetry.api.trace.SpanContext;
48-
import io.opentelemetry.api.trace.StatusCode;
49-
import io.opentelemetry.api.trace.Tracer;
5048
import io.opentelemetry.context.Context;
5149
import io.opentelemetry.context.Scope;
5250
import java.io.BufferedReader;
@@ -864,7 +862,8 @@ public ResultSet getProcedures(
864862
procedureNamePattern,
865863
procedureNameRegex,
866864
LOG);
867-
Future<List<Routine>> apiFuture = apiExecutor.submit(apiCallable);
865+
Future<List<Routine>> apiFuture =
866+
apiExecutor.submit(Context.current().wrap(apiCallable));
868867
apiFutures.add(apiFuture);
869868
}
870869
LOG.fine("Finished submitting " + apiFutures.size() + " findMatchingRoutines tasks.");
@@ -888,9 +887,13 @@ public ResultSet getProcedures(
888887
final Routine finalRoutine = routine;
889888
Future<?> processFuture =
890889
routineProcessorExecutor.submit(
891-
() ->
892-
processProcedureInfo(
893-
finalRoutine, collectedResults, localResultSchemaFields));
890+
Context.current()
891+
.wrap(
892+
() ->
893+
processProcedureInfo(
894+
finalRoutine,
895+
collectedResults,
896+
localResultSchemaFields)));
894897
processingTaskFutures.add(processFuture);
895898
} else {
896899
LOG.finer("Skipping non-procedure routine: " + routine.getRoutineId());
@@ -1280,7 +1283,7 @@ List<RoutineId> listMatchingProcedureIdsFromDatasets(
12801283
procedureNamePattern,
12811284
procedureNameRegex,
12821285
logger);
1283-
listRoutineFutures.add(listRoutinesExecutor.submit(listCallable));
1286+
listRoutineFutures.add(listRoutinesExecutor.submit(Context.current().wrap(listCallable)));
12841287
}
12851288
logger.fine(
12861289
"Submitted "
@@ -1357,7 +1360,7 @@ List<Routine> fetchFullRoutineDetailsForIds(
13571360
return null;
13581361
}
13591362
};
1360-
getRoutineFutures.add(getRoutineDetailsExecutor.submit(getCallable));
1363+
getRoutineFutures.add(getRoutineDetailsExecutor.submit(Context.current().wrap(getCallable)));
13611364
}
13621365
logger.fine("Submitted " + getRoutineFutures.size() + " getRoutine detail tasks.");
13631366

@@ -1407,9 +1410,14 @@ void submitProcedureArgumentProcessingJobs(
14071410
final Routine finalFullRoutine = fullRoutine;
14081411
Future<?> processFuture =
14091412
processArgsExecutor.submit(
1410-
() ->
1411-
processProcedureArguments(
1412-
finalFullRoutine, columnNameRegex, collectedResults, resultSchemaFields));
1413+
Context.current()
1414+
.wrap(
1415+
() ->
1416+
processProcedureArguments(
1417+
finalFullRoutine,
1418+
columnNameRegex,
1419+
collectedResults,
1420+
resultSchemaFields)));
14131421
outArgumentProcessingFutures.add(processFuture);
14141422
} else {
14151423
logger.warning(
@@ -4080,7 +4088,8 @@ public ResultSet getFunctions(String catalog, String schemaPattern, String funct
40804088
functionNameRegex,
40814089
LOG);
40824090
};
4083-
Future<List<Routine>> apiFuture = apiExecutor.submit(apiCallable);
4091+
Future<List<Routine>> apiFuture =
4092+
apiExecutor.submit(Context.current().wrap(apiCallable));
40844093
apiFutures.add(apiFuture);
40854094
}
40864095
LOG.fine(
@@ -4515,7 +4524,7 @@ List<RoutineId> listMatchingFunctionIdsFromDatasets(
45154524
functionNamePattern,
45164525
functionNameRegex,
45174526
logger);
4518-
listRoutineFutures.add(listRoutinesExecutor.submit(listCallable));
4527+
listRoutineFutures.add(listRoutinesExecutor.submit(Context.current().wrap(listCallable)));
45194528
}
45204529
logger.fine(
45214530
"Submitted "
@@ -5443,16 +5452,7 @@ private interface TracedMetadataOperation<T> {
54435452

54445453
private <T> T withTracing(String spanName, TracedMetadataOperation<T> operation)
54455454
throws SQLException {
5446-
Tracer tracer = this.connection.getTracer();
5447-
Span span = tracer.spanBuilder(spanName).startSpan();
5448-
try (Scope scope = span.makeCurrent()) {
5449-
return operation.run();
5450-
} catch (Exception ex) {
5451-
span.recordException(ex);
5452-
span.setStatus(StatusCode.ERROR, ex.getMessage());
5453-
throw ex;
5454-
} finally {
5455-
span.end();
5456-
}
5455+
return BigQueryJdbcOpenTelemetry.withTracing(
5456+
spanName, this.connection, null, () -> operation.run());
54575457
}
54585458
}

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

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,14 +22,22 @@
2222
import com.google.cloud.logging.LoggingOptions;
2323
import com.google.common.hash.Hashing;
2424
import io.opentelemetry.api.OpenTelemetry;
25+
import io.opentelemetry.api.baggage.Baggage;
26+
import io.opentelemetry.api.trace.Span;
27+
import io.opentelemetry.api.trace.SpanKind;
28+
import io.opentelemetry.api.trace.StatusCode;
2529
import io.opentelemetry.api.trace.Tracer;
30+
import io.opentelemetry.context.Context;
31+
import io.opentelemetry.context.Scope;
2632
import io.opentelemetry.sdk.OpenTelemetrySdk;
2733
import io.opentelemetry.sdk.autoconfigure.AutoConfiguredOpenTelemetrySdk;
2834
import java.nio.charset.StandardCharsets;
35+
import java.sql.SQLException;
2936
import java.util.Collection;
3037
import java.util.HashMap;
3138
import java.util.Map;
3239
import java.util.Objects;
40+
import java.util.concurrent.Callable;
3341
import java.util.concurrent.ConcurrentHashMap;
3442
import java.util.logging.Handler;
3543
import java.util.logging.Logger;
@@ -39,6 +47,14 @@ public class BigQueryJdbcOpenTelemetry {
3947
static final String INSTRUMENTATION_SCOPE_NAME = "com.google.cloud.bigquery.jdbc";
4048
static final String BIGQUERY_NAMESPACE = "com.google.cloud.bigquery";
4149
public static final String CONNECTION_ID_BAGGAGE_KEY = "jdbc.connection_id";
50+
public static final String DB_SYSTEM_KEY = "db.system";
51+
public static final String DB_SYSTEM_VALUE = "bigquery";
52+
public static final String DB_CONNECTION_ID_KEY = "db.connection_id";
53+
public static final String DB_APPLICATION_KEY = "db.application";
54+
public static final String DEFAULT_APPLICATION_NAME = "Google-BigQuery-JDBC-Driver";
55+
public static final String DB_STATEMENT_KEY = "db.statement";
56+
public static final String DB_STATEMENT_COUNT_KEY = "db.statement.count";
57+
public static final String DB_BATCH_STATEMENTS_KEY = "db.batch.statements";
4258
private static final String OTEL_TRACES_EXPORTER = "otel.traces.exporter";
4359
private static final String OTEL_EXPORTER_OTLP_ENDPOINT = "otel.exporter.otlp.endpoint";
4460
private static final String OTEL_LOGS_EXPORTER = "otel.logs.exporter";
@@ -287,4 +303,54 @@ public static OpenTelemetry getOpenTelemetry(
287303
public static Tracer getTracer(OpenTelemetry openTelemetry) {
288304
return openTelemetry.getTracer(INSTRUMENTATION_SCOPE_NAME);
289305
}
306+
307+
public static <T> T withTracing(
308+
String spanName, BigQueryConnection connection, String sql, Callable<T> operation)
309+
throws SQLException {
310+
311+
Tracer tracer = connection.getTracer();
312+
Span span = tracer.spanBuilder(spanName).setSpanKind(SpanKind.CLIENT).startSpan();
313+
314+
span.setAttribute(DB_SYSTEM_KEY, DB_SYSTEM_VALUE);
315+
span.setAttribute(DB_CONNECTION_ID_KEY, connection.getConnectionId());
316+
317+
String appName = connection.getPartnerToken();
318+
if (appName == null || appName.isEmpty()) {
319+
appName = DEFAULT_APPLICATION_NAME;
320+
}
321+
span.setAttribute(DB_APPLICATION_KEY, appName);
322+
323+
if (sql != null) {
324+
span.setAttribute(DB_STATEMENT_KEY, sql);
325+
}
326+
327+
Baggage updatedBaggage =
328+
Baggage.fromContext(Context.current()).toBuilder()
329+
.put(CONNECTION_ID_BAGGAGE_KEY, connection.getConnectionId())
330+
.build();
331+
332+
// Create full context with new span and updated baggage
333+
Context fullContext = Context.current().with(span).with(updatedBaggage);
334+
335+
try (Scope scope = fullContext.makeCurrent()) {
336+
return operation.call();
337+
} catch (Exception ex) {
338+
span.recordException(ex);
339+
span.setStatus(StatusCode.ERROR, ex.getMessage());
340+
341+
if (ex instanceof SQLException) {
342+
throw (SQLException) ex;
343+
}
344+
if (ex instanceof RuntimeException) {
345+
throw (RuntimeException) ex;
346+
}
347+
if (ex instanceof InterruptedException) {
348+
Thread.currentThread().interrupt();
349+
throw new BigQueryJdbcRuntimeException("Operation interrupted", ex);
350+
}
351+
throw new BigQueryJdbcRuntimeException(ex);
352+
} finally {
353+
span.end();
354+
}
355+
}
290356
}

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

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,15 @@ private int getParameterCount(String query) {
9191
@Override
9292
public ResultSet executeQuery() throws SQLException {
9393
LOG.finest("++enter++");
94+
checkClosed();
95+
return BigQueryJdbcOpenTelemetry.withTracing(
96+
"BigQueryPreparedStatement.executeQuery",
97+
this.connection,
98+
this.currentQuery,
99+
() -> executeQueryImpl());
100+
}
101+
102+
private ResultSet executeQueryImpl() throws SQLException {
94103
logQueryExecutionStart(this.currentQuery);
95104
try {
96105
QueryJobConfiguration.Builder jobConfiguration = getJobConfig(this.currentQuery);
@@ -106,6 +115,15 @@ public ResultSet executeQuery() throws SQLException {
106115
@Override
107116
public long executeLargeUpdate() throws SQLException {
108117
LOG.finest("++enter++");
118+
checkClosed();
119+
return BigQueryJdbcOpenTelemetry.withTracing(
120+
"BigQueryPreparedStatement.executeLargeUpdate",
121+
this.connection,
122+
this.currentQuery,
123+
() -> executeLargeUpdateImpl());
124+
}
125+
126+
private long executeLargeUpdateImpl() throws SQLException {
109127
logQueryExecutionStart(this.currentQuery);
110128
try {
111129
QueryJobConfiguration.Builder jobConfiguration = getJobConfig(this.currentQuery);
@@ -127,6 +145,15 @@ public int executeUpdate() throws SQLException {
127145
@Override
128146
public boolean execute() throws SQLException {
129147
LOG.finest("++enter++");
148+
checkClosed();
149+
return BigQueryJdbcOpenTelemetry.withTracing(
150+
"BigQueryPreparedStatement.execute",
151+
this.connection,
152+
this.currentQuery,
153+
() -> executeImpl());
154+
}
155+
156+
private boolean executeImpl() throws SQLException {
130157
logQueryExecutionStart(this.currentQuery);
131158
try {
132159
QueryJobConfiguration.Builder jobConfiguration = getJobConfig(this.currentQuery);

0 commit comments

Comments
 (0)