feat: allow DML batches in transactions to execute analyzeUpdate (#3114)

olavloite · web-flow · commit dee7cdabe740 · 2024-05-23T08:14:02.000+02:00
Executing analyzeUpdate in a DML Batch inside a transaction should
be possible, as the parent transaction can be used for that.
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
@@ -176,8 +176,6 @@ private Void endUnitOfWorkSpan() {
     return null;
   }
 
-  abstract boolean isSingleUse();
-
   /**
    * Returns a descriptive name for the type of transaction / unit of work. This is used in error
    * messages.
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractMultiUseTransaction.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractMultiUseTransaction.java
@@ -96,7 +96,7 @@ public String toString() {
   }
 
   @Override
-  boolean isSingleUse() {
+  public boolean isSingleUse() {
     return false;
   }
 
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DdlBatch.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DdlBatch.java
@@ -97,7 +97,7 @@ private DdlBatch(Builder builder) {
   }
 
   @Override
-  boolean isSingleUse() {
+  public boolean isSingleUse() {
     return false;
   }
 
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DmlBatch.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DmlBatch.java
@@ -78,12 +78,12 @@ static Builder newBuilder() {
 
   private DmlBatch(Builder builder) {
     super(builder);
-    this.transaction = builder.transaction;
+    this.transaction = Preconditions.checkNotNull(builder.transaction);
     this.statementTag = builder.statementTag;
   }
 
   @Override
-  boolean isSingleUse() {
+  public boolean isSingleUse() {
     return false;
   }
 
@@ -174,8 +174,11 @@ public ApiFuture<Long> executeUpdateAsync(
   @Override
   public ApiFuture<ResultSet> analyzeUpdateAsync(
       CallType callType, ParsedStatement update, AnalyzeMode analyzeMode, UpdateOption... options) {
-    throw SpannerExceptionFactory.newSpannerException(
-        ErrorCode.FAILED_PRECONDITION, "Analyzing updates is not allowed for DML batches.");
+    if (transaction.isSingleUse()) {
+      throw SpannerExceptionFactory.newSpannerException(
+          ErrorCode.FAILED_PRECONDITION, "Analyzing updates is not allowed for DML batches.");
+    }
+    return transaction.analyzeUpdateAsync(callType, update, analyzeMode, options);
   }
 
   @Override
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SingleUseTransaction.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SingleUseTransaction.java
@@ -180,7 +180,7 @@ private SingleUseTransaction(Builder builder) {
   }
 
   @Override
-  boolean isSingleUse() {
+  public boolean isSingleUse() {
     return true;
   }
 
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
@@ -77,6 +77,9 @@ public boolean isActive() {
   /** @return <code>true</code> if this unit of work is still active. */
   boolean isActive();
 
+  /** Returns true if this transaction can only be used for a single statement. */
+  boolean isSingleUse();
+
   /**
    * Commits the changes in this unit of work to the database. For read-only transactions, this only
    * closes the {@link ReadContext}. This method will throw a {@link SpannerException} if called for
diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AnalyzeStatementsTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AnalyzeStatementsTest.java
@@ -355,8 +355,9 @@ public void testAnalyzeUpdateStatementDdlBatch() {
   }
 
   @Test
-  public void testAnalyzeUpdateDmlBatch() {
+  public void testAnalyzeUpdateDmlBatch_AutoCommit() {
     try (Connection connection = createConnection()) {
+      connection.setAutocommit(true);
       connection.startBatchDml();
 
       SpannerException exception =
@@ -371,8 +372,23 @@ public void testAnalyzeUpdateDmlBatch() {
   }
 
   @Test
-  public void testAnalyzeUpdateStatementDmlBatch() {
+  public void testAnalyzeUpdateDmlBatch_Transactional() {
     try (Connection connection = createConnection()) {
+      connection.setAutocommit(false);
+      connection.startBatchDml();
+
+      assertNotNull(connection.analyzeUpdate(PLAN_UPDATE, QueryAnalyzeMode.PLAN));
+      assertEquals(-1L, connection.executeUpdate(INSERT_STATEMENT));
+      connection.runBatch();
+
+      assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
+    }
+  }
+
+  @Test
+  public void testAnalyzeUpdateStatementDmlBatch_AutoCommit() {
+    try (Connection connection = createConnection()) {
+      connection.setAutocommit(true);
       connection.startBatchDml();
 
       SpannerException exception =
@@ -385,4 +401,18 @@ public void testAnalyzeUpdateStatementDmlBatch() {
     assertEquals(0, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
     assertEquals(0, mockSpanner.countRequestsOfType(CommitRequest.class));
   }
+
+  @Test
+  public void testAnalyzeUpdateStatementDmlBatch_Transactional() {
+    try (Connection connection = createConnection()) {
+      connection.setAutocommit(false);
+      connection.startBatchDml();
+
+      connection.analyzeUpdateStatement(PLAN_UPDATE, QueryAnalyzeMode.PLAN);
+      assertEquals(-1L, connection.executeUpdate(INSERT_STATEMENT));
+      connection.runBatch();
+
+      assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
+    }
+  }
 }

Original file line number	Diff line number	Diff line change
`@@ -176,8 +176,6 @@ private Void endUnitOfWorkSpan() {`
`176`	`176`	`return null;`
`177`	`177`	`}`
`178`	`178`
`179`		`- abstract boolean isSingleUse();`
`180`		`-`
`181`	`179`	`/**`
`182`	`180`	`* Returns a descriptive name for the type of transaction / unit of work. This is used in error`
`183`	`181`	`* messages.`
Original file line number	Diff line number	Diff line change
`@@ -96,7 +96,7 @@ public String toString() {`
`96`	`96`	`}`
`97`	`97`
`98`	`98`	`@Override`
`99`		`- boolean isSingleUse() {`
	`99`	`+ public boolean isSingleUse() {`
`100`	`100`	`return false;`
`101`	`101`	`}`
`102`	`102`
Original file line number	Diff line number	Diff line change
`@@ -97,7 +97,7 @@ private DdlBatch(Builder builder) {`
`97`	`97`	`}`
`98`	`98`
`99`	`99`	`@Override`
`100`		`- boolean isSingleUse() {`
	`100`	`+ public boolean isSingleUse() {`
`101`	`101`	`return false;`
`102`	`102`	`}`
`103`	`103`
Original file line number	Diff line number	Diff line change
`@@ -180,7 +180,7 @@ private SingleUseTransaction(Builder builder) {`
`180`	`180`	`}`
`181`	`181`
`182`	`182`	`@Override`
`183`		`- boolean isSingleUse() {`
	`183`	`+ public boolean isSingleUse() {`
`184`	`184`	`return true;`
`185`	`185`	`}`
`186`	`186`