Skip to content

Commit 119ad10

Browse files
keshavdandevakeshavdandeva
committed
test(jdbc): add unit tests for OTel tracing
1 parent ef20f9e commit 119ad10

5 files changed

Lines changed: 376 additions & 39 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: 93 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,57 @@ 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 catalog,
1753+
String schemaPattern,
1754+
String tableNamePattern,
1755+
String[] types,
1756+
Schema resultSchema,
1757+
BlockingQueue<BigQueryFieldValueListWrapper> queue)
1758+
throws SQLException {
1759+
1760+
Tuple<String, String> effectiveIdentifiers =
1761+
determineEffectiveCatalogAndSchema(catalog, schemaPattern);
1762+
String effectiveCatalog = effectiveIdentifiers.x();
1763+
String effectiveSchemaPattern = effectiveIdentifiers.y();
1764+
17341765
final Pattern schemaRegex = compileSqlLikePattern(effectiveSchemaPattern);
17351766
final Pattern tableNameRegex = compileSqlLikePattern(tableNamePattern);
17361767
final Set<String> requestedTypes =
17371768
(types == null || types.length == 0) ? null : new HashSet<>(Arrays.asList(types));
17381769

1739-
final Schema resultSchema = defineGetTablesSchema();
17401770
final FieldList resultSchemaFields = resultSchema.getFields();
1741-
1742-
final BlockingQueue<BigQueryFieldValueListWrapper> queue =
1743-
new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY);
17441771
final List<FieldValueList> collectedResults = Collections.synchronizedList(new ArrayList<>());
17451772
final String catalogParam = effectiveCatalog;
17461773
final String schemaParam = effectiveSchemaPattern;
1747-
1748-
Tracer tracer = this.connection.getTracer();
17491774
SpanContext parentSpanContext = Span.current().getSpanContext();
17501775
Runnable tableFetcher =
17511776
() -> {
17521777
Span backgroundSpan =
1753-
tracer
1778+
this.connection
1779+
.getTracer()
17541780
.spanBuilder("BigQueryDatabaseMetaData.getTables.background")
17551781
.setNoParent()
17561782
.addLink(parentSpanContext)
17571783
.startSpan();
17581784

1759-
try (Scope backgroundScope = backgroundSpan.makeCurrent()) {
1785+
try (Scope scope = backgroundSpan.makeCurrent()) {
17601786
ExecutorService apiExecutor = null;
17611787
ExecutorService tableProcessorExecutor = null;
17621788
final FieldList localResultSchemaFields = resultSchemaFields;
@@ -1898,12 +1924,8 @@ private ResultSet getTablesImpl(
18981924

18991925
Runnable wrappedTableFetcher = Context.current().wrap(tableFetcher);
19001926
Thread fetcherThread = new Thread(wrappedTableFetcher, "getTables-fetcher-" + effectiveCatalog);
1901-
BigQueryJsonResultSet resultSet =
1902-
BigQueryJsonResultSet.of(resultSchema, -1, queue, null, new Thread[] {fetcherThread});
1903-
19041927
fetcherThread.start();
1905-
LOG.info("Started background thread for getTables");
1906-
return resultSet;
1928+
return fetcherThread;
19071929
}
19081930

19091931
Schema defineGetTablesSchema() {
@@ -2127,24 +2149,51 @@ private ResultSet getColumnsImpl(
21272149
+ " columnNamePattern: %s",
21282150
effectiveCatalog, effectiveSchemaPattern, tableNamePattern, columnNamePattern);
21292151

2152+
final Schema resultSchema = defineGetColumnsSchema();
2153+
final BlockingQueue<BigQueryFieldValueListWrapper> queue =
2154+
new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY);
2155+
2156+
Thread fetcherThread =
2157+
runGetColumnsTaskAsync(
2158+
catalog, schemaPattern, tableNamePattern, columnNamePattern, resultSchema, 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 catalog,
2170+
String schemaPattern,
2171+
String tableNamePattern,
2172+
String columnNamePattern,
2173+
Schema resultSchema,
2174+
BlockingQueue<BigQueryFieldValueListWrapper> queue)
2175+
throws SQLException {
2176+
2177+
Tuple<String, String> effectiveIdentifiers =
2178+
determineEffectiveCatalogAndSchema(catalog, schemaPattern);
2179+
String effectiveCatalog = effectiveIdentifiers.x();
2180+
String effectiveSchemaPattern = effectiveIdentifiers.y();
2181+
21302182
Pattern schemaRegex = compileSqlLikePattern(effectiveSchemaPattern);
21312183
Pattern tableNameRegex = compileSqlLikePattern(tableNamePattern);
21322184
Pattern columnNameRegex = compileSqlLikePattern(columnNamePattern);
21332185

2134-
final Schema resultSchema = defineGetColumnsSchema();
21352186
final FieldList resultSchemaFields = resultSchema.getFields();
2136-
final BlockingQueue<BigQueryFieldValueListWrapper> queue =
2137-
new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY);
21382187
final List<FieldValueList> collectedResults = Collections.synchronizedList(new ArrayList<>());
21392188
final String catalogParam = effectiveCatalog;
21402189
final String schemaParam = effectiveSchemaPattern;
21412190

2142-
Tracer tracer = this.connection.getTracer();
21432191
SpanContext parentSpanContext = Span.current().getSpanContext();
21442192
Runnable columnFetcher =
21452193
() -> {
21462194
Span backgroundSpan =
2147-
tracer
2195+
this.connection
2196+
.getTracer()
21482197
.spanBuilder("BigQueryDatabaseMetaData.getColumns.background")
21492198
.setNoParent()
21502199
.addLink(parentSpanContext)
@@ -2252,12 +2301,8 @@ private ResultSet getColumnsImpl(
22522301
Runnable wrappedColumnFetcher = Context.current().wrap(columnFetcher);
22532302
Thread fetcherThread =
22542303
new Thread(wrappedColumnFetcher, "getColumns-fetcher-" + effectiveCatalog);
2255-
BigQueryJsonResultSet resultSet =
2256-
BigQueryJsonResultSet.of(resultSchema, -1, queue, null, new Thread[] {fetcherThread});
2257-
22582304
fetcherThread.start();
2259-
LOG.info("Started background thread for getColumns");
2260-
return resultSet;
2305+
return fetcherThread;
22612306
}
22622307

22632308
private void processTableColumns(
@@ -2324,7 +2369,7 @@ private void processTableColumns(
23242369
}
23252370
}
23262371

2327-
private Schema defineGetColumnsSchema() {
2372+
Schema defineGetColumnsSchema() {
23282373
List<Field> fields = new ArrayList<>(24);
23292374
fields.add(
23302375
Field.newBuilder("TABLE_CAT", StandardSQLTypeName.STRING)
@@ -3690,27 +3735,44 @@ private ResultSet getSchemasImpl(String catalog, String schemaPattern) throws SQ
36903735

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

3693-
final Pattern schemaRegex = compileSqlLikePattern(schemaPattern);
36943738
final Schema resultSchema = defineGetSchemasSchema();
3695-
final FieldList resultSchemaFields = resultSchema.getFields();
3696-
36973739
final BlockingQueue<BigQueryFieldValueListWrapper> queue =
36983740
new LinkedBlockingQueue<>(DEFAULT_QUEUE_CAPACITY);
3741+
3742+
Thread fetcherThread = runGetSchemasTaskAsync(catalog, schemaPattern, resultSchema, queue);
3743+
3744+
BigQueryJsonResultSet resultSet =
3745+
BigQueryJsonResultSet.of(resultSchema, -1, queue, null, new Thread[] {fetcherThread});
3746+
3747+
LOG.info("Started background thread for getSchemas");
3748+
return resultSet;
3749+
}
3750+
3751+
@VisibleForTesting
3752+
Thread runGetSchemasTaskAsync(
3753+
String catalog,
3754+
String schemaPattern,
3755+
Schema resultSchema,
3756+
BlockingQueue<BigQueryFieldValueListWrapper> queue)
3757+
throws SQLException {
3758+
3759+
final Pattern schemaRegex = compileSqlLikePattern(schemaPattern);
3760+
final FieldList resultSchemaFields = resultSchema.getFields();
36993761
final List<FieldValueList> collectedResults = Collections.synchronizedList(new ArrayList<>());
37003762
final String catalogParam = catalog;
37013763

3702-
Tracer tracer = this.connection.getTracer();
37033764
SpanContext parentSpanContext = Span.current().getSpanContext();
37043765
Runnable schemaFetcher =
37053766
() -> {
37063767
Span backgroundSpan =
3707-
tracer
3768+
this.connection
3769+
.getTracer()
37083770
.spanBuilder("BigQueryDatabaseMetaData.getSchemas.background")
37093771
.setNoParent()
37103772
.addLink(parentSpanContext)
37113773
.startSpan();
37123774

3713-
try (Scope backgroundScope = backgroundSpan.makeCurrent()) {
3775+
try (Scope scope = backgroundSpan.makeCurrent()) {
37143776
final FieldList localResultSchemaFields = resultSchemaFields;
37153777
List<String> projectsToScanList = new ArrayList<>();
37163778

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

37923854
Runnable wrappedFetcher = Context.current().wrap(schemaFetcher);
37933855
Thread fetcherThread = new Thread(wrappedFetcher, "getSchemas-fetcher-" + catalog);
3794-
BigQueryJsonResultSet resultSet =
3795-
BigQueryJsonResultSet.of(resultSchema, -1, queue, null, new Thread[] {fetcherThread});
3796-
37973856
fetcherThread.start();
3798-
LOG.info("Started background thread for getSchemas");
3799-
return resultSet;
3857+
return fetcherThread;
38003858
}
38013859

38023860
Schema defineGetSchemasSchema() {

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1561,7 +1561,6 @@ private void fetchNextPages(
15611561
BlockingQueue<Tuple<TableResult, Boolean>> rpcResponseQueue,
15621562
BlockingQueue<BigQueryFieldValueListWrapper> bigQueryFieldValueListWrapperBlockingQueue,
15631563
TableResult result) {
1564-
Tracer tracer = this.connection.getTracer();
15651564
SpanContext parentSpanContext = Span.current().getSpanContext();
15661565
String currentPageToken = firstPageToken;
15671566
TableResult currentResults = result;
@@ -1577,7 +1576,8 @@ private void fetchNextPages(
15771576
break;
15781577
}
15791578

1580-
SpanBuilder spanBuilder = tracer.spanBuilder("BigQueryStatement.pagination");
1579+
SpanBuilder spanBuilder =
1580+
this.connection.getTracer().spanBuilder("BigQueryStatement.pagination");
15811581
if (parentSpanContext.isValid()) {
15821582
spanBuilder.addLink(parentSpanContext);
15831583
}

0 commit comments

Comments
 (0)