Skip to content
This repository was archived by the owner on Feb 24, 2026. It is now read-only.

Commit 9bb337c

Browse files
jimit-j-shahjimit-j-shah
authored
feat: Added tracing for Transaction.RunQuery (#1499)
* feat: Added span for Transactional RunQuery - tested * fix: lint * fix: patch apply issues * fix: refactor using boolean flag
1 parent 5995b3e commit 9bb337c

File tree

3 files changed

+52
-6
lines changed

3 files changed

+52
-6
lines changed

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

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -285,11 +285,14 @@ public AggregationResults runAggregation(
285285

286286
com.google.datastore.v1.RunQueryResponse runQuery(
287287
final com.google.datastore.v1.RunQueryRequest requestPb) {
288-
com.google.cloud.datastore.telemetry.TraceUtil.Span span =
289-
otelTraceUtil.startSpan(com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_RUN_QUERY);
290288
ReadOptions readOptions = requestPb.getReadOptions();
291-
span.setAttribute(
292-
"isTransactional", readOptions.hasTransaction() || readOptions.hasNewTransaction());
289+
boolean isTransactional = readOptions.hasTransaction() || readOptions.hasNewTransaction();
290+
String spanName =
291+
(isTransactional
292+
? com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN_QUERY
293+
: com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_RUN_QUERY);
294+
com.google.cloud.datastore.telemetry.TraceUtil.Span span = otelTraceUtil.startSpan(spanName);
295+
span.setAttribute("isTransactional", isTransactional);
293296
span.setAttribute("readConsistency", readOptions.getReadConsistency().toString());
294297

295298
try (com.google.cloud.datastore.telemetry.TraceUtil.Scope ignored = span.makeCurrent()) {
@@ -302,7 +305,7 @@ com.google.datastore.v1.RunQueryResponse runQuery(
302305
: TRANSACTION_OPERATION_EXCEPTION_HANDLER,
303306
getOptions().getClock());
304307
span.addEvent(
305-
com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_RUN_QUERY + ": Completed",
308+
spanName + ": Completed",
306309
new ImmutableMap.Builder<String, Object>()
307310
.put("Received", response.getBatch().getEntityResultsCount())
308311
.put("More results", response.getBatch().getMoreResults().toString())

‎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
@@ -42,6 +42,7 @@ public interface TraceUtil {
4242
static final String SPAN_NAME_BEGIN_TRANSACTION = "Transaction.Begin";
4343
static final String SPAN_NAME_TRANSACTION_LOOKUP = "Transaction.Lookup";
4444
static final String SPAN_NAME_TRANSACTION_COMMIT = "Transaction.Commit";
45+
static final String SPAN_NAME_TRANSACTION_RUN_QUERY = "Transaction.RunQuery";
4546
static final String SPAN_NAME_ROLLBACK = "Transaction.Rollback";
4647
static final String SPAN_NAME_TRANSACTION_RUN_AGGREGATION_QUERY =
4748
"Transaction.RunAggregationQuery";

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

Lines changed: 43 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_COMMIT;
2727
import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_LOOKUP;
2828
import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN;
29+
import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN_QUERY;
2930
import static com.google.common.truth.Truth.assertThat;
3031
import static io.opentelemetry.semconv.resource.attributes.ResourceAttributes.SERVICE_NAME;
3132
import static org.junit.Assert.assertEquals;
@@ -795,8 +796,49 @@ public void transactionalLookupTest() throws Exception {
795796
Collections.singletonList(SPAN_NAME_TRANSACTION_COMMIT)));
796797
}
797798

799+
@Test
800+
public void transactionQueryTest() throws Exception {
801+
// Set up
802+
Entity entity1 = Entity.newBuilder(KEY1).set("test_field", "test_value1").build();
803+
Entity entity2 = Entity.newBuilder(KEY2).set("test_field", "test_value2").build();
804+
List<Entity> entityList = new ArrayList<>();
805+
entityList.add(entity1);
806+
entityList.add(entity2);
807+
808+
List<Entity> response = datastore.add(entity1, entity2);
809+
assertEquals(entityList, response);
810+
811+
assertNotNull(customSpanContext);
812+
813+
// Test
814+
Span rootSpan = getNewRootSpanWithContext();
815+
try (Scope ignored = rootSpan.makeCurrent()) {
816+
Transaction transaction = datastore.newTransaction();
817+
PropertyFilter filter = PropertyFilter.eq("test_field", entity1.getValue("test_field"));
818+
Query<Entity> query =
819+
Query.newEntityQueryBuilder().setKind(KEY1.getKind()).setFilter(filter).build();
820+
QueryResults<Entity> queryResults = transaction.run(query);
821+
transaction.commit();
822+
assertTrue(queryResults.hasNext());
823+
assertEquals(entity1, queryResults.next());
824+
assertFalse(queryResults.hasNext());
825+
} finally {
826+
rootSpan.end();
827+
}
828+
waitForTracesToComplete();
829+
830+
fetchAndValidateTrace(
831+
customSpanContext.getTraceId(),
832+
/*numExpectedSpans=*/ 3,
833+
Arrays.asList(
834+
Collections.singletonList(SPAN_NAME_BEGIN_TRANSACTION),
835+
Collections.singletonList(SPAN_NAME_TRANSACTION_RUN_QUERY),
836+
Collections.singletonList(SPAN_NAME_TRANSACTION_COMMIT)));
837+
}
838+
798839
@Test
799840
public void runInTransactionQueryTest() throws Exception {
841+
// Set up
800842
Entity entity1 = Entity.newBuilder(KEY1).set("test_field", "test_value1").build();
801843
Entity entity2 = Entity.newBuilder(KEY2).set("test_field", "test_value2").build();
802844
List<Entity> entityList = new ArrayList<>();
@@ -837,7 +879,7 @@ public void runInTransactionQueryTest() throws Exception {
837879
}
838880

839881
@Test
840-
public void runInTransactionAggregationQueryTest() throws Exception {}
882+
public void transactionRunQueryTest() throws Exception {}
841883

842884
@Test
843885
public void readWriteTransactionTraceTest() throws Exception {}

0 commit comments

Comments
 (0)