Merge branch 'tracing-main' into jimit/runquery-trace-1

googleapis · May 20, 2024 · 940a01c · 940a01c
2 parents 1bfb706 + 727c398
commit 940a01c
Show file tree

Hide file tree

Showing 3 changed files with 92 additions and 5 deletions.
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
@@ -634,8 +634,11 @@ private com.google.datastore.v1.CommitResponse commitMutation(
 
   com.google.datastore.v1.CommitResponse commit(
       final com.google.datastore.v1.CommitRequest requestPb) {
-    Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_COMMIT);
-    try (Scope scope = traceUtil.getTracer().withSpan(span)) {
+    com.google.cloud.datastore.telemetry.TraceUtil.Span span =
+        otelTraceUtil.startSpan(com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_COMMIT);
+    span.setAttribute("isTransactional", requestPb.hasTransaction());
+
+    try (com.google.cloud.datastore.telemetry.TraceUtil.Scope ignored = span.makeCurrent()) {
       return RetryHelper.runWithRetries(
           () -> datastoreRpc.commit(requestPb),
           retrySettings,
@@ -644,10 +647,10 @@ com.google.datastore.v1.CommitResponse commit(
               : TRANSACTION_OPERATION_EXCEPTION_HANDLER,
           getOptions().getClock());
     } 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
@@ -32,6 +32,7 @@ public interface TraceUtil {
   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";
 
   /**

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
@@ -16,6 +16,7 @@
 
 package com.google.cloud.datastore.it;
 
+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;
@@ -548,7 +549,45 @@ public void lookupTraceTest() throws Exception {
   }
 
   @Test
-  public void runQueryTraceTest() throws Exception {
+  public void commitTraceTest() throws Exception {
+    assertNotNull(customSpanContext);
+
+    Span rootSpan = getNewRootSpanWithContext();
+
+    Entity entity1 = Entity.newBuilder(KEY1).set("test_key", "test_value").build();
+    try (Scope ignored = rootSpan.makeCurrent()) {
+      Entity response = datastore.add(entity1);
+      assertEquals(entity1, response);
+    } finally {
+      rootSpan.end();
+    }
+    waitForTracesToComplete();
+
+    fetchAndValidateTrace(customSpanContext.getTraceId(), SPAN_NAME_COMMIT);
+  }
+
+  @Test
+  public void putTraceTest() throws Exception {
+    assertNotNull(customSpanContext);
+
+    Span rootSpan = getNewRootSpanWithContext();
+
+    Entity entity1 = Entity.newBuilder(KEY1).set("test_key", "test_value").build();
+    try (Scope ignored = rootSpan.makeCurrent()) {
+      Entity response = datastore.put(entity1);
+      assertEquals(entity1, response);
+    } finally {
+      rootSpan.end();
+    }
+    waitForTracesToComplete();
+
+    fetchAndValidateTrace(customSpanContext.getTraceId(), SPAN_NAME_COMMIT);
+  }
+
+  @Test
+  public void updateTraceTest() throws Exception {
+    assertNotNull(customSpanContext);
+
     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<>();
@@ -558,6 +597,50 @@ public void runQueryTraceTest() throws Exception {
     List<Entity> response = datastore.add(entity1, entity2);
     assertEquals(entityList, response);
 
+    Span rootSpan = getNewRootSpanWithContext();
+    try (Scope ignored = rootSpan.makeCurrent()) {
+      Entity entity1_update =
+          Entity.newBuilder(entity1).set("test_field", "new_test_value1").build();
+      Entity entity2_update =
+          Entity.newBuilder(entity2).set("test_field", "new_test_value1").build();
+      datastore.update(entity1_update, entity2_update);
+    } finally {
+      rootSpan.end();
+    }
+    waitForTracesToComplete();
+
+    fetchAndValidateTrace(customSpanContext.getTraceId(), SPAN_NAME_COMMIT);
+  }
+
+  @Test
+  public void deleteTraceTest() throws Exception {
+    assertNotNull(customSpanContext);
+
+    Entity entity1 = Entity.newBuilder(KEY1).set("test_key", "test_value").build();
+    Entity response = datastore.put(entity1);
+    assertEquals(entity1, response);
+
+    Span rootSpan = getNewRootSpanWithContext();
+
+    try (Scope ignored = rootSpan.makeCurrent()) {
+      datastore.delete(entity1.getKey());
+    } finally {
+      rootSpan.end();
+    }
+    waitForTracesToComplete();
+    fetchAndValidateTrace(customSpanContext.getTraceId(), SPAN_NAME_COMMIT);
+  }
+
+  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"));