Skip to content

Commit 8c4a832

Browse files
jimit-j-shahjimit-j-shah
authored
feat: RunAggregationQuery instrumentation (#1447)
* feat: RunQuery trace instrumentation * Formatting * Formatting * Refactor: s/RUNQUERY/RUN_QUERY * feat: RunAggregationQuery Trace Instrumentation * Build: retiring test assertions for OpenCensus spans - will be replacing this in hermetic integration tests for OpenTelemetry using in-memory span exports (in addition to ITE2ETraceTest.java). * Formatting * Fixing @test annotation missed after merge * Formatting
1 parent e44653c commit 8c4a832

File tree

4 files changed

+89
-24
lines changed

4 files changed

+89
-24
lines changed

‎google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecorator.java

Lines changed: 8 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@
1616
package com.google.cloud.datastore;
1717

1818
import static com.google.cloud.BaseService.EXCEPTION_HANDLER;
19-
import static com.google.cloud.datastore.TraceUtil.SPAN_NAME_RUN_AGGREGATION_QUERY;
2019

2120
import com.google.api.core.InternalApi;
2221
import com.google.api.gax.retrying.RetrySettings;
@@ -39,9 +38,6 @@
3938
import com.google.datastore.v1.RunAggregationQueryResponse;
4039
import com.google.datastore.v1.RunQueryRequest;
4140
import com.google.datastore.v1.RunQueryResponse;
42-
import io.opencensus.common.Scope;
43-
import io.opencensus.trace.Span;
44-
import io.opencensus.trace.Status;
4541
import java.util.concurrent.Callable;
4642

4743
/**
@@ -52,7 +48,7 @@
5248
public class RetryAndTraceDatastoreRpcDecorator implements DatastoreRpc {
5349

5450
private final DatastoreRpc datastoreRpc;
55-
private final TraceUtil traceUtil;
51+
private final com.google.cloud.datastore.telemetry.TraceUtil otelTraceUtil;
5652
private final RetrySettings retrySettings;
5753
private final DatastoreOptions datastoreOptions;
5854

@@ -62,9 +58,9 @@ public RetryAndTraceDatastoreRpcDecorator(
6258
RetrySettings retrySettings,
6359
DatastoreOptions datastoreOptions) {
6460
this.datastoreRpc = datastoreRpc;
65-
this.traceUtil = traceUtil;
6661
this.retrySettings = retrySettings;
6762
this.datastoreOptions = datastoreOptions;
63+
this.otelTraceUtil = datastoreOptions.getTraceUtil();
6864
}
6965

7066
@Override
@@ -106,19 +102,20 @@ public RunQueryResponse runQuery(RunQueryRequest request) {
106102
@Override
107103
public RunAggregationQueryResponse runAggregationQuery(RunAggregationQueryRequest request) {
108104
return invokeRpc(
109-
() -> datastoreRpc.runAggregationQuery(request), SPAN_NAME_RUN_AGGREGATION_QUERY);
105+
() -> datastoreRpc.runAggregationQuery(request),
106+
com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_RUN_AGGREGATION_QUERY);
110107
}
111108

112109
public <O> O invokeRpc(Callable<O> block, String startSpan) {
113-
Span span = traceUtil.startSpan(startSpan);
114-
try (Scope scope = traceUtil.getTracer().withSpan(span)) {
110+
com.google.cloud.datastore.telemetry.TraceUtil.Span span = otelTraceUtil.startSpan(startSpan);
111+
try (com.google.cloud.datastore.telemetry.TraceUtil.Scope ignored = span.makeCurrent()) {
115112
return RetryHelper.runWithRetries(
116113
block, this.retrySettings, EXCEPTION_HANDLER, this.datastoreOptions.getClock());
117114
} catch (RetryHelperException e) {
118-
span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
115+
span.end(e);
119116
throw DatastoreException.translateAndThrow(e);
120117
} finally {
121-
span.end(TraceUtil.END_SPAN_OPTIONS);
118+
span.end();
122119
}
123120
}
124121
}

‎google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/TraceUtil.java

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ public interface TraceUtil {
3434
static final String SPAN_NAME_LOOKUP = "Lookup";
3535
static final String SPAN_NAME_COMMIT = "Commit";
3636
static final String SPAN_NAME_RUN_QUERY = "RunQuery";
37+
static final String SPAN_NAME_RUN_AGGREGATION_QUERY = "RunAggregationQuery";
3738

3839
/**
3940
* Creates and returns an instance of the TraceUtil class.

‎google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecoratorTest.java

Lines changed: 2 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -15,11 +15,8 @@
1515
*/
1616
package com.google.cloud.datastore;
1717

18-
import static com.google.cloud.datastore.TraceUtil.END_SPAN_OPTIONS;
19-
import static com.google.cloud.datastore.TraceUtil.SPAN_NAME_RUN_AGGREGATION_QUERY;
2018
import static com.google.common.truth.Truth.assertThat;
2119
import static com.google.rpc.Code.UNAVAILABLE;
22-
import static org.easymock.EasyMock.createNiceMock;
2320
import static org.easymock.EasyMock.createStrictMock;
2421
import static org.easymock.EasyMock.expect;
2522
import static org.easymock.EasyMock.replay;
@@ -29,8 +26,6 @@
2926
import com.google.cloud.datastore.spi.v1.DatastoreRpc;
3027
import com.google.datastore.v1.RunAggregationQueryRequest;
3128
import com.google.datastore.v1.RunAggregationQueryResponse;
32-
import io.opencensus.trace.Span;
33-
import io.opencensus.trace.Tracer;
3429
import org.junit.Before;
3530
import org.junit.Test;
3631

@@ -49,15 +44,13 @@ public class RetryAndTraceDatastoreRpcDecoratorTest {
4944
@Before
5045
public void setUp() throws Exception {
5146
mockDatastoreRpc = createStrictMock(DatastoreRpc.class);
52-
mockTraceUtil = createStrictMock(TraceUtil.class);
5347
datastoreRpcDecorator =
5448
new RetryAndTraceDatastoreRpcDecorator(
5549
mockDatastoreRpc, mockTraceUtil, retrySettings, datastoreOptions);
5650
}
5751

5852
@Test
5953
public void testRunAggregationQuery() {
60-
Span mockSpan = createStrictMock(Span.class);
6154
RunAggregationQueryRequest aggregationQueryRequest =
6255
RunAggregationQueryRequest.getDefaultInstance();
6356
RunAggregationQueryResponse aggregationQueryResponse =
@@ -69,16 +62,13 @@ public void testRunAggregationQuery() {
6962
UNAVAILABLE.getNumber(), "API not accessible currently", UNAVAILABLE.name()))
7063
.times(2)
7164
.andReturn(aggregationQueryResponse);
72-
expect(mockTraceUtil.startSpan(SPAN_NAME_RUN_AGGREGATION_QUERY)).andReturn(mockSpan);
73-
expect(mockTraceUtil.getTracer()).andReturn(createNiceMock(Tracer.class));
74-
mockSpan.end(END_SPAN_OPTIONS);
7565

76-
replay(mockDatastoreRpc, mockTraceUtil, mockSpan);
66+
replay(mockDatastoreRpc);
7767

7868
RunAggregationQueryResponse actualAggregationQueryResponse =
7969
datastoreRpcDecorator.runAggregationQuery(aggregationQueryRequest);
8070

8171
assertThat(actualAggregationQueryResponse).isSameInstanceAs(aggregationQueryResponse);
82-
verify(mockDatastoreRpc, mockTraceUtil, mockSpan);
72+
verify(mockDatastoreRpc);
8373
}
8474
}

‎google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java

Lines changed: 78 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,9 +16,12 @@
1616

1717
package com.google.cloud.datastore.it;
1818

19+
import static com.google.cloud.datastore.aggregation.Aggregation.count;
1920
import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_COMMIT;
2021
import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_LOOKUP;
22+
import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_RUN_AGGREGATION_QUERY;
2123
import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_RUN_QUERY;
24+
import static com.google.common.truth.Truth.assertThat;
2225
import static io.opentelemetry.semconv.resource.attributes.ResourceAttributes.SERVICE_NAME;
2326
import static org.junit.Assert.assertEquals;
2427
import static org.junit.Assert.assertFalse;
@@ -27,12 +30,16 @@
2730
import static org.junit.Assert.assertTrue;
2831

2932
import com.google.api.gax.rpc.NotFoundException;
33+
import com.google.cloud.datastore.AggregationQuery;
34+
import com.google.cloud.datastore.AggregationResult;
35+
import com.google.cloud.datastore.AggregationResults;
3036
import com.google.cloud.datastore.Datastore;
3137
import com.google.cloud.datastore.DatastoreOptions;
3238
import com.google.cloud.datastore.Entity;
3339
import com.google.cloud.datastore.Key;
3440
import com.google.cloud.datastore.Query;
3541
import com.google.cloud.datastore.QueryResults;
42+
import com.google.cloud.datastore.StructuredQuery;
3643
import com.google.cloud.datastore.StructuredQuery.PropertyFilter;
3744
import com.google.cloud.datastore.testing.RemoteDatastoreHelper;
3845
import com.google.cloud.opentelemetry.trace.TraceConfiguration;
@@ -95,6 +102,7 @@
95102
// 5. Traces are read-back using TraceServiceClient and verified against expected Call Stacks.
96103
@RunWith(TestParameterInjector.class)
97104
public class ITE2ETracingTest {
105+
98106
protected boolean isUsingGlobalOpenTelemetrySDK() {
99107
return useGlobalOpenTelemetrySDK;
100108
}
@@ -214,6 +222,10 @@ private boolean dfsContainsCallStack(long spanId, List<String> expectedCallStack
214222

215223
private static Key KEY2;
216224

225+
private static Key KEY3;
226+
227+
private static Key KEY4;
228+
217229
// Random int generator for trace ID and span ID
218230
private static Random random;
219231

@@ -309,10 +321,17 @@ public void before() throws Exception {
309321
.setNamespace(options.getNamespace())
310322
.build();
311323
KEY2 =
324+
Key.newBuilder(projectId, kind1, "key3", options.getDatabaseId())
325+
.setNamespace(options.getNamespace())
326+
.build();
327+
KEY3 =
328+
Key.newBuilder(projectId, kind1, "key4", options.getDatabaseId())
329+
.setNamespace(options.getNamespace())
330+
.build();
331+
KEY4 =
312332
Key.newBuilder(projectId, kind1, "key2", options.getDatabaseId())
313333
.setNamespace(options.getNamespace())
314334
.build();
315-
316335
// Set up the tracer for custom TraceID injection
317336
rootSpanName =
318337
String.format("%s%d", this.getClass().getSimpleName(), System.currentTimeMillis());
@@ -658,4 +677,62 @@ public void runQueryTraceTest() throws Exception {
658677

659678
fetchAndValidateTrace(customSpanContext.getTraceId(), SPAN_NAME_RUN_QUERY);
660679
}
680+
681+
@Test
682+
public void runAggregationQueryTraceTest() throws Exception {
683+
Entity entity1 =
684+
Entity.newBuilder(KEY1)
685+
.set("pepper_name", "jalapeno")
686+
.set("max_scoville_level", 10000)
687+
.build();
688+
Entity entity2 =
689+
Entity.newBuilder(KEY2)
690+
.set("pepper_name", "serrano")
691+
.set("max_scoville_level", 25000)
692+
.build();
693+
Entity entity3 =
694+
Entity.newBuilder(KEY3)
695+
.set("pepper_name", "habanero")
696+
.set("max_scoville_level", 350000)
697+
.build();
698+
Entity entity4 =
699+
Entity.newBuilder(KEY4)
700+
.set("pepper_name", "ghost")
701+
.set("max_scoville_level", 1500000)
702+
.build();
703+
704+
List<Entity> entityList = new ArrayList<>();
705+
entityList.add(entity1);
706+
entityList.add(entity2);
707+
entityList.add(entity3);
708+
entityList.add(entity4);
709+
710+
List<Entity> response = datastore.add(entity1, entity2, entity3, entity4);
711+
assertEquals(entityList, response);
712+
713+
Span rootSpan = getNewRootSpanWithContext();
714+
try (Scope ignored = rootSpan.makeCurrent()) {
715+
PropertyFilter mediumSpicyFilters = PropertyFilter.lt("max_scoville_level", 100000);
716+
StructuredQuery<Entity> mediumSpicyQuery =
717+
Query.newEntityQueryBuilder()
718+
.setKind(KEY1.getKind())
719+
.setFilter(mediumSpicyFilters)
720+
.build();
721+
AggregationQuery countSpicyPeppers =
722+
Query.newAggregationQueryBuilder()
723+
.addAggregation(count().as("count"))
724+
.over(mediumSpicyQuery)
725+
.build();
726+
AggregationResults results = datastore.runAggregation(countSpicyPeppers);
727+
assertThat(results.size()).isEqualTo(1);
728+
AggregationResult result = results.get(0);
729+
assertThat(result.getLong("count")).isEqualTo(2L);
730+
} finally {
731+
rootSpan.end();
732+
}
733+
734+
waitForTracesToComplete();
735+
736+
fetchAndValidateTrace(customSpanContext.getTraceId(), SPAN_NAME_RUN_AGGREGATION_QUERY);
737+
}
661738
}

0 commit comments

Comments
 (0)