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 cb60685c7..ebf5aa7e9 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_RUN_QUERY); + 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_RUN_QUERY + ": 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 fe4effb1d..ccb5cfe7a 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,10 +31,9 @@ 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_COMMIT = "Commit"; + 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 2fb139289..0212b5e42 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 @@ -18,6 +18,7 @@ import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_COMMIT; import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_LOOKUP; +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; @@ -30,6 +31,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; @@ -301,11 +305,11 @@ public void before() throws Exception { String projectId = options.getProjectId(); String kind1 = "kind1"; KEY1 = - Key.newBuilder(projectId, kind1, "name1", options.getDatabaseId()) + Key.newBuilder(projectId, kind1, "key1", options.getDatabaseId()) .setNamespace(options.getNamespace()) .build(); KEY2 = - Key.newBuilder(projectId, kind1, "name2", options.getDatabaseId()) + Key.newBuilder(projectId, kind1, "key2", options.getDatabaseId()) .setNamespace(options.getNamespace()) .build(); @@ -594,7 +598,6 @@ public void updateTraceTest() throws Exception { assertEquals(entityList, response); Span rootSpan = getNewRootSpanWithContext(); - try (Scope ignored = rootSpan.makeCurrent()) { Entity entity1_update = Entity.newBuilder(entity1).set("test_field", "new_test_value1").build(); @@ -625,7 +628,34 @@ public void deleteTraceTest() throws Exception { rootSpan.end(); } waitForTracesToComplete(); - 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(); + 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_RUN_QUERY); + } }