Skip to content

Commit 20bb4f9

Browse files
keshavdandevakeshavdandeva
authored
test(bqjdbc): add unit tests for OTel tracing (#12943)
b/491245568 ### Changes #### Enhanced Unit Testing Infrastructure - `OpenTelemetryTestUtility`: Introduced a shared utility to simplify complex OTel assertions - `BigQueryStatementTest`: Added comprehensive parameterised tracing tests for all execution paths, including attribute validation for queries and batch operations. - `BigQueryDatabaseMetaDataTest`: Instrumented all metadata discovery methods and added corresponding unit tests to verify span generation using parameterised tests
1 parent 57a0207 commit 20bb4f9

6 files changed

Lines changed: 525 additions & 51 deletions

File tree

‎java-bigquery/google-cloud-bigquery-jdbc/pom.xml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -340,6 +340,11 @@
340340
<artifactId>junit-platform-suite-engine</artifactId>
341341
<scope>test</scope>
342342
</dependency>
343+
<dependency>
344+
<groupId>io.opentelemetry</groupId>
345+
<artifactId>opentelemetry-sdk-testing</artifactId>
346+
<scope>test</scope>
347+
</dependency>
343348
</dependencies>
344349

345350
<profiles>

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

Lines changed: 88 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@
4242
import com.google.cloud.bigquery.TableDefinition;
4343
import com.google.cloud.bigquery.TableId;
4444
import com.google.cloud.bigquery.exception.BigQueryJdbcException;
45+
import com.google.common.annotations.VisibleForTesting;
4546
import io.opentelemetry.api.trace.Span;
4647
import io.opentelemetry.api.trace.SpanContext;
4748
import io.opentelemetry.api.trace.StatusCode;
@@ -1731,32 +1732,52 @@ private ResultSet getTablesImpl(
17311732
"getTables called for catalog: %s, schemaPattern: %s, tableNamePattern: %s, types: %s",
17321733
effectiveCatalog, effectiveSchemaPattern, tableNamePattern, Arrays.toString(types));
17331734

1735+
final Schema resultSchema = defineGetTablesSchema();
1736+
final BlockingQueue<BigQueryFieldValueListWrapper> queue =
1737+
new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY);
1738+
1739+
Thread fetcherThread =
1740+
runGetTablesTaskAsync(
1741+
effectiveCatalog, effectiveSchemaPattern, tableNamePattern, types, resultSchema, queue);
1742+
1743+
BigQueryJsonResultSet resultSet =
1744+
BigQueryJsonResultSet.of(resultSchema, -1, queue, null, new Thread[] {fetcherThread});
1745+
1746+
LOG.info("Started background thread for getTables");
1747+
return resultSet;
1748+
}
1749+
1750+
@VisibleForTesting
1751+
Thread runGetTablesTaskAsync(
1752+
String effectiveCatalog,
1753+
String effectiveSchemaPattern,
1754+
String tableNamePattern,
1755+
String[] types,
1756+
Schema resultSchema,
1757+
BlockingQueue<BigQueryFieldValueListWrapper> queue)
1758+
throws SQLException {
1759+
17341760
final Pattern schemaRegex = compileSqlLikePattern(effectiveSchemaPattern);
17351761
final Pattern tableNameRegex = compileSqlLikePattern(tableNamePattern);
17361762
final Set<String> requestedTypes =
17371763
(types == null || types.length == 0) ? null : new HashSet<>(Arrays.asList(types));
17381764

1739-
final Schema resultSchema = defineGetTablesSchema();
17401765
final FieldList resultSchemaFields = resultSchema.getFields();
1741-
1742-
final BlockingQueue<BigQueryFieldValueListWrapper> queue =
1743-
new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY);
17441766
final List<FieldValueList> collectedResults = Collections.synchronizedList(new ArrayList<>());
17451767
final String catalogParam = effectiveCatalog;
17461768
final String schemaParam = effectiveSchemaPattern;
1747-
1748-
Tracer tracer = this.connection.getTracer();
17491769
SpanContext parentSpanContext = Span.current().getSpanContext();
17501770
Runnable tableFetcher =
17511771
() -> {
17521772
Span backgroundSpan =
1753-
tracer
1773+
this.connection
1774+
.getTracer()
17541775
.spanBuilder("BigQueryDatabaseMetaData.getTables.background")
17551776
.setNoParent()
17561777
.addLink(parentSpanContext)
17571778
.startSpan();
17581779

1759-
try (Scope backgroundScope = backgroundSpan.makeCurrent()) {
1780+
try (Scope scope = backgroundSpan.makeCurrent()) {
17601781
ExecutorService apiExecutor = null;
17611782
ExecutorService tableProcessorExecutor = null;
17621783
final FieldList localResultSchemaFields = resultSchemaFields;
@@ -1898,12 +1919,8 @@ private ResultSet getTablesImpl(
18981919

18991920
Runnable wrappedTableFetcher = Context.current().wrap(tableFetcher);
19001921
Thread fetcherThread = new Thread(wrappedTableFetcher, "getTables-fetcher-" + effectiveCatalog);
1901-
BigQueryJsonResultSet resultSet =
1902-
BigQueryJsonResultSet.of(resultSchema, -1, queue, null, new Thread[] {fetcherThread});
1903-
19041922
fetcherThread.start();
1905-
LOG.info("Started background thread for getTables");
1906-
return resultSet;
1923+
return fetcherThread;
19071924
}
19081925

19091926
Schema defineGetTablesSchema() {
@@ -2127,24 +2144,51 @@ private ResultSet getColumnsImpl(
21272144
+ " columnNamePattern: %s",
21282145
effectiveCatalog, effectiveSchemaPattern, tableNamePattern, columnNamePattern);
21292146

2147+
final Schema resultSchema = defineGetColumnsSchema();
2148+
final BlockingQueue<BigQueryFieldValueListWrapper> queue =
2149+
new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY);
2150+
2151+
Thread fetcherThread =
2152+
runGetColumnsTaskAsync(
2153+
effectiveCatalog,
2154+
effectiveSchemaPattern,
2155+
tableNamePattern,
2156+
columnNamePattern,
2157+
resultSchema,
2158+
queue);
2159+
2160+
BigQueryJsonResultSet resultSet =
2161+
BigQueryJsonResultSet.of(resultSchema, -1, queue, null, new Thread[] {fetcherThread});
2162+
2163+
LOG.info("Started background thread for getColumns");
2164+
return resultSet;
2165+
}
2166+
2167+
@VisibleForTesting
2168+
Thread runGetColumnsTaskAsync(
2169+
String effectiveCatalog,
2170+
String effectiveSchemaPattern,
2171+
String tableNamePattern,
2172+
String columnNamePattern,
2173+
Schema resultSchema,
2174+
BlockingQueue<BigQueryFieldValueListWrapper> queue)
2175+
throws SQLException {
2176+
21302177
Pattern schemaRegex = compileSqlLikePattern(effectiveSchemaPattern);
21312178
Pattern tableNameRegex = compileSqlLikePattern(tableNamePattern);
21322179
Pattern columnNameRegex = compileSqlLikePattern(columnNamePattern);
21332180

2134-
final Schema resultSchema = defineGetColumnsSchema();
21352181
final FieldList resultSchemaFields = resultSchema.getFields();
2136-
final BlockingQueue<BigQueryFieldValueListWrapper> queue =
2137-
new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY);
21382182
final List<FieldValueList> collectedResults = Collections.synchronizedList(new ArrayList<>());
21392183
final String catalogParam = effectiveCatalog;
21402184
final String schemaParam = effectiveSchemaPattern;
21412185

2142-
Tracer tracer = this.connection.getTracer();
21432186
SpanContext parentSpanContext = Span.current().getSpanContext();
21442187
Runnable columnFetcher =
21452188
() -> {
21462189
Span backgroundSpan =
2147-
tracer
2190+
this.connection
2191+
.getTracer()
21482192
.spanBuilder("BigQueryDatabaseMetaData.getColumns.background")
21492193
.setNoParent()
21502194
.addLink(parentSpanContext)
@@ -2252,12 +2296,8 @@ private ResultSet getColumnsImpl(
22522296
Runnable wrappedColumnFetcher = Context.current().wrap(columnFetcher);
22532297
Thread fetcherThread =
22542298
new Thread(wrappedColumnFetcher, "getColumns-fetcher-" + effectiveCatalog);
2255-
BigQueryJsonResultSet resultSet =
2256-
BigQueryJsonResultSet.of(resultSchema, -1, queue, null, new Thread[] {fetcherThread});
2257-
22582299
fetcherThread.start();
2259-
LOG.info("Started background thread for getColumns");
2260-
return resultSet;
2300+
return fetcherThread;
22612301
}
22622302

22632303
private void processTableColumns(
@@ -2324,7 +2364,7 @@ private void processTableColumns(
23242364
}
23252365
}
23262366

2327-
private Schema defineGetColumnsSchema() {
2367+
Schema defineGetColumnsSchema() {
23282368
List<Field> fields = new ArrayList<>(24);
23292369
fields.add(
23302370
Field.newBuilder("TABLE_CAT", StandardSQLTypeName.STRING)
@@ -3690,27 +3730,44 @@ private ResultSet getSchemasImpl(String catalog, String schemaPattern) throws SQ
36903730

36913731
LOG.info("getSchemas called for catalog: %s, schemaPattern: %s", catalog, schemaPattern);
36923732

3693-
final Pattern schemaRegex = compileSqlLikePattern(schemaPattern);
36943733
final Schema resultSchema = defineGetSchemasSchema();
3695-
final FieldList resultSchemaFields = resultSchema.getFields();
3696-
36973734
final BlockingQueue<BigQueryFieldValueListWrapper> queue =
36983735
new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY);
3736+
3737+
Thread fetcherThread = runGetSchemasTaskAsync(catalog, schemaPattern, resultSchema, queue);
3738+
3739+
BigQueryJsonResultSet resultSet =
3740+
BigQueryJsonResultSet.of(resultSchema, -1, queue, null, new Thread[] {fetcherThread});
3741+
3742+
LOG.info("Started background thread for getSchemas");
3743+
return resultSet;
3744+
}
3745+
3746+
@VisibleForTesting
3747+
Thread runGetSchemasTaskAsync(
3748+
String catalog,
3749+
String schemaPattern,
3750+
Schema resultSchema,
3751+
BlockingQueue<BigQueryFieldValueListWrapper> queue)
3752+
throws SQLException {
3753+
3754+
final Pattern schemaRegex = compileSqlLikePattern(schemaPattern);
3755+
final FieldList resultSchemaFields = resultSchema.getFields();
36993756
final List<FieldValueList> collectedResults = Collections.synchronizedList(new ArrayList<>());
37003757
final String catalogParam = catalog;
37013758

3702-
Tracer tracer = this.connection.getTracer();
37033759
SpanContext parentSpanContext = Span.current().getSpanContext();
37043760
Runnable schemaFetcher =
37053761
() -> {
37063762
Span backgroundSpan =
3707-
tracer
3763+
this.connection
3764+
.getTracer()
37083765
.spanBuilder("BigQueryDatabaseMetaData.getSchemas.background")
37093766
.setNoParent()
37103767
.addLink(parentSpanContext)
37113768
.startSpan();
37123769

3713-
try (Scope backgroundScope = backgroundSpan.makeCurrent()) {
3770+
try (Scope scope = backgroundSpan.makeCurrent()) {
37143771
final FieldList localResultSchemaFields = resultSchemaFields;
37153772
List<String> projectsToScanList = new ArrayList<>();
37163773

@@ -3791,12 +3848,8 @@ private ResultSet getSchemasImpl(String catalog, String schemaPattern) throws SQ
37913848

37923849
Runnable wrappedFetcher = Context.current().wrap(schemaFetcher);
37933850
Thread fetcherThread = new Thread(wrappedFetcher, "getSchemas-fetcher-" + catalog);
3794-
BigQueryJsonResultSet resultSet =
3795-
BigQueryJsonResultSet.of(resultSchema, -1, queue, null, new Thread[] {fetcherThread});
3796-
37973851
fetcherThread.start();
3798-
LOG.info("Started background thread for getSchemas");
3799-
return resultSet;
3852+
return fetcherThread;
38003853
}
38013854

38023855
Schema defineGetSchemasSchema() {

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1356,7 +1356,9 @@ public int[] executeBatch() throws SQLException {
13561356
"BigQueryStatement.executeBatch",
13571357
(span) -> {
13581358
span.setAttribute("db.statement.count", this.batchQueries.size());
1359-
span.setAttribute(AttributeKey.stringArrayKey("db.batch.statements"), this.batchQueries);
1359+
span.setAttribute(
1360+
AttributeKey.stringArrayKey("db.batch.statements"),
1361+
new ArrayList<>(this.batchQueries));
13601362

13611363
StringBuilder sb = new StringBuilder();
13621364
for (String query : this.batchQueries) {
@@ -1561,7 +1563,6 @@ private void fetchNextPages(
15611563
BlockingQueue<Tuple<TableResult, Boolean>> rpcResponseQueue,
15621564
BlockingQueue<BigQueryFieldValueListWrapper> bigQueryFieldValueListWrapperBlockingQueue,
15631565
TableResult result) {
1564-
Tracer tracer = this.connection.getTracer();
15651566
SpanContext parentSpanContext = Span.current().getSpanContext();
15661567
String currentPageToken = firstPageToken;
15671568
TableResult currentResults = result;
@@ -1571,6 +1572,7 @@ private void fetchNextPages(
15711572
}
15721573

15731574
try {
1575+
Tracer tracer = this.connection.getTracer();
15741576
while (currentPageToken != null) {
15751577
if (Thread.currentThread().isInterrupted() || queryTaskExecutor.isShutdown()) {
15761578
LOG.warning("%s Interrupted @ runNextPageTaskAsync", Thread.currentThread().getName());

0 commit comments

Comments
 (0)