test: Additional Transaction Testing and cleanup OpenCensus usage

google-cloud-datastore/clirr-ignored-differences.xml

@@ -55,4 +55,16 @@
         <method>java.lang.Object execute(com.google.cloud.datastore.Query, com.google.cloud.datastore.ReadOption[])</method>
         <differenceType>7004</differenceType>
     </difference>
+    <difference>
+        <className>com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecorator</className>
+        <method>RetryAndTraceDatastoreRpcDecorator(com.google.cloud.datastore.spi.v1.DatastoreRpc, com.google.cloud.datastore.TraceUtil, com.google.api.gax.retrying.RetrySettings, com.google.cloud.datastore.DatastoreOptions)</method>
+        <to>RetryAndTraceDatastoreRpcDecorator(com.google.cloud.datastore.spi.v1.DatastoreRpc, com.google.cloud.datastore.telemetry.TraceUtil, com.google.api.gax.retrying.RetrySettings, com.google.cloud.datastore.DatastoreOptions)</to>
+        <differenceType>7005</differenceType>
+    </difference>
+
+    <!-- Class removed -->
+    <difference>
+        <className>com/google/cloud/datastore/TraceUtil</className>
+        <differenceType>8001</differenceType>
+    </difference>
 </differences>

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

@@ -25,6 +25,7 @@
 import com.google.cloud.ServiceOptions;
 import com.google.cloud.datastore.execution.AggregationQueryExecutor;
 import com.google.cloud.datastore.spi.v1.DatastoreRpc;
+import com.google.cloud.datastore.telemetry.TraceUtil.Scope;
 import com.google.common.base.MoreObjects;
 import com.google.common.base.Preconditions;
 import com.google.common.base.Throwables;
@@ -39,9 +40,6 @@
 import com.google.datastore.v1.RunQueryResponse;
 import com.google.datastore.v1.TransactionOptions;
 import com.google.protobuf.ByteString;
-import io.opencensus.common.Scope;
-import io.opencensus.trace.Span;
-import io.opencensus.trace.Status;
 import io.opentelemetry.api.common.Attributes;
 import io.opentelemetry.api.trace.SpanBuilder;
 import io.opentelemetry.api.trace.SpanKind;
@@ -68,7 +66,6 @@ final class DatastoreImpl extends BaseService<DatastoreOptions> implements Datas
       TransactionExceptionHandler.build();
   private static final ExceptionHandler TRANSACTION_OPERATION_EXCEPTION_HANDLER =
       TransactionOperationExceptionHandler.build();
-  private final TraceUtil traceUtil = TraceUtil.getInstance();
 
   private final com.google.cloud.datastore.telemetry.TraceUtil otelTraceUtil =
       getOptions().getTraceUtil();
@@ -85,7 +82,8 @@ final class DatastoreImpl extends BaseService<DatastoreOptions> implements Datas
     readOptionProtoPreparer = new ReadOptionProtoPreparer();
     aggregationQueryExecutor =
         new AggregationQueryExecutor(
-            new RetryAndTraceDatastoreRpcDecorator(datastoreRpc, traceUtil, retrySettings, options),
+            new RetryAndTraceDatastoreRpcDecorator(
+                datastoreRpc, otelTraceUtil, retrySettings, options),
             options);
   }
 
@@ -744,8 +742,9 @@ void rollbackTransaction(ByteString transaction) {
   }
 
   void rollback(final com.google.datastore.v1.RollbackRequest requestPb) {
-    Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_ROLLBACK);
-    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_ROLLBACK);
+    try (Scope scope = span.makeCurrent()) {
       RetryHelper.runWithRetries(
           new Callable<Void>() {
             @Override
@@ -758,10 +757,10 @@ public Void call() throws DatastoreException {
           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();
     }
   }
 }

google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecorator.java

@@ -22,6 +22,7 @@
 import com.google.cloud.RetryHelper;
 import com.google.cloud.RetryHelper.RetryHelperException;
 import com.google.cloud.datastore.spi.v1.DatastoreRpc;
+import com.google.cloud.datastore.telemetry.TraceUtil;
 import com.google.datastore.v1.AllocateIdsRequest;
 import com.google.datastore.v1.AllocateIdsResponse;
 import com.google.datastore.v1.BeginTransactionRequest;
@@ -55,13 +56,13 @@ public class RetryAndTraceDatastoreRpcDecorator implements DatastoreRpc {
 
   public RetryAndTraceDatastoreRpcDecorator(
       DatastoreRpc datastoreRpc,
-      TraceUtil traceUtil,
+      TraceUtil otelTraceUtil,
       RetrySettings retrySettings,
       DatastoreOptions datastoreOptions) {
     this.datastoreRpc = datastoreRpc;
     this.retrySettings = retrySettings;
     this.datastoreOptions = datastoreOptions;
-    this.otelTraceUtil = datastoreOptions.getTraceUtil();
+    this.otelTraceUtil = otelTraceUtil;
   }
 
   @Override

google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/HttpDatastoreRpc.java

@@ -21,7 +21,6 @@
 import com.google.api.client.http.HttpTransport;
 import com.google.cloud.datastore.DatastoreException;
 import com.google.cloud.datastore.DatastoreOptions;
-import com.google.cloud.datastore.TraceUtil;
 import com.google.cloud.http.CensusHttpModule;
 import com.google.cloud.http.HttpTransportOptions;
 import com.google.datastore.v1.AllocateIdsRequest;
@@ -40,6 +39,7 @@
 import com.google.datastore.v1.RunAggregationQueryResponse;
 import com.google.datastore.v1.RunQueryRequest;
 import com.google.datastore.v1.RunQueryResponse;
+import io.opencensus.trace.Tracing;
 import java.io.IOException;
 import java.net.InetAddress;
 import java.net.URL;
@@ -80,8 +80,7 @@ public HttpDatastoreRpc(DatastoreOptions options) {
   private HttpRequestInitializer getHttpRequestInitializer(
       final DatastoreOptions options, HttpTransportOptions httpTransportOptions) {
     // Open Census initialization
-    CensusHttpModule censusHttpModule =
-        new CensusHttpModule(TraceUtil.getInstance().getTracer(), true);
+    CensusHttpModule censusHttpModule = new CensusHttpModule(Tracing.getTracer(), true);
     final HttpRequestInitializer censusHttpModuleHttpRequestInitializer =
         censusHttpModule.getHttpRequestInitializer(
             httpTransportOptions.getHttpRequestInitializer(options));

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

@@ -22,6 +22,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_RESERVE_IDS;
+import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_ROLLBACK;
 import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_RUN_AGGREGATION_QUERY;
 import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_RUN_QUERY;
 import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_COMMIT;
@@ -374,6 +375,7 @@ public void after() throws Exception {
     tracer = null;
     retrievedTrace = null;
     customSpanContext = null;
+    openTelemetrySdk = null;
   }
 
   @AfterClass
@@ -793,7 +795,7 @@ public void runAggregationQueryTraceTest() throws Exception {
   }
 
   @Test
-  public void transactionalLookupTest() throws Exception {
+  public void newTransactionReadTest() throws Exception {
     assertNotNull(customSpanContext);
 
     Span rootSpan = getNewRootSpanWithContext();
@@ -817,7 +819,7 @@ public void transactionalLookupTest() throws Exception {
   }
 
   @Test
-  public void transactionQueryTest() throws Exception {
+  public void newTransactionQueryTest() 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();
@@ -856,6 +858,116 @@ public void transactionQueryTest() throws Exception {
             Collections.singletonList(SPAN_NAME_TRANSACTION_COMMIT)));
   }
 
+  @Test
+  public void newTransactionReadWriteTraceTest() throws Exception {
+    // Set up
+    Entity entity1 = Entity.newBuilder(KEY1).set("pepper_type", "jalapeno").build();
+    Entity entity2 = Entity.newBuilder(KEY2).set("pepper_type", "habanero").build();
+    List<Entity> entityList = new ArrayList<>();
+    entityList.add(entity1);
+    entityList.add(entity2);
+
+    List<Entity> response = datastore.add(entity1, entity2);
+    assertEquals(entityList, response);
+
+    String simplified_spice_level = "not_spicy";
+    Entity entity1update =
+        Entity.newBuilder(entity1).set("spice_level", simplified_spice_level).build();
+
+    assertNotNull(customSpanContext);
+
+    // Test
+    Span rootSpan = getNewRootSpanWithContext();
+    try (Scope ignored = rootSpan.makeCurrent()) {
+      Transaction transaction = datastore.newTransaction();
+      entity1 = transaction.get(KEY1);
+      switch (entity1.getString("pepper_type")) {
+        case "jalapeno":
+          simplified_spice_level = "mild";
+          break;
+
+        case "jabanero":
+          simplified_spice_level = "hot";
+          break;
+      }
+      transaction.update(entity1update);
+      transaction.delete(KEY2);
+      transaction.commit();
+      assertFalse(transaction.isActive());
+    } finally {
+      rootSpan.end();
+    }
+
+    waitForTracesToComplete();
+
+    List<Entity> list = datastore.fetch(KEY1, KEY2);
+    assertEquals(list.get(0), entity1update);
+    assertNull(list.get(1));
+
+    fetchAndValidateTrace(
+        customSpanContext.getTraceId(),
+        /*numExpectedSpans=*/ 3,
+        Arrays.asList(
+            Collections.singletonList(SPAN_NAME_BEGIN_TRANSACTION),
+            Collections.singletonList(SPAN_NAME_TRANSACTION_LOOKUP),
+            Collections.singletonList(SPAN_NAME_TRANSACTION_COMMIT)));
+  }
+
+  @Test
+  public void newTransactionRollbackTest() throws Exception {
+    // Set up
+    Entity entity1 = Entity.newBuilder(KEY1).set("pepper_type", "jalapeno").build();
+    Entity entity2 = Entity.newBuilder(KEY2).set("pepper_type", "habanero").build();
+    List<Entity> entityList = new ArrayList<>();
+    entityList.add(entity1);
+    entityList.add(entity2);
+
+    List<Entity> response = datastore.add(entity1, entity2);
+    assertEquals(entityList, response);
+
+    String simplified_spice_level = "not_spicy";
+    Entity entity1update =
+        Entity.newBuilder(entity1).set("spice_level", simplified_spice_level).build();
+
+    assertNotNull(customSpanContext);
+
+    // Test
+    Span rootSpan = getNewRootSpanWithContext();
+    try (Scope ignored = rootSpan.makeCurrent()) {
+      Transaction transaction = datastore.newTransaction();
+      entity1 = transaction.get(KEY1);
+      switch (entity1.getString("pepper_type")) {
+        case "jalapeno":
+          simplified_spice_level = "mild";
+          break;
+
+        case "jabanero":
+          simplified_spice_level = "hot";
+          break;
+      }
+      transaction.update(entity1update);
+      transaction.delete(KEY2);
+      transaction.rollback();
+      assertFalse(transaction.isActive());
+    } finally {
+      rootSpan.end();
+    }
+
+    waitForTracesToComplete();
+
+    List<Entity> list = datastore.fetch(KEY1, KEY2);
+    assertEquals(list.get(0), entity1);
+    assertEquals(list.get(1), entity2);
+
+    fetchAndValidateTrace(
+        customSpanContext.getTraceId(),
+        /*numExpectedSpans=*/ 3,
+        Arrays.asList(
+            Collections.singletonList(SPAN_NAME_BEGIN_TRANSACTION),
+            Collections.singletonList(SPAN_NAME_TRANSACTION_LOOKUP),
+            Collections.singletonList(SPAN_NAME_ROLLBACK)));
+  }
+
   @Test
   public void runInTransactionQueryTest() throws Exception {
     // Set up
@@ -897,10 +1009,4 @@ public void runInTransactionQueryTest() throws Exception {
             Arrays.asList(SPAN_NAME_TRANSACTION_RUN, SPAN_NAME_RUN_QUERY),
             Arrays.asList(SPAN_NAME_TRANSACTION_RUN, SPAN_NAME_TRANSACTION_COMMIT)));
   }
-
-  @Test
-  public void transactionRunQueryTest() throws Exception {}
-
-  @Test
-  public void readWriteTransactionTraceTest() throws Exception {}
 }