Skip to content

Commit e124bfb

Browse files
ldetmerldetmer
committed
chore(bigquery): add resend attribute + integration tests
1 parent 965761a commit e124bfb

File tree

6 files changed

+348
-4
lines changed

6 files changed

+348
-4
lines changed

‎java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryRetryHelper.java‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,15 +27,21 @@
2727
import com.google.cloud.RetryHelper;
2828
import io.opentelemetry.api.trace.Span;
2929
import io.opentelemetry.api.trace.Tracer;
30+
import io.opentelemetry.context.Context;
31+
import io.opentelemetry.context.ContextKey;
3032
import io.opentelemetry.context.Scope;
3133
import java.io.IOException;
3234
import java.util.concurrent.Callable;
3335
import java.util.concurrent.ExecutionException;
36+
import java.util.concurrent.atomic.AtomicInteger;
3437
import java.util.logging.Level;
3538
import java.util.logging.Logger;
3639

3740
public class BigQueryRetryHelper extends RetryHelper {
3841

42+
public static final ContextKey<AtomicInteger> RETRY_ATTEMPT_KEY =
43+
ContextKey.named("bq_retry_attempt");
44+
3945
private static final Logger LOG = Logger.getLogger(BigQueryRetryHelper.class.getName());
4046

4147
public static <V> V runWithRetries(
@@ -54,7 +60,11 @@ public static <V> V runWithRetries(
5460
.spanBuilder("com.google.cloud.bigquery.BigQueryRetryHelper.runWithRetries")
5561
.startSpan();
5662
}
57-
try (Scope runWithRetriesScope = runWithRetries != null ? runWithRetries.makeCurrent() : null) {
63+
Context retryContext = Context.current().with(RETRY_ATTEMPT_KEY, new AtomicInteger(0));
64+
if (runWithRetries != null) {
65+
retryContext = retryContext.with(runWithRetries);
66+
}
67+
try (Scope runWithRetriesScope = retryContext != null ? retryContext.makeCurrent() : null) {
5868
// Suppressing should be ok as a workaraund. Current and only ResultRetryAlgorithm
5969
// implementation does not use response at all, so ignoring its type is ok.
6070
@SuppressWarnings("unchecked")

‎java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/telemetry/ErrorTypeUtil.java‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,16 +16,19 @@
1616
package com.google.cloud.bigquery.telemetry;
1717

1818
import com.google.api.core.BetaApi;
19+
import com.google.common.annotations.VisibleForTesting;
1920

2021
/**
2122
* Utility class for identifying exception types for telemetry tracking. TODO: this class should get
2223
* replaced with gax version when ready work tracked in
2324
* https://github.com/googleapis/google-cloud-java/issues/12105
2425
*/
2526
@BetaApi
26-
class ErrorTypeUtil {
27+
@VisibleForTesting
28+
public class ErrorTypeUtil {
2729

28-
enum ErrorType {
30+
@VisibleForTesting
31+
public enum ErrorType {
2932
CLIENT_TIMEOUT,
3033
CLIENT_CONNECTION_ERROR,
3134
CLIENT_REQUEST_ERROR,

‎java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/telemetry/HttpTracingRequestInitializer.java‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,13 +19,15 @@
1919
import com.google.api.client.http.*;
2020
import com.google.api.core.BetaApi;
2121
import com.google.api.core.InternalApi;
22+
import com.google.cloud.bigquery.BigQueryRetryHelper;
2223
import com.google.common.annotations.VisibleForTesting;
2324
import io.opentelemetry.api.common.AttributeKey;
2425
import io.opentelemetry.api.trace.Span;
2526
import io.opentelemetry.api.trace.Tracer;
2627
import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
2728
import io.opentelemetry.context.Context;
2829
import java.io.IOException;
30+
import java.util.concurrent.atomic.AtomicInteger;
2931

3032
/**
3133
* HttpRequestInitializer that wraps a delegate initializer, intercepts all HTTP requests, adds
@@ -50,7 +52,7 @@ public class HttpTracingRequestInitializer implements HttpRequestInitializer {
5052
public static final AttributeKey<Long> HTTP_RESPONSE_BODY_SIZE =
5153
AttributeKey.longKey("http.response.body.size");
5254

53-
@VisibleForTesting static final String HTTP_RPC_SYSTEM_NAME = "http";
55+
@VisibleForTesting public static final String HTTP_RPC_SYSTEM_NAME = "http";
5456

5557
private static final java.util.Set<String> REDACTED_QUERY_PARAMETERS =
5658
com.google.common.collect.ImmutableSet.of(
@@ -84,6 +86,14 @@ public void initialize(HttpRequest request) throws IOException {
8486

8587
addInitialHttpAttributesToSpan(span, request);
8688

89+
AtomicInteger attemptTracker = Context.current().get(BigQueryRetryHelper.RETRY_ATTEMPT_KEY);
90+
if (attemptTracker != null) {
91+
int attempt = attemptTracker.getAndIncrement();
92+
if (attempt > 0) {
93+
span.setAttribute(HTTP_REQUEST_RESEND_COUNT, (long) attempt);
94+
}
95+
}
96+
8797
HttpResponseInterceptor originalInterceptor = request.getResponseInterceptor();
8898
request.setResponseInterceptor(
8999
response -> {

‎java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITOpenTelemetryTest.java‎

Lines changed: 286 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,286 @@
1+
/*
2+
* Copyright 2026 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.google.cloud.bigquery.it;
17+
18+
import static com.google.cloud.bigquery.telemetry.ErrorTypeUtil.ErrorType;
19+
import static org.junit.jupiter.api.Assertions.assertEquals;
20+
import static org.junit.jupiter.api.Assertions.assertFalse;
21+
import static org.junit.jupiter.api.Assertions.assertNotNull;
22+
import static org.junit.jupiter.api.Assertions.assertTrue;
23+
import static org.junit.jupiter.api.Assertions.fail;
24+
25+
import com.google.api.gax.retrying.RetrySettings;
26+
import com.google.cloud.bigquery.BigQuery;
27+
import com.google.cloud.bigquery.BigQueryException;
28+
import com.google.cloud.bigquery.telemetry.BigQueryTelemetryTracer;
29+
import com.google.cloud.bigquery.telemetry.HttpTracingRequestInitializer;
30+
import com.google.cloud.bigquery.testing.RemoteBigQueryHelper;
31+
import io.opentelemetry.api.OpenTelemetry;
32+
import io.opentelemetry.api.common.AttributeKey;
33+
import io.opentelemetry.api.trace.Tracer;
34+
import io.opentelemetry.sdk.OpenTelemetrySdk;
35+
import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter;
36+
import io.opentelemetry.sdk.trace.SdkTracerProvider;
37+
import io.opentelemetry.sdk.trace.data.SpanData;
38+
import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor;
39+
import java.io.IOException;
40+
import java.util.List;
41+
import java.util.Map;
42+
import org.junit.jupiter.api.BeforeAll;
43+
import org.junit.jupiter.api.BeforeEach;
44+
import org.junit.jupiter.api.Test;
45+
46+
public class ITOpenTelemetryTest {
47+
48+
private static RemoteBigQueryHelper bigqueryHelper;
49+
private InMemorySpanExporter memoryExporter;
50+
private Tracer tracer;
51+
52+
@BeforeAll
53+
public static void setUpClass() throws IOException {
54+
System.setProperty("com.google.cloud.bigquery.http.tracing.dev.enabled", "true");
55+
bigqueryHelper = RemoteBigQueryHelper.create();
56+
}
57+
58+
@BeforeEach
59+
public void setUp() {
60+
memoryExporter = InMemorySpanExporter.create();
61+
SdkTracerProvider tracerProvider =
62+
SdkTracerProvider.builder()
63+
.addSpanProcessor(SimpleSpanProcessor.create(memoryExporter))
64+
.build();
65+
OpenTelemetry openTelemetry =
66+
OpenTelemetrySdk.builder().setTracerProvider(tracerProvider).build();
67+
tracer = openTelemetry.getTracer("it-otel-test");
68+
}
69+
70+
@Test
71+
public void testListDatasetsTraced() {
72+
BigQuery bq =
73+
bigqueryHelper.getOptions().toBuilder()
74+
.setEnableOpenTelemetryTracing(true)
75+
.setOpenTelemetryTracer(tracer)
76+
.build()
77+
.getService();
78+
79+
bq.listDatasets();
80+
81+
List<SpanData> spans = memoryExporter.getFinishedSpanItems();
82+
assertNotNull(spans);
83+
assertFalse(spans.isEmpty(), "Expected at least one span collected");
84+
85+
boolean foundRpcSpan = false;
86+
for (SpanData span : spans) {
87+
if (span.getName().equals("com.google.cloud.bigquery.BigQueryRpc.listDatasets")) {
88+
foundRpcSpan = true;
89+
Map<AttributeKey<?>, Object> attrs = span.getAttributes().asMap();
90+
checkGeneralAttributes(attrs);
91+
assertEquals("GET", attrs.get(HttpTracingRequestInitializer.HTTP_REQUEST_METHOD));
92+
assertEquals("DatasetService", attrs.get(AttributeKey.stringKey("bq.rpc.service")));
93+
assertEquals("ListDatasets", attrs.get(AttributeKey.stringKey("bq.rpc.method")));
94+
assertEquals("bigquery.googleapis.com", attrs.get(BigQueryTelemetryTracer.SERVER_ADDRESS));
95+
assertEquals(200L, attrs.get(HttpTracingRequestInitializer.HTTP_RESPONSE_STATUS_CODE));
96+
assertEquals("bigquery.googleapis.com", attrs.get(BigQueryTelemetryTracer.URL_DOMAIN));
97+
assertEquals(
98+
"https://bigquery.googleapis.com/bigquery/v2/projects/gcloud-devel/datasets?prettyPrint=false",
99+
attrs.get(HttpTracingRequestInitializer.URL_FULL));
100+
assertEquals(
101+
"//bigquery.googleapis.com/projects/gcloud-devel/datasets",
102+
attrs.get(BigQueryTelemetryTracer.GCP_RESOURCE_DESTINATION_ID));
103+
assertEquals(
104+
"projects/{+projectId}/datasets", attrs.get(BigQueryTelemetryTracer.URL_TEMPLATE));
105+
}
106+
}
107+
assertTrue(foundRpcSpan, "Expected to find BigQueryRpc.listDatasets span");
108+
}
109+
110+
@Test
111+
public void testGetDatasetNotFoundTraced() {
112+
BigQuery bq =
113+
bigqueryHelper.getOptions().toBuilder()
114+
.setEnableOpenTelemetryTracing(true)
115+
.setOpenTelemetryTracer(tracer)
116+
.build()
117+
.getService();
118+
119+
bq.getDataset("non_existent_dataset");
120+
121+
List<SpanData> spans = memoryExporter.getFinishedSpanItems();
122+
assertNotNull(spans);
123+
assertFalse(spans.isEmpty(), "Expected at least one span collected");
124+
125+
boolean foundRpcSpan = false;
126+
for (SpanData span : spans) {
127+
if (span.getName().equals("com.google.cloud.bigquery.BigQueryRpc.getDataset")) {
128+
foundRpcSpan = true;
129+
Map<AttributeKey<?>, Object> attrs = span.getAttributes().asMap();
130+
checkGeneralAttributes(attrs);
131+
assertEquals("GET", attrs.get(HttpTracingRequestInitializer.HTTP_REQUEST_METHOD));
132+
assertEquals("DatasetService", attrs.get(AttributeKey.stringKey("bq.rpc.service")));
133+
assertEquals("GetDataset", attrs.get(AttributeKey.stringKey("bq.rpc.method")));
134+
assertEquals(404L, attrs.get(HttpTracingRequestInitializer.HTTP_RESPONSE_STATUS_CODE));
135+
assertEquals(
136+
"projects/{+projectId}/datasets/{+datasetId}",
137+
attrs.get(BigQueryTelemetryTracer.URL_TEMPLATE));
138+
assertEquals(
139+
"https://bigquery.googleapis.com/bigquery/v2/projects/gcloud-devel/datasets/non_existent_dataset?prettyPrint=false",
140+
attrs.get(HttpTracingRequestInitializer.URL_FULL));
141+
assertEquals("bigquery.googleapis.com", attrs.get(BigQueryTelemetryTracer.SERVER_ADDRESS));
142+
assertEquals("bigquery.googleapis.com", attrs.get(BigQueryTelemetryTracer.URL_DOMAIN));
143+
assertEquals(
144+
"//bigquery.googleapis.com/projects/gcloud-devel/datasets/non_existent_dataset",
145+
attrs.get(BigQueryTelemetryTracer.GCP_RESOURCE_DESTINATION_ID));
146+
147+
// Error attributes
148+
assertEquals("notFound", attrs.get(BigQueryTelemetryTracer.ERROR_TYPE));
149+
assertEquals(
150+
"Not found: Dataset gcloud-devel:non_existent_dataset",
151+
attrs.get(BigQueryTelemetryTracer.STATUS_MESSAGE));
152+
}
153+
}
154+
assertTrue(foundRpcSpan, "Expected to find BigQueryRpc.getDataset span");
155+
}
156+
157+
@Test
158+
public void testConnectionErrorRetriesTraced() {
159+
// Pass invalid host to force connection error and retries
160+
BigQuery bq =
161+
bigqueryHelper.getOptions().toBuilder()
162+
.setRetrySettings(RetrySettings.newBuilder().setMaxAttempts(5).build())
163+
.setEnableOpenTelemetryTracing(true)
164+
.setOpenTelemetryTracer(tracer)
165+
.setHost("https://invalid-host-name-12345.com:8080")
166+
.build()
167+
.getService();
168+
169+
try {
170+
bq.listDatasets();
171+
fail("Expected BigQueryException due to invalid host");
172+
} catch (BigQueryException e) {
173+
// Expected
174+
}
175+
176+
List<SpanData> spans = memoryExporter.getFinishedSpanItems();
177+
assertNotNull(spans);
178+
assertFalse(spans.isEmpty(), "Expected at least one span collected");
179+
180+
int rpcSpanCount = 0;
181+
for (SpanData span : spans) {
182+
if (span.getName().equals("com.google.cloud.bigquery.BigQueryRpc.listDatasets")) {
183+
rpcSpanCount++;
184+
Map<AttributeKey<?>, Object> attrs = span.getAttributes().asMap();
185+
checkGeneralAttributes(attrs);
186+
assertEquals(
187+
"https://invalid-host-name-12345.com:8080/bigquery/v2/projects/gcloud-devel/datasets?prettyPrint=false",
188+
(String) attrs.get(HttpTracingRequestInitializer.URL_FULL));
189+
assertEquals(
190+
"invalid-host-name-12345.com", attrs.get(BigQueryTelemetryTracer.SERVER_ADDRESS));
191+
assertEquals(8080L, attrs.get(BigQueryTelemetryTracer.SERVER_PORT));
192+
assertEquals("invalid-host-name-12345.com", attrs.get(BigQueryTelemetryTracer.URL_DOMAIN));
193+
assertEquals(
194+
"projects/{+projectId}/datasets", attrs.get(BigQueryTelemetryTracer.URL_TEMPLATE));
195+
assertEquals(
196+
"//bigquery.googleapis.com/projects/gcloud-devel/datasets",
197+
attrs.get(BigQueryTelemetryTracer.GCP_RESOURCE_DESTINATION_ID));
198+
checkRetryAttribute(span, rpcSpanCount);
199+
200+
// Error attributes
201+
assertEquals(
202+
"java.net.UnknownHostException", attrs.get(BigQueryTelemetryTracer.EXCEPTION_TYPE));
203+
assertEquals(
204+
ErrorType.CLIENT_CONNECTION_ERROR.toString(),
205+
attrs.get(BigQueryTelemetryTracer.ERROR_TYPE));
206+
assertEquals(
207+
"UnknownHostException: invalid-host-name-12345.com",
208+
attrs.get(BigQueryTelemetryTracer.STATUS_MESSAGE));
209+
}
210+
}
211+
assertEquals(5, rpcSpanCount, "Expected 5 attempts total");
212+
}
213+
214+
@Test
215+
public void testSimultaneousCallsDoNotAffectResendCountForEachother() {
216+
BigQuery bq =
217+
bigqueryHelper.getOptions().toBuilder()
218+
.setRetrySettings(RetrySettings.newBuilder().setMaxAttempts(5).build())
219+
.setEnableOpenTelemetryTracing(true)
220+
.setOpenTelemetryTracer(tracer)
221+
.setHost("https://invalid-host-name-123456.com")
222+
.build()
223+
.getService();
224+
225+
try {
226+
bq.listDatasets();
227+
fail("Expected BigQueryException due to invalid host");
228+
} catch (BigQueryException e) {
229+
// Expected
230+
}
231+
try {
232+
bq.cancel("test-job-id");
233+
fail("Expected BigQueryException due to invalid host");
234+
} catch (BigQueryException e) {
235+
// Expected
236+
}
237+
238+
List<SpanData> spans = memoryExporter.getFinishedSpanItems();
239+
assertNotNull(spans);
240+
assertFalse(spans.isEmpty(), "Expected at least one span collected");
241+
242+
int listDataSpanCount = 0;
243+
int cancelJobSpanCount = 0;
244+
for (SpanData span : spans) {
245+
if (span.getName().equals("com.google.cloud.bigquery.BigQueryRpc.listDatasets")) {
246+
listDataSpanCount++;
247+
checkRetryAttribute(span, listDataSpanCount);
248+
} else if (span.getName().equals("com.google.cloud.bigquery.BigQueryRpc.cancelJob")) {
249+
cancelJobSpanCount++;
250+
checkRetryAttribute(span, cancelJobSpanCount);
251+
}
252+
}
253+
assertEquals(5, listDataSpanCount, "Expected 5 attempts total for listDatasets call");
254+
assertEquals(5, cancelJobSpanCount, "Expected 5 attempts total for cancelJob call");
255+
}
256+
257+
private static void checkRetryAttribute(SpanData span, int listDataSpanCount) {
258+
Map<AttributeKey<?>, Object> attrs = span.getAttributes().asMap();
259+
Long resendCount = (Long) attrs.get(HttpTracingRequestInitializer.HTTP_REQUEST_RESEND_COUNT);
260+
if (listDataSpanCount == 1) {
261+
assertTrue(resendCount == null || resendCount == 0);
262+
} else {
263+
assertNotNull(resendCount, "Expected resend count for retry attempt " + listDataSpanCount);
264+
assertEquals((long) (listDataSpanCount - 1), resendCount.longValue());
265+
}
266+
}
267+
268+
private void checkGeneralAttributes(Map<AttributeKey<?>, Object> attrs) {
269+
assertEquals(
270+
HttpTracingRequestInitializer.HTTP_RPC_SYSTEM_NAME,
271+
attrs.get(BigQueryTelemetryTracer.RPC_SYSTEM_NAME));
272+
assertEquals(
273+
BigQueryTelemetryTracer.BQ_GCP_CLIENT_SERVICE,
274+
attrs.get(BigQueryTelemetryTracer.GCP_CLIENT_SERVICE));
275+
assertEquals(
276+
BigQueryTelemetryTracer.BQ_GCP_CLIENT_REPO,
277+
attrs.get(BigQueryTelemetryTracer.GCP_CLIENT_REPO));
278+
assertEquals(
279+
BigQueryTelemetryTracer.BQ_GCP_CLIENT_LANGUAGE,
280+
attrs.get(BigQueryTelemetryTracer.GCP_CLIENT_LANGUAGE));
281+
assertEquals(
282+
BigQueryTelemetryTracer.BQ_GCP_CLIENT_ARTIFACT,
283+
attrs.get(BigQueryTelemetryTracer.GCP_CLIENT_ARTIFACT));
284+
assertNotNull(attrs.get(BigQueryTelemetryTracer.GCP_CLIENT_VERSION));
285+
}
286+
}

0 commit comments

Comments
 (0)