From b739d6919dbe85d92bb45c41055a944e0dcc96a7 Mon Sep 17 00:00:00 2001 From: jimit-j-shah <57637300+jimit-j-shah@users.noreply.github.com> Date: Wed, 15 May 2024 16:25:18 -0700 Subject: [PATCH 1/9] feat: RunQuery trace instrumentation --- .../google/cloud/datastore/DatastoreImpl.java | 37 +++++++++---- .../cloud/datastore/telemetry/TraceUtil.java | 3 ++ .../cloud/datastore/it/ITE2ETracingTest.java | 54 +++++++++++++++++-- 3 files changed, 79 insertions(+), 15 deletions(-) diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java index 401e3ed54..cf84d62b6 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java @@ -35,6 +35,7 @@ import com.google.datastore.v1.ExplainOptions; import com.google.datastore.v1.ReadOptions; import com.google.datastore.v1.ReserveIdsRequest; +import com.google.datastore.v1.RunQueryResponse; import com.google.datastore.v1.TransactionOptions; import com.google.protobuf.ByteString; import io.opencensus.common.Scope; @@ -240,20 +241,34 @@ public AggregationResults runAggregation( com.google.datastore.v1.RunQueryResponse runQuery( final com.google.datastore.v1.RunQueryRequest requestPb) { - Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_RUNQUERY); - try (Scope scope = traceUtil.getTracer().withSpan(span)) { - return RetryHelper.runWithRetries( - () -> datastoreRpc.runQuery(requestPb), - retrySettings, - requestPb.getReadOptions().getTransaction().isEmpty() - ? EXCEPTION_HANDLER - : TRANSACTION_OPERATION_EXCEPTION_HANDLER, - getOptions().getClock()); + com.google.cloud.datastore.telemetry.TraceUtil.Span span = + otelTraceUtil.startSpan(com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_RUNQUERY); + ReadOptions readOptions = requestPb.getReadOptions(); + span.setAttribute( + "isTransactional", readOptions.hasTransaction() || readOptions.hasNewTransaction()); + span.setAttribute("readConsistency", readOptions.getReadConsistency().toString()); + + try (com.google.cloud.datastore.telemetry.TraceUtil.Scope ignored = span.makeCurrent()) { + RunQueryResponse response = + RetryHelper.runWithRetries( + () -> datastoreRpc.runQuery(requestPb), + retrySettings, + requestPb.getReadOptions().getTransaction().isEmpty() + ? EXCEPTION_HANDLER + : TRANSACTION_OPERATION_EXCEPTION_HANDLER, + getOptions().getClock()); + span.addEvent( + com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_RUNQUERY + ": Completed", + new ImmutableMap.Builder() + .put("Received", response.getBatch().getEntityResultsCount()) + .put("More results", response.getBatch().getMoreResults().toString()) + .build()); + return response; } catch (RetryHelperException e) { - span.setStatus(Status.UNKNOWN.withDescription(e.getMessage())); + span.end(e); throw DatastoreException.translateAndThrow(e); } finally { - span.end(TraceUtil.END_SPAN_OPTIONS); + span.end(); } } diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/TraceUtil.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/TraceUtil.java index c867a5525..14d8801fc 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/TraceUtil.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/TraceUtil.java @@ -33,6 +33,9 @@ public interface TraceUtil { static final String LIBRARY_NAME = "com.google.cloud.datastore"; static final String SPAN_NAME_LOOKUP = "Lookup"; + + static final String SPAN_NAME_RUNQUERY = "RunQuery"; + /** * Creates and returns an instance of the TraceUtil class. * diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java index 24a9a7149..d498d38ba 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java @@ -17,6 +17,7 @@ package com.google.cloud.datastore.it; import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_LOOKUP; +import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_RUNQUERY; import static io.opentelemetry.semconv.resource.attributes.ResourceAttributes.SERVICE_NAME; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; @@ -29,6 +30,9 @@ import com.google.cloud.datastore.DatastoreOptions; import com.google.cloud.datastore.Entity; import com.google.cloud.datastore.Key; +import com.google.cloud.datastore.Query; +import com.google.cloud.datastore.QueryResults; +import com.google.cloud.datastore.StructuredQuery.PropertyFilter; import com.google.cloud.datastore.testing.RemoteDatastoreHelper; import com.google.cloud.opentelemetry.trace.TraceConfiguration; import com.google.cloud.opentelemetry.trace.TraceExporter; @@ -207,6 +211,8 @@ private boolean dfsContainsCallStack(long spanId, List expectedCallStack private static Key KEY1; + private static Key KEY2; + // Random int generator for trace ID and span ID private static Random random; @@ -233,9 +239,14 @@ private boolean dfsContainsCallStack(long spanId, List expectedCallStack private static Datastore datastore; + private static RemoteDatastoreHelper remoteDatastoreHelper; + @TestParameter boolean useGlobalOpenTelemetrySDK; - @TestParameter({"default", "test-db"}) + @TestParameter({ + /*(default)*/ + "", "test-db" + }) String datastoreNamedDatabase; @BeforeClass @@ -280,8 +291,7 @@ public void before() throws Exception { // but because gRPC traces need to be deterministically force-flushed for every test String namedDb = datastoreNamedDatabase(); logger.log(Level.INFO, "Integration test using named database " + namedDb); - RemoteDatastoreHelper remoteDatastoreHelper = - RemoteDatastoreHelper.create(namedDb, openTelemetrySdk); + remoteDatastoreHelper = RemoteDatastoreHelper.create(namedDb, openTelemetrySdk); options = remoteDatastoreHelper.getOptions(); datastore = options.getService(); @@ -292,7 +302,14 @@ public void before() throws Exception { String projectId = options.getProjectId(); String kind1 = "kind1"; - KEY1 = Key.newBuilder(projectId, kind1, "name", options.getDatabaseId()).build(); + KEY1 = + Key.newBuilder(projectId, kind1, "key1", options.getDatabaseId()) + .setNamespace(options.getNamespace()) + .build(); + KEY2 = + Key.newBuilder(projectId, kind1, "key2", options.getDatabaseId()) + .setNamespace(options.getNamespace()) + .build(); // Set up the tracer for custom TraceID injection rootSpanName = @@ -319,6 +336,7 @@ public void after() throws Exception { if (isUsingGlobalOpenTelemetrySDK()) { GlobalOpenTelemetry.resetForTest(); } + remoteDatastoreHelper.deleteNamespace(); rootSpanName = null; tracer = null; retrievedTrace = null; @@ -527,4 +545,32 @@ public void lookupTraceTest() throws Exception { fetchAndValidateTrace(customSpanContext.getTraceId(), SPAN_NAME_LOOKUP); } + + @Test + public void runQueryTraceTest() throws Exception { + Entity entity1 = Entity.newBuilder(KEY1).set("test_field", "test_value1").build(); + Entity entity2 = Entity.newBuilder(KEY2).set("test_field", "test_value2").build(); + List entityList = new ArrayList<>(); + entityList.add(entity1); + entityList.add(entity2); + + List response = datastore.add(entity1, entity2); + assertEquals(entityList, response); + + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ignored = rootSpan.makeCurrent()) { + PropertyFilter filter = PropertyFilter.eq("test_field", entity1.getValue("test_field")); + Query query = + Query.newEntityQueryBuilder().setKind(KEY1.getKind()).setFilter(filter).build(); + QueryResults queryResults = datastore.run(query); + assertTrue(queryResults.hasNext()); + assertEquals(entity1, queryResults.next()); + assertFalse(queryResults.hasNext()); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + fetchAndValidateTrace(customSpanContext.getTraceId(), SPAN_NAME_RUNQUERY); + } } From ec346cc8aea2b80489839ff9e7c4039986ae55d9 Mon Sep 17 00:00:00 2001 From: jimit-j-shah <57637300+jimit-j-shah@users.noreply.github.com> Date: Wed, 15 May 2024 16:38:29 -0700 Subject: [PATCH 2/9] Formatting --- .../java/com/google/cloud/datastore/telemetry/TraceUtil.java | 2 -- 1 file changed, 2 deletions(-) diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/TraceUtil.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/TraceUtil.java index 14d8801fc..e59584958 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/TraceUtil.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/TraceUtil.java @@ -31,9 +31,7 @@ public interface TraceUtil { static final String ATTRIBUTE_SERVICE_PREFIX = "gcp.datastore."; static final String ENABLE_TRACING_ENV_VAR = "DATASTORE_ENABLE_TRACING"; static final String LIBRARY_NAME = "com.google.cloud.datastore"; - static final String SPAN_NAME_LOOKUP = "Lookup"; - static final String SPAN_NAME_RUNQUERY = "RunQuery"; /** From cb4cd39b81f80002cd1d5cb82a1df839590924df Mon Sep 17 00:00:00 2001 From: jimit-j-shah <57637300+jimit-j-shah@users.noreply.github.com> Date: Wed, 15 May 2024 16:39:55 -0700 Subject: [PATCH 3/9] Formatting --- .../java/com/google/cloud/datastore/it/ITE2ETracingTest.java | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java index d498d38ba..476ce23fc 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java @@ -244,8 +244,9 @@ private boolean dfsContainsCallStack(long spanId, List expectedCallStack @TestParameter boolean useGlobalOpenTelemetrySDK; @TestParameter({ - /*(default)*/ - "", "test-db" + /*(default)*/ + "", + "test-db" }) String datastoreNamedDatabase; From 1bfb706869a1e7d50ea43fc2e583e4511d5ae449 Mon Sep 17 00:00:00 2001 From: jimit-j-shah <57637300+jimit-j-shah@users.noreply.github.com> Date: Thu, 16 May 2024 10:49:34 -0700 Subject: [PATCH 4/9] Refactor: s/RUNQUERY/RUN_QUERY --- .../main/java/com/google/cloud/datastore/DatastoreImpl.java | 4 ++-- .../java/com/google/cloud/datastore/telemetry/TraceUtil.java | 2 +- .../java/com/google/cloud/datastore/it/ITE2ETracingTest.java | 4 ++-- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java index cf84d62b6..f025518ab 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java @@ -242,7 +242,7 @@ public AggregationResults runAggregation( com.google.datastore.v1.RunQueryResponse runQuery( final com.google.datastore.v1.RunQueryRequest requestPb) { com.google.cloud.datastore.telemetry.TraceUtil.Span span = - otelTraceUtil.startSpan(com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_RUNQUERY); + otelTraceUtil.startSpan(com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_RUN_QUERY); ReadOptions readOptions = requestPb.getReadOptions(); span.setAttribute( "isTransactional", readOptions.hasTransaction() || readOptions.hasNewTransaction()); @@ -258,7 +258,7 @@ com.google.datastore.v1.RunQueryResponse runQuery( : TRANSACTION_OPERATION_EXCEPTION_HANDLER, getOptions().getClock()); span.addEvent( - com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_RUNQUERY + ": Completed", + com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_RUN_QUERY + ": Completed", new ImmutableMap.Builder() .put("Received", response.getBatch().getEntityResultsCount()) .put("More results", response.getBatch().getMoreResults().toString()) diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/TraceUtil.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/TraceUtil.java index e59584958..5bf505638 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/TraceUtil.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/TraceUtil.java @@ -32,7 +32,7 @@ public interface TraceUtil { static final String ENABLE_TRACING_ENV_VAR = "DATASTORE_ENABLE_TRACING"; static final String LIBRARY_NAME = "com.google.cloud.datastore"; static final String SPAN_NAME_LOOKUP = "Lookup"; - static final String SPAN_NAME_RUNQUERY = "RunQuery"; + static final String SPAN_NAME_RUN_QUERY = "RunQuery"; /** * Creates and returns an instance of the TraceUtil class. diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java index 476ce23fc..6bea8827d 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java @@ -17,7 +17,7 @@ package com.google.cloud.datastore.it; import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_LOOKUP; -import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_RUNQUERY; +import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_RUN_QUERY; import static io.opentelemetry.semconv.resource.attributes.ResourceAttributes.SERVICE_NAME; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; @@ -572,6 +572,6 @@ public void runQueryTraceTest() throws Exception { } waitForTracesToComplete(); - fetchAndValidateTrace(customSpanContext.getTraceId(), SPAN_NAME_RUNQUERY); + fetchAndValidateTrace(customSpanContext.getTraceId(), SPAN_NAME_RUN_QUERY); } } From 339824d4dc8ccf8f7452d137efa86608cab476cb Mon Sep 17 00:00:00 2001 From: jimit-j-shah <57637300+jimit-j-shah@users.noreply.github.com> Date: Thu, 16 May 2024 15:14:25 -0700 Subject: [PATCH 5/9] feat: RunAggregationQuery Trace Instrumentation --- .../RetryAndTraceDatastoreRpcDecorator.java | 19 ++--- .../cloud/datastore/telemetry/TraceUtil.java | 1 + .../cloud/datastore/it/ITE2ETracingTest.java | 79 ++++++++++++++++++- 3 files changed, 87 insertions(+), 12 deletions(-) diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecorator.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecorator.java index c4a85caab..ba4adb5ff 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecorator.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecorator.java @@ -16,7 +16,6 @@ package com.google.cloud.datastore; import static com.google.cloud.BaseService.EXCEPTION_HANDLER; -import static com.google.cloud.datastore.TraceUtil.SPAN_NAME_RUN_AGGREGATION_QUERY; import com.google.api.core.InternalApi; import com.google.api.gax.retrying.RetrySettings; @@ -39,9 +38,6 @@ import com.google.datastore.v1.RunAggregationQueryResponse; import com.google.datastore.v1.RunQueryRequest; import com.google.datastore.v1.RunQueryResponse; -import io.opencensus.common.Scope; -import io.opencensus.trace.Span; -import io.opencensus.trace.Status; import java.util.concurrent.Callable; /** @@ -52,7 +48,7 @@ public class RetryAndTraceDatastoreRpcDecorator implements DatastoreRpc { private final DatastoreRpc datastoreRpc; - private final TraceUtil traceUtil; + private final com.google.cloud.datastore.telemetry.TraceUtil otelTraceUtil; private final RetrySettings retrySettings; private final DatastoreOptions datastoreOptions; @@ -62,9 +58,9 @@ public RetryAndTraceDatastoreRpcDecorator( RetrySettings retrySettings, DatastoreOptions datastoreOptions) { this.datastoreRpc = datastoreRpc; - this.traceUtil = traceUtil; this.retrySettings = retrySettings; this.datastoreOptions = datastoreOptions; + this.otelTraceUtil = datastoreOptions.getTraceUtil(); } @Override @@ -106,19 +102,20 @@ public RunQueryResponse runQuery(RunQueryRequest request) { @Override public RunAggregationQueryResponse runAggregationQuery(RunAggregationQueryRequest request) { return invokeRpc( - () -> datastoreRpc.runAggregationQuery(request), SPAN_NAME_RUN_AGGREGATION_QUERY); + () -> datastoreRpc.runAggregationQuery(request), + com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_RUN_AGGREGATION_QUERY); } public O invokeRpc(Callable block, String startSpan) { - Span span = traceUtil.startSpan(startSpan); - try (Scope scope = traceUtil.getTracer().withSpan(span)) { + com.google.cloud.datastore.telemetry.TraceUtil.Span span = otelTraceUtil.startSpan(startSpan); + try (com.google.cloud.datastore.telemetry.TraceUtil.Scope ignored = span.makeCurrent()) { return RetryHelper.runWithRetries( block, this.retrySettings, EXCEPTION_HANDLER, this.datastoreOptions.getClock()); } catch (RetryHelperException e) { - span.setStatus(Status.UNKNOWN.withDescription(e.getMessage())); + span.end(e); throw DatastoreException.translateAndThrow(e); } finally { - span.end(TraceUtil.END_SPAN_OPTIONS); + span.end(); } } } diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/TraceUtil.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/TraceUtil.java index 5bf505638..19c955955 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/TraceUtil.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/TraceUtil.java @@ -33,6 +33,7 @@ public interface TraceUtil { static final String LIBRARY_NAME = "com.google.cloud.datastore"; static final String SPAN_NAME_LOOKUP = "Lookup"; static final String SPAN_NAME_RUN_QUERY = "RunQuery"; + static final String SPAN_NAME_RUN_AGGREGATION_QUERY = "RunAggregationQuery"; /** * Creates and returns an instance of the TraceUtil class. diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java index 6bea8827d..64f00ec22 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java @@ -16,8 +16,11 @@ package com.google.cloud.datastore.it; +import static com.google.cloud.datastore.aggregation.Aggregation.count; import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_LOOKUP; +import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_RUN_AGGREGATION_QUERY; import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_RUN_QUERY; +import static com.google.common.truth.Truth.assertThat; import static io.opentelemetry.semconv.resource.attributes.ResourceAttributes.SERVICE_NAME; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; @@ -26,12 +29,16 @@ import static org.junit.Assert.assertTrue; import com.google.api.gax.rpc.NotFoundException; +import com.google.cloud.datastore.AggregationQuery; +import com.google.cloud.datastore.AggregationResult; +import com.google.cloud.datastore.AggregationResults; import com.google.cloud.datastore.Datastore; import com.google.cloud.datastore.DatastoreOptions; import com.google.cloud.datastore.Entity; import com.google.cloud.datastore.Key; import com.google.cloud.datastore.Query; import com.google.cloud.datastore.QueryResults; +import com.google.cloud.datastore.StructuredQuery; import com.google.cloud.datastore.StructuredQuery.PropertyFilter; import com.google.cloud.datastore.testing.RemoteDatastoreHelper; import com.google.cloud.opentelemetry.trace.TraceConfiguration; @@ -94,6 +101,7 @@ // 5. Traces are read-back using TraceServiceClient and verified against expected Call Stacks. @RunWith(TestParameterInjector.class) public class ITE2ETracingTest { + protected boolean isUsingGlobalOpenTelemetrySDK() { return useGlobalOpenTelemetrySDK; } @@ -213,6 +221,10 @@ private boolean dfsContainsCallStack(long spanId, List expectedCallStack private static Key KEY2; + private static Key KEY3; + + private static Key KEY4; + // Random int generator for trace ID and span ID private static Random random; @@ -308,10 +320,17 @@ public void before() throws Exception { .setNamespace(options.getNamespace()) .build(); KEY2 = + Key.newBuilder(projectId, kind1, "key3", options.getDatabaseId()) + .setNamespace(options.getNamespace()) + .build(); + KEY3 = + Key.newBuilder(projectId, kind1, "key4", options.getDatabaseId()) + .setNamespace(options.getNamespace()) + .build(); + KEY4 = Key.newBuilder(projectId, kind1, "key2", options.getDatabaseId()) .setNamespace(options.getNamespace()) .build(); - // Set up the tracer for custom TraceID injection rootSpanName = String.format("%s%d", this.getClass().getSimpleName(), System.currentTimeMillis()); @@ -574,4 +593,62 @@ public void runQueryTraceTest() throws Exception { fetchAndValidateTrace(customSpanContext.getTraceId(), SPAN_NAME_RUN_QUERY); } + + @Test + public void runAggregationQueryTraceTest() throws Exception { + Entity entity1 = + Entity.newBuilder(KEY1) + .set("pepper_name", "jalapeno") + .set("max_scoville_level", 10000) + .build(); + Entity entity2 = + Entity.newBuilder(KEY2) + .set("pepper_name", "serrano") + .set("max_scoville_level", 25000) + .build(); + Entity entity3 = + Entity.newBuilder(KEY3) + .set("pepper_name", "habanero") + .set("max_scoville_level", 350000) + .build(); + Entity entity4 = + Entity.newBuilder(KEY4) + .set("pepper_name", "ghost") + .set("max_scoville_level", 1500000) + .build(); + + List entityList = new ArrayList<>(); + entityList.add(entity1); + entityList.add(entity2); + entityList.add(entity3); + entityList.add(entity4); + + List response = datastore.add(entity1, entity2, entity3, entity4); + assertEquals(entityList, response); + + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ignored = rootSpan.makeCurrent()) { + PropertyFilter mediumSpicyFilters = PropertyFilter.lt("max_scoville_level", 100000); + StructuredQuery mediumSpicyQuery = + Query.newEntityQueryBuilder() + .setKind(KEY1.getKind()) + .setFilter(mediumSpicyFilters) + .build(); + AggregationQuery countSpicyPeppers = + Query.newAggregationQueryBuilder() + .addAggregation(count().as("count")) + .over(mediumSpicyQuery) + .build(); + AggregationResults results = datastore.runAggregation(countSpicyPeppers); + assertThat(results.size()).isEqualTo(1); + AggregationResult result = results.get(0); + assertThat(result.getLong("count")).isEqualTo(2L); + } finally { + rootSpan.end(); + } + + waitForTracesToComplete(); + + fetchAndValidateTrace(customSpanContext.getTraceId(), SPAN_NAME_RUN_AGGREGATION_QUERY); + } } From ec777fc80d68e7e7f1ed9d2f699bb2ba61875039 Mon Sep 17 00:00:00 2001 From: jimit-j-shah <57637300+jimit-j-shah@users.noreply.github.com> Date: Thu, 16 May 2024 16:13:14 -0700 Subject: [PATCH 6/9] 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). --- .../RetryAndTraceDatastoreRpcDecoratorTest.java | 13 +++---------- 1 file changed, 3 insertions(+), 10 deletions(-) diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecoratorTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecoratorTest.java index b86355afa..b4a8ddbfe 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecoratorTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecoratorTest.java @@ -19,7 +19,6 @@ import static com.google.cloud.datastore.TraceUtil.SPAN_NAME_RUN_AGGREGATION_QUERY; import static com.google.common.truth.Truth.assertThat; import static com.google.rpc.Code.UNAVAILABLE; -import static org.easymock.EasyMock.createNiceMock; import static org.easymock.EasyMock.createStrictMock; import static org.easymock.EasyMock.expect; import static org.easymock.EasyMock.replay; @@ -29,8 +28,7 @@ import com.google.cloud.datastore.spi.v1.DatastoreRpc; import com.google.datastore.v1.RunAggregationQueryRequest; import com.google.datastore.v1.RunAggregationQueryResponse; -import io.opencensus.trace.Span; -import io.opencensus.trace.Tracer; + import org.junit.Before; import org.junit.Test; @@ -49,7 +47,6 @@ public class RetryAndTraceDatastoreRpcDecoratorTest { @Before public void setUp() throws Exception { mockDatastoreRpc = createStrictMock(DatastoreRpc.class); - mockTraceUtil = createStrictMock(TraceUtil.class); datastoreRpcDecorator = new RetryAndTraceDatastoreRpcDecorator( mockDatastoreRpc, mockTraceUtil, retrySettings, datastoreOptions); @@ -57,7 +54,6 @@ public void setUp() throws Exception { @Test public void testRunAggregationQuery() { - Span mockSpan = createStrictMock(Span.class); RunAggregationQueryRequest aggregationQueryRequest = RunAggregationQueryRequest.getDefaultInstance(); RunAggregationQueryResponse aggregationQueryResponse = @@ -69,16 +65,13 @@ public void testRunAggregationQuery() { UNAVAILABLE.getNumber(), "API not accessible currently", UNAVAILABLE.name())) .times(2) .andReturn(aggregationQueryResponse); - expect(mockTraceUtil.startSpan(SPAN_NAME_RUN_AGGREGATION_QUERY)).andReturn(mockSpan); - expect(mockTraceUtil.getTracer()).andReturn(createNiceMock(Tracer.class)); - mockSpan.end(END_SPAN_OPTIONS); - replay(mockDatastoreRpc, mockTraceUtil, mockSpan); + replay(mockDatastoreRpc); RunAggregationQueryResponse actualAggregationQueryResponse = datastoreRpcDecorator.runAggregationQuery(aggregationQueryRequest); assertThat(actualAggregationQueryResponse).isSameInstanceAs(aggregationQueryResponse); - verify(mockDatastoreRpc, mockTraceUtil, mockSpan); + verify(mockDatastoreRpc); } } From a2a4788514dc8e84d49bdbd3810fc268ed76c3d7 Mon Sep 17 00:00:00 2001 From: jimit-j-shah <57637300+jimit-j-shah@users.noreply.github.com> Date: Fri, 17 May 2024 09:48:52 -0700 Subject: [PATCH 7/9] Formatting --- .../datastore/RetryAndTraceDatastoreRpcDecoratorTest.java | 3 --- 1 file changed, 3 deletions(-) diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecoratorTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecoratorTest.java index b4a8ddbfe..b01bcab82 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecoratorTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecoratorTest.java @@ -15,8 +15,6 @@ */ package com.google.cloud.datastore; -import static com.google.cloud.datastore.TraceUtil.END_SPAN_OPTIONS; -import static com.google.cloud.datastore.TraceUtil.SPAN_NAME_RUN_AGGREGATION_QUERY; import static com.google.common.truth.Truth.assertThat; import static com.google.rpc.Code.UNAVAILABLE; import static org.easymock.EasyMock.createStrictMock; @@ -28,7 +26,6 @@ import com.google.cloud.datastore.spi.v1.DatastoreRpc; import com.google.datastore.v1.RunAggregationQueryRequest; import com.google.datastore.v1.RunAggregationQueryResponse; - import org.junit.Before; import org.junit.Test; From 6e5e9caca8d3702393bfb62364853b817f8b8736 Mon Sep 17 00:00:00 2001 From: jimit-j-shah <57637300+jimit-j-shah@users.noreply.github.com> Date: Mon, 20 May 2024 11:37:22 -0700 Subject: [PATCH 8/9] Fixing @Test annotation missed after merge --- .../java/com/google/cloud/datastore/it/ITE2ETracingTest.java | 1 + 1 file changed, 1 insertion(+) diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java index f33a7e23a..cc733fbdb 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java @@ -631,6 +631,7 @@ public void deleteTraceTest() throws Exception { fetchAndValidateTrace(customSpanContext.getTraceId(), SPAN_NAME_COMMIT); } + @Test public void runQueryTraceTest() throws Exception { Entity entity1 = Entity.newBuilder(KEY1).set("test_field", "test_value1").build(); Entity entity2 = Entity.newBuilder(KEY2).set("test_field", "test_value2").build(); From e48c36d638a86f431d8989deb030ef8b66de7a73 Mon Sep 17 00:00:00 2001 From: jimit-j-shah <57637300+jimit-j-shah@users.noreply.github.com> Date: Mon, 20 May 2024 14:01:37 -0700 Subject: [PATCH 9/9] Formatting --- .../java/com/google/cloud/datastore/it/ITE2ETracingTest.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java index 4309c918f..357e748f1 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java @@ -652,7 +652,7 @@ public void deleteTraceTest() throws Exception { @Test public void runQueryTraceTest() throws Exception { - Entity entity1 = Entity.newBuilder(KEY1).set("test_field", "test_value1").build(); + Entity entity1 = Entity.newBuilder(KEY1).set("test_field", "test_value1").build(); Entity entity2 = Entity.newBuilder(KEY2).set("test_field", "test_value2").build(); List entityList = new ArrayList<>(); entityList.add(entity1);