From 9f2d397d3ab7496f64ee575db4a89819a1bd2fdc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= <koloite@gmail.com>
Date: Fri, 7 Jun 2024 12:20:13 +0200
Subject: [PATCH] fix: remove unused DmlBatch span

The DmlBatch span that was created for DML batches was never used,
as all traces were registered on the underlying transaction of the
batch. The DmlBatch span was therefore always just an empty
sibling span of the Transaction span, which did contain the traces
of the statements that were executed.
---
 .../cloud/spanner/connection/AbstractBaseUnitOfWork.java     | 5 +++++
 .../com/google/cloud/spanner/connection/ConnectionImpl.java  | 4 ++--
 .../java/com/google/cloud/spanner/connection/UnitOfWork.java | 4 ++++
 3 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractBaseUnitOfWork.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractBaseUnitOfWork.java
index ef624ea8416..ff159c681c5 100644
--- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractBaseUnitOfWork.java
+++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractBaseUnitOfWork.java
@@ -165,6 +165,11 @@ B setSpan(@Nullable Span span) {
     this.span = Preconditions.checkNotNull(builder.span);
   }
 
+  @Override
+  public Span getSpan() {
+    return this.span;
+  }
+
   ApiFuture<Void> asyncEndUnitOfWorkSpan() {
     return this.statementExecutor.submit(this::endUnitOfWorkSpan);
   }
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java
index 745188c17fb..4ab3f7fa868 100644
--- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java
+++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java
@@ -96,7 +96,6 @@ class ConnectionImpl implements Connection {
   private static final String SINGLE_USE_TRANSACTION = "SingleUseTransaction";
   private static final String READ_ONLY_TRANSACTION = "ReadOnlyTransaction";
   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";
 
@@ -1932,7 +1931,8 @@ UnitOfWork createNewUnitOfWork(
               .setStatementTag(statementTag)
               .setExcludeTxnFromChangeStreams(excludeTxnFromChangeStreams)
               .setRpcPriority(rpcPriority)
-              .setSpan(createSpanForUnitOfWork(DML_BATCH))
+              // Use the transaction Span for the DML batch.
+              .setSpan(transactionStack.peek().getSpan())
               .build();
         case DDL_BATCH:
           return DdlBatch.newBuilder()
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/UnitOfWork.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/UnitOfWork.java
index 577cd39d201..9a5e9197245 100644
--- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/UnitOfWork.java
+++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/UnitOfWork.java
@@ -31,6 +31,7 @@
 import com.google.cloud.spanner.TransactionContext;
 import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStatement;
 import com.google.spanner.v1.ResultSetStats;
+import io.opentelemetry.api.trace.Span;
 import java.util.concurrent.ExecutionException;
 import javax.annotation.Nonnull;
 
@@ -77,6 +78,9 @@ public boolean isActive() {
   /** @return <code>true</code> if this unit of work is still active. */
   boolean isActive();
 
+  /** @return the {@link Span} that is used by this {@link UnitOfWork}. */
+  Span getSpan();
+
   /** Returns true if this transaction can only be used for a single statement. */
   boolean isSingleUse();