Skip to content

Commit 57a0207

Browse files
feat(bqjdbc): extend OpenTelemetry instrumentation for metadata and pagination (#12918)
b/491245568 ### Key Changes #### Core Instrumentation Logic * **Database Metadata Tracing**: Added OTel spans to key methods in `BigQueryDatabaseMetaData.java` (`getCatalogs`, `getSchemas`, `getTables`, `getColumns`) to capture underlying API calls. * **Pagination Span Links**: Captured the parent span context at the start of `fetchNextPages` in `BigQueryStatement.java` and linked background pagination spans back to it, avoiding timeline anomalies. * **Cross-Thread Context Propagation**: Stored the `SpanContext` in `BigQueryBaseResultSet.java` at creation time and made it current during `next()` in `BigQueryJsonResultSet.java` and `BigQueryArrowResultSet.java` to survive thread hops. * **Tracer Reuse**: Extracted `getSafeTracer` to `BigQueryJdbcOpenTelemetry.java` as a static utility to ensure consistent fallback behavior across the driver. * **Lambda Extraction**: Extracted the large lambda function in `populateArrowBufferedQueue` in `BigQueryStatement.java` to its own private method `processArrowStream` to improve readability and maintainability.
1 parent c7131f9 commit 57a0207

6 files changed

Lines changed: 471 additions & 362 deletions

File tree

java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java

Lines changed: 23 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
import com.google.cloud.bigquery.exception.BigQueryJdbcRuntimeException;
2828
import com.google.cloud.bigquery.storage.v1.ArrowRecordBatch;
2929
import com.google.cloud.bigquery.storage.v1.ArrowSchema;
30+
import io.opentelemetry.context.Scope;
3031
import java.io.IOException;
3132
import java.math.BigDecimal;
3233
import java.sql.Date;
@@ -236,29 +237,31 @@ public boolean next() throws SQLException {
236237
|| this.currentBatchRowIndex == (this.vectorSchemaRoot.getRowCount() - 1)) {
237238
/* Start of iteration or we have exhausted the current batch */
238239
// Advance the cursor. Potentially blocking operation.
239-
BigQueryArrowBatchWrapper batchWrapper = this.buffer.take();
240-
if (batchWrapper.getException() != null) {
241-
throw new BigQueryJdbcRuntimeException(batchWrapper.getException());
242-
}
243-
if (batchWrapper.isLast()) {
244-
/* Marks the end of the records */
245-
if (this.vectorSchemaRoot != null) {
246-
// IMP: To avoid memory leak: clear vectorSchemaRoot as it still holds
247-
// the last batch
248-
this.vectorSchemaRoot.clear();
240+
try (Scope scope = makeOriginalContextCurrent()) {
241+
BigQueryArrowBatchWrapper batchWrapper = this.buffer.take();
242+
if (batchWrapper.getException() != null) {
243+
throw new BigQueryJdbcRuntimeException(batchWrapper.getException());
244+
}
245+
if (batchWrapper.isLast()) {
246+
/* Marks the end of the records */
247+
if (this.vectorSchemaRoot != null) {
248+
// IMP: To avoid memory leak: clear vectorSchemaRoot as it still holds
249+
// the last batch
250+
this.vectorSchemaRoot.clear();
251+
}
252+
this.hasReachedEnd = true;
253+
this.rowCount++;
254+
return false;
249255
}
250-
this.hasReachedEnd = true;
256+
// Valid batch, process it
257+
ArrowRecordBatch arrowBatch = batchWrapper.getCurrentArrowBatch();
258+
// Populates vectorSchemaRoot
259+
this.arrowDeserializer.deserializeArrowBatch(arrowBatch);
260+
// Pointing to the first row in this fresh batch
261+
this.currentBatchRowIndex = 0;
251262
this.rowCount++;
252-
return false;
263+
return true;
253264
}
254-
// Valid batch, process it
255-
ArrowRecordBatch arrowBatch = batchWrapper.getCurrentArrowBatch();
256-
// Populates vectorSchemaRoot
257-
this.arrowDeserializer.deserializeArrowBatch(arrowBatch);
258-
// Pointing to the first row in this fresh batch
259-
this.currentBatchRowIndex = 0;
260-
this.rowCount++;
261-
return true;
262265
}
263266
// There are rows left in the current batch.
264267
else if (this.currentBatchRowIndex < this.vectorSchemaRoot.getRowCount()) {

java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryBaseResultSet.java

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,10 @@
2727
import com.google.cloud.bigquery.exception.BigQueryConversionException;
2828
import com.google.cloud.bigquery.exception.BigQueryJdbcCoercionException;
2929
import com.google.cloud.bigquery.exception.BigQueryJdbcCoercionNotFoundException;
30+
import io.opentelemetry.api.trace.Span;
31+
import io.opentelemetry.api.trace.SpanContext;
32+
import io.opentelemetry.context.Context;
33+
import io.opentelemetry.context.Scope;
3034
import java.io.InputStream;
3135
import java.io.Reader;
3236
import java.io.StringReader;
@@ -58,6 +62,7 @@ public abstract class BigQueryBaseResultSet extends BigQueryNoOpsResultSet
5862
protected boolean isClosed = false;
5963
protected boolean wasNull = false;
6064
protected final BigQueryTypeCoercer bigQueryTypeCoercer = BigQueryTypeCoercionUtility.INSTANCE;
65+
protected final SpanContext originalSpanContext;
6166

6267
protected BigQueryBaseResultSet(
6368
BigQuery bigQuery, BigQueryStatement statement, Schema schema, boolean isNested) {
@@ -66,6 +71,11 @@ protected BigQueryBaseResultSet(
6671
this.schema = schema;
6772
this.schemaFieldList = schema != null ? schema.getFields() : null;
6873
this.isNested = isNested;
74+
this.originalSpanContext = Span.current().getSpanContext();
75+
}
76+
77+
protected Scope makeOriginalContextCurrent() {
78+
return Context.current().with(Span.wrap(this.originalSpanContext)).makeCurrent();
6979
}
7080

7181
public QueryStatistics getQueryStatistics() {

0 commit comments

Comments
 (0)