feat: RunQuery trace instrumentation

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<String, Object>()
+              .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();
     }
   }
 

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

@@ -31,8 +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_RUN_QUERY = "RunQuery";
+
   /**
    * Creates and returns an instance of the TraceUtil class.
    *

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_RUN_QUERY;
 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<String> 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,15 @@ private boolean dfsContainsCallStack(long spanId, List<String> 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 +292,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 +303,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 +337,7 @@ public void after() throws Exception {
     if (isUsingGlobalOpenTelemetrySDK()) {
       GlobalOpenTelemetry.resetForTest();
     }
+    remoteDatastoreHelper.deleteNamespace();
     rootSpanName = null;
     tracer = null;
     retrievedTrace = null;
@@ -527,4 +546,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<Entity> entityList = new ArrayList<>();
+    entityList.add(entity1);
+    entityList.add(entity2);
+
+    List<Entity> 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<Entity> query =
+          Query.newEntityQueryBuilder().setKind(KEY1.getKind()).setFilter(filter).build();
+      QueryResults<Entity> 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);
+  }
 }