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 804673b3a..7a81a469d 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 @@ -285,11 +285,14 @@ 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_RUN_QUERY); ReadOptions readOptions = requestPb.getReadOptions(); - span.setAttribute( - "isTransactional", readOptions.hasTransaction() || readOptions.hasNewTransaction()); + boolean isTransactional = readOptions.hasTransaction() || readOptions.hasNewTransaction(); + String spanName = + (isTransactional + ? com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN_QUERY + : com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_RUN_QUERY); + com.google.cloud.datastore.telemetry.TraceUtil.Span span = otelTraceUtil.startSpan(spanName); + span.setAttribute("isTransactional", isTransactional); span.setAttribute("readConsistency", readOptions.getReadConsistency().toString()); try (com.google.cloud.datastore.telemetry.TraceUtil.Scope ignored = span.makeCurrent()) { @@ -302,7 +305,7 @@ com.google.datastore.v1.RunQueryResponse runQuery( : TRANSACTION_OPERATION_EXCEPTION_HANDLER, getOptions().getClock()); span.addEvent( - com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_RUN_QUERY + ": Completed", + spanName + ": 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 a1901fe0d..d63cdda34 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 @@ -42,6 +42,7 @@ public interface TraceUtil { static final String SPAN_NAME_BEGIN_TRANSACTION = "Transaction.Begin"; static final String SPAN_NAME_TRANSACTION_LOOKUP = "Transaction.Lookup"; static final String SPAN_NAME_TRANSACTION_COMMIT = "Transaction.Commit"; + static final String SPAN_NAME_TRANSACTION_RUN_QUERY = "Transaction.RunQuery"; static final String SPAN_NAME_ROLLBACK = "Transaction.Rollback"; static final String SPAN_NAME_TRANSACTION_RUN_AGGREGATION_QUERY = "Transaction.RunAggregationQuery"; 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 57f4d5b5c..161476af1 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 @@ -26,6 +26,7 @@ import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_COMMIT; import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_LOOKUP; import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN; +import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_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; @@ -795,8 +796,49 @@ public void transactionalLookupTest() throws Exception { Collections.singletonList(SPAN_NAME_TRANSACTION_COMMIT))); } + @Test + public void transactionQueryTest() throws Exception { + // Set up + 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); + + assertNotNull(customSpanContext); + + // Test + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ignored = rootSpan.makeCurrent()) { + Transaction transaction = datastore.newTransaction(); + PropertyFilter filter = PropertyFilter.eq("test_field", entity1.getValue("test_field")); + Query query = + Query.newEntityQueryBuilder().setKind(KEY1.getKind()).setFilter(filter).build(); + QueryResults queryResults = transaction.run(query); + transaction.commit(); + assertTrue(queryResults.hasNext()); + assertEquals(entity1, queryResults.next()); + assertFalse(queryResults.hasNext()); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + fetchAndValidateTrace( + customSpanContext.getTraceId(), + /*numExpectedSpans=*/ 3, + Arrays.asList( + Collections.singletonList(SPAN_NAME_BEGIN_TRANSACTION), + Collections.singletonList(SPAN_NAME_TRANSACTION_RUN_QUERY), + Collections.singletonList(SPAN_NAME_TRANSACTION_COMMIT))); + } + @Test public void runInTransactionQueryTest() throws Exception { + // Set up 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<>(); @@ -837,7 +879,7 @@ public void runInTransactionQueryTest() throws Exception { } @Test - public void runInTransactionAggregationQueryTest() throws Exception {} + public void transactionRunQueryTest() throws Exception {} @Test public void readWriteTransactionTraceTest() throws Exception {}