feat: RunQuery trace instrumentation (#1441)

* feat: RunQuery trace instrumentation

Author: jimit-j-shah

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,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.

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<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);
+  }
 }