test: add tracing tests

google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java

@@ -748,7 +748,9 @@ private ResultSet internalAnalyzeStatement(
 
     @Override
     public long executeUpdate(Statement statement, UpdateOption... options) {
-      ISpan span = tracer.spanBuilder(UPDATE, Attributes.of(DB_STATEMENT_KEY, statement.getSql()));
+      ISpan span =
+          tracer.spanBuilderWithExplicitParent(
+              UPDATE, this.span, Attributes.of(DB_STATEMENT_KEY, statement.getSql()));
       try (IScope ignore = tracer.withSpan(span)) {
         ResultSet resultSet = internalExecuteUpdate(statement, QueryMode.NORMAL, options);
         // For standard DML, using the exact row count.
@@ -790,7 +792,9 @@ private ResultSet internalExecuteUpdate(
 
     @Override
     public ApiFuture<Long> executeUpdateAsync(Statement statement, UpdateOption... updateOptions) {
-      ISpan span = tracer.spanBuilder(UPDATE, Attributes.of(DB_STATEMENT_KEY, statement.getSql()));
+      ISpan span =
+          tracer.spanBuilderWithExplicitParent(
+              UPDATE, this.span, Attributes.of(DB_STATEMENT_KEY, statement.getSql()));
       try (IScope ignore = tracer.withSpan(span)) {
         beforeReadOrQuery();
         final Options options = Options.fromUpdateOptions(updateOptions);
@@ -880,8 +884,9 @@ private SpannerException createAbortedExceptionForBatchDml(ExecuteBatchDmlRespon
     @Override
     public long[] batchUpdate(Iterable<Statement> statements, UpdateOption... updateOptions) {
       ISpan span =
-          tracer.spanBuilder(
+          tracer.spanBuilderWithExplicitParent(
               BATCH_UPDATE,
+              this.span,
               Attributes.of(
                   DB_STATEMENT_ARRAY_KEY,
                   StreamSupport.stream(statements.spliterator(), false)
@@ -937,8 +942,9 @@ public long[] batchUpdate(Iterable<Statement> statements, UpdateOption... update
     public ApiFuture<long[]> batchUpdateAsync(
         Iterable<Statement> statements, UpdateOption... updateOptions) {
       ISpan span =
-          tracer.spanBuilder(
+          tracer.spanBuilderWithExplicitParent(
               BATCH_UPDATE,
+              this.span,
               Attributes.of(
                   DB_STATEMENT_ARRAY_KEY,
                   StreamSupport.stream(statements.spliterator(), false)

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractBaseUnitOfWork.java

@@ -16,7 +16,6 @@
 
 package com.google.cloud.spanner.connection;
 
-
 import com.google.api.core.ApiFuture;
 import com.google.api.core.ApiFutures;
 import com.google.api.gax.grpc.GrpcCallContext;

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java

@@ -93,6 +93,7 @@ class ConnectionImpl implements Connection {
   private static final String READ_WRITE_TRANSACTION = "ReadWriteTransaction";
   private static final String DML_BATCH = "DmlBatch";
   private static final String DDL_BATCH = "DdlBatch";
+  private static final String DDL_STATEMENT = "DdlStatement";
 
   private static final String CLOSED_ERROR_MSG = "This connection is closed";
   private static final String ONLY_ALLOWED_IN_AUTOCOMMIT =
@@ -1734,7 +1735,8 @@ UnitOfWork getCurrentUnitOfWorkOrStartNewUnitOfWork(
               /* isInternalMetadataQuery = */ false,
               /* forceSingleUse = */ statementType == StatementType.DDL
                   && this.ddlInTransactionMode != DdlInTransactionMode.FAIL
-                  && !this.transactionBeginMarked);
+                  && !this.transactionBeginMarked,
+              statementType);
     }
     return this.currentUnitOfWork;
   }
@@ -1757,6 +1759,12 @@ void maybeAutoCommitCurrentTransaction(StatementType statementType) {
 
   @VisibleForTesting
   UnitOfWork createNewUnitOfWork(boolean isInternalMetadataQuery, boolean forceSingleUse) {
+    return createNewUnitOfWork(isInternalMetadataQuery, forceSingleUse, null);
+  }
+
+  @VisibleForTesting
+  UnitOfWork createNewUnitOfWork(
+      boolean isInternalMetadataQuery, boolean forceSingleUse, StatementType statementType) {
     if (isInternalMetadataQuery
         || (isAutocommit() && !isInTransaction() && !isInBatch())
         || forceSingleUse) {
@@ -1772,7 +1780,9 @@ UnitOfWork createNewUnitOfWork(boolean isInternalMetadataQuery, boolean forceSin
           .setMaxCommitDelay(maxCommitDelay)
           .setStatementTimeout(statementTimeout)
           .withStatementExecutor(statementExecutor)
-          .setSpan(createSpanForUnitOfWork(SINGLE_USE_TRANSACTION))
+          .setSpan(
+              createSpanForUnitOfWork(
+                  statementType == StatementType.DDL ? DDL_STATEMENT : SINGLE_USE_TRANSACTION))
           .build();
     } else {
       switch (getUnitOfWorkType()) {

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SingleUseTransaction.java

@@ -385,6 +385,7 @@ public ApiFuture<Void> executeDdlAsync(CallType callType, final ParsedStatement
         !isReadOnly(), "DDL statements are not allowed in read-only mode");
     try (Scope ignore = span.makeCurrent()) {
       checkAndMarkUsed();
+      span.setAttribute(DB_STATEMENT_KEY, ddl.getStatement().getSql());
 
       Callable<Void> callable =
           () -> {

google-cloud-spanner/src/test/java/com/google/cloud/spanner/OpenTelemetrySpanTest.java

@@ -424,11 +424,13 @@ public void transactionRunner() {
             ? ImmutableList.of(
                 "CloudSpannerOperation.CreateMultiplexedSession",
                 "CloudSpannerOperation.BatchCreateSessionsRequest",
+                "CloudSpannerOperation.ExecuteUpdate",
                 "CloudSpannerOperation.Commit",
                 "CloudSpannerOperation.BatchCreateSessions",
                 "CloudSpanner.ReadWriteTransaction")
             : ImmutableList.of(
                 "CloudSpannerOperation.BatchCreateSessionsRequest",
+                "CloudSpannerOperation.ExecuteUpdate",
                 "CloudSpannerOperation.Commit",
                 "CloudSpannerOperation.BatchCreateSessions",
                 "CloudSpanner.ReadWriteTransaction");
@@ -460,6 +462,7 @@ public void transactionRunner() {
                       expectedBatchCreateSessionsEventsCount);
                   break;
                 case "CloudSpannerOperation.Commit":
+                case "CloudSpannerOperation.ExecuteUpdate":
                   assertEquals(0, spanItem.getEvents().size());
                   break;
                 case "CloudSpanner.ReadWriteTransaction":
@@ -484,10 +487,12 @@ public void transactionRunnerWithError() {
                 "CloudSpannerOperation.CreateMultiplexedSession",
                 "CloudSpannerOperation.BatchCreateSessionsRequest",
                 "CloudSpannerOperation.BatchCreateSessions",
+                "CloudSpannerOperation.ExecuteUpdate",
                 "CloudSpanner.ReadWriteTransaction")
             : ImmutableList.of(
                 "CloudSpannerOperation.BatchCreateSessionsRequest",
                 "CloudSpannerOperation.BatchCreateSessions",
+                "CloudSpannerOperation.ExecuteUpdate",
                 "CloudSpanner.ReadWriteTransaction");
     TransactionRunner runner = client.readWriteTransaction();
     SpannerException e =
@@ -527,6 +532,9 @@ public void transactionRunnerWithError() {
                       expectedReadWriteTransactionErrorEvents,
                       expectedReadWriteTransactionErrorEventsCount);
                   break;
+                case "CloudSpannerOperation.ExecuteUpdate":
+                  assertEquals(0, spanItem.getEvents().size());
+                  break;
                 default:
                   assert false;
               }
@@ -543,12 +551,14 @@ public void transactionRunnerWithFailedAndBeginTransaction() {
                 "CloudSpannerOperation.CreateMultiplexedSession",
                 "CloudSpannerOperation.BeginTransaction",
                 "CloudSpannerOperation.BatchCreateSessionsRequest",
+                "CloudSpannerOperation.ExecuteUpdate",
                 "CloudSpannerOperation.Commit",
                 "CloudSpannerOperation.BatchCreateSessions",
                 "CloudSpanner.ReadWriteTransaction")
             : ImmutableList.of(
                 "CloudSpannerOperation.BeginTransaction",
                 "CloudSpannerOperation.BatchCreateSessionsRequest",
+                "CloudSpannerOperation.ExecuteUpdate",
                 "CloudSpannerOperation.Commit",
                 "CloudSpannerOperation.BatchCreateSessions",
                 "CloudSpanner.ReadWriteTransaction");
@@ -606,6 +616,7 @@ public void transactionRunnerWithFailedAndBeginTransaction() {
                   break;
                 case "CloudSpannerOperation.Commit":
                 case "CloudSpannerOperation.BeginTransaction":
+                case "CloudSpannerOperation.ExecuteUpdate":
                   assertEquals(0, spanItem.getEvents().size());
                   break;
                 case "CloudSpanner.ReadWriteTransaction":

google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpanTest.java

@@ -506,12 +506,13 @@ public void transactionRunnerWithError() {
     Map<String, Boolean> spans = failOnOverkillTraceComponent.getSpans();
 
     if (isMultiplexedSessionsEnabled()) {
-      assertEquals(spans.toString(), 4, spans.size());
+      assertEquals(spans.toString(), 5, spans.size());
       assertThat(spans).containsEntry("CloudSpannerOperation.CreateMultiplexedSession", true);
     } else {
-      assertThat(spans.size()).isEqualTo(3);
+      assertThat(spans.size()).isEqualTo(4);
     }
     assertThat(spans).containsEntry("CloudSpanner.ReadWriteTransaction", true);
+    assertThat(spans).containsEntry("CloudSpannerOperation.ExecuteUpdate", true);
     assertThat(spans).containsEntry("CloudSpannerOperation.BatchCreateSessions", true);
     assertThat(spans).containsEntry("CloudSpannerOperation.BatchCreateSessionsRequest", true);
 

google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerOptionsTestHelper.java

@@ -0,0 +1,23 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.cloud.spanner;
+
+public class SpannerOptionsTestHelper {
+
+  public static void resetActiveTracingFramework() {
+    SpannerOptions.resetActiveTracingFramework();
+  }
+}

google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionContextImplTest.java

@@ -17,6 +17,7 @@
 package com.google.cloud.spanner;
 
 import static org.junit.Assert.assertThrows;
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.anyMap;
 import static org.mockito.Mockito.doNothing;
@@ -34,6 +35,7 @@
 import com.google.spanner.v1.CommitRequest;
 import com.google.spanner.v1.ExecuteBatchDmlRequest;
 import com.google.spanner.v1.ExecuteBatchDmlResponse;
+import io.opentelemetry.api.common.Attributes;
 import java.util.Collections;
 import org.junit.Before;
 import org.junit.Test;
@@ -68,6 +70,12 @@ public void setup() {
     doNothing().when(span).end();
     doNothing().when(span).addAnnotation("Starting Commit");
     when(tracer.spanBuilderWithExplicitParent(SpannerImpl.COMMIT, span)).thenReturn(span);
+    when(tracer.spanBuilderWithExplicitParent(
+            eq(SpannerImpl.COMMIT), eq(span), any(Attributes.class)))
+        .thenReturn(span);
+    when(tracer.spanBuilderWithExplicitParent(
+            eq(SpannerImpl.BATCH_UPDATE), eq(span), any(Attributes.class)))
+        .thenReturn(span);
   }
 
   private TransactionContextImpl createContext() {
@@ -187,7 +195,7 @@ public void testReturnCommitStats() {
               .setSession(session.getName())
               .setTransactionId(transactionId)
               .build();
-      verify(rpc).commitAsync(Mockito.eq(request), anyMap());
+      verify(rpc).commitAsync(eq(request), anyMap());
     }
   }
 
@@ -210,6 +218,8 @@ private void batchDml(int status) {
             .setRpc(rpc)
             .setTransactionId(ByteString.copyFromUtf8("test"))
             .setOptions(Options.fromTransactionOptions())
+            .setTracer(tracer)
+            .setSpan(span)
             .build()) {
       impl.batchUpdate(Collections.singletonList(statement));
     }

google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionImplTest.java

@@ -69,6 +69,7 @@
 import com.google.cloud.spanner.TransactionRunner;
 import com.google.cloud.spanner.Type;
 import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStatement;
+import com.google.cloud.spanner.connection.AbstractStatementParser.StatementType;
 import com.google.cloud.spanner.connection.ConnectionImpl.UnitOfWorkType;
 import com.google.cloud.spanner.connection.ConnectionStatementExecutorImpl.StatementTimeoutGetter;
 import com.google.cloud.spanner.connection.ReadOnlyStalenessUtil.GetExactStaleness;
@@ -1751,7 +1752,10 @@ public void testTransactionTagNotAllowedAfterTransactionStarted() {
         new ConnectionImpl(
             connectionOptions, spannerPool, ddlClient, dbClient, mock(BatchClient.class)) {
           @Override
-          UnitOfWork createNewUnitOfWork(boolean isInternalMetadataQuery, boolean forceSingleUse) {
+          UnitOfWork createNewUnitOfWork(
+              boolean isInternalMetadataQuery,
+              boolean forceSingleUse,
+              StatementType statementType) {
             return unitOfWork;
           }
         }) {

google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/DdlBatchTest.java

@@ -51,6 +51,7 @@
 import com.google.protobuf.Timestamp;
 import com.google.spanner.admin.database.v1.UpdateDatabaseDdlMetadata;
 import io.grpc.Status;
+import io.opentelemetry.api.trace.Span;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
@@ -128,6 +129,7 @@ private DdlBatch createSubject(DdlClient ddlClient, DatabaseClient dbClient) {
         .setDdlClient(ddlClient)
         .setDatabaseClient(dbClient)
         .withStatementExecutor(new StatementExecutor())
+        .setSpan(Span.getInvalid())
         .build();
   }
 
@@ -387,6 +389,7 @@ public void testUpdateCount() throws InterruptedException, ExecutionException {
             .withStatementExecutor(new StatementExecutor())
             .setDdlClient(client)
             .setDatabaseClient(mock(DatabaseClient.class))
+            .setSpan(Span.getInvalid())
             .build();
     batch.executeDdlAsync(
         CallType.SYNC,
@@ -425,6 +428,7 @@ public void testFailedUpdateCount() throws InterruptedException, ExecutionExcept
             .withStatementExecutor(new StatementExecutor())
             .setDdlClient(client)
             .setDatabaseClient(mock(DatabaseClient.class))
+            .setSpan(Span.getInvalid())
             .build();
     batch.executeDdlAsync(
         CallType.SYNC,
@@ -467,6 +471,7 @@ public void testFailedAfterFirstStatement() throws InterruptedException, Executi
             .withStatementExecutor(new StatementExecutor())
             .setDdlClient(client)
             .setDatabaseClient(mock(DatabaseClient.class))
+            .setSpan(Span.getInvalid())
             .build();
     batch.executeDdlAsync(
         CallType.SYNC,

google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/DmlBatchTest.java

@@ -38,6 +38,7 @@
 import com.google.cloud.spanner.connection.AbstractStatementParser.StatementType;
 import com.google.cloud.spanner.connection.UnitOfWork.CallType;
 import com.google.cloud.spanner.connection.UnitOfWork.UnitOfWorkState;
+import io.opentelemetry.api.trace.Span;
 import java.util.Arrays;
 import java.util.Collections;
 import org.junit.Test;
@@ -65,6 +66,7 @@ private DmlBatch createSubject(UnitOfWork transaction) {
     return DmlBatch.newBuilder()
         .setTransaction(transaction)
         .withStatementExecutor(new StatementExecutor())
+        .setSpan(Span.getInvalid())
         .build();
   }