chore: Support dml_batch_update_count option in PG dialect (#4142)

sakthivelmanii · cloud-java-bot · web-flow · commit ece13906ae06 · 2025-10-10T13:44:46.000+02:00
* chore: Support dml_batch_update_count option in PG dialect

* chore: generate libraries at Wed Oct  8 21:30:02 UTC 2025

* Support in GoogleSQL dialect

* Address comments

* Fix tests

* Update connection property name

* Address comments

---------

Co-authored-by: cloud-java-bot &lt;cloud-java-bot@google.com&gt;
diff --git a/.github/workflows/update_generation_config.yaml b/.github/workflows/update_generation_config.yaml
@@ -26,7 +26,7 @@ jobs:
       # the branch into which the pull request is merged
       base_branch: main
     steps:
-    - uses: actions/checkout@v5
+    - uses: actions/checkout@v4
       with:
         fetch-depth: 0
         token: ${{ secrets.CLOUD_JAVA_BOT_TOKEN }}
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
@@ -25,6 +25,7 @@
 import static com.google.cloud.spanner.connection.ConnectionProperties.AUTO_BATCH_DML_UPDATE_COUNT;
 import static com.google.cloud.spanner.connection.ConnectionProperties.AUTO_BATCH_DML_UPDATE_COUNT_VERIFICATION;
 import static com.google.cloud.spanner.connection.ConnectionProperties.AUTO_PARTITION_MODE;
+import static com.google.cloud.spanner.connection.ConnectionProperties.BATCH_DML_UPDATE_COUNT;
 import static com.google.cloud.spanner.connection.ConnectionProperties.DATA_BOOST_ENABLED;
 import static com.google.cloud.spanner.connection.ConnectionProperties.DDL_IN_TRANSACTION_MODE;
 import static com.google.cloud.spanner.connection.ConnectionProperties.DEFAULT_ISOLATION_LEVEL;
@@ -1607,6 +1608,10 @@ public long getAutoBatchDmlUpdateCount() {
     return getConnectionPropertyValue(AUTO_BATCH_DML_UPDATE_COUNT);
   }
 
+  long getDmlBatchUpdateCount() {
+    return getConnectionPropertyValue(BATCH_DML_UPDATE_COUNT);
+  }
+
   @Override
   public void setAutoBatchDmlUpdateCountVerification(boolean verification) {
     setConnectionPropertyValue(AUTO_BATCH_DML_UPDATE_COUNT_VERIFICATION, verification);
@@ -1617,6 +1622,10 @@ public boolean isAutoBatchDmlUpdateCountVerification() {
     return getConnectionPropertyValue(AUTO_BATCH_DML_UPDATE_COUNT_VERIFICATION);
   }
 
+  void setBatchDmlUpdateCount(long updateCount, boolean local) {
+    setConnectionPropertyValue(BATCH_DML_UPDATE_COUNT, updateCount, local);
+  }
+
   @Override
   public void setDataBoostEnabled(boolean dataBoostEnabled) {
     setConnectionPropertyValue(DATA_BOOST_ENABLED, dataBoostEnabled);
@@ -2340,6 +2349,7 @@ UnitOfWork createNewUnitOfWork(
               .setAutoBatchUpdateCountSupplier(this::getAutoBatchDmlUpdateCount)
               .setAutoBatchUpdateCountVerificationSupplier(
                   this::isAutoBatchDmlUpdateCountVerification)
+              .setDmlBatchUpdateCountSupplier(this::getDmlBatchUpdateCount)
               .setTransaction(currentUnitOfWork)
               .setStatementTimeout(statementTimeout)
               .withStatementExecutor(statementExecutor)
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionOptions.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionOptions.java
@@ -181,6 +181,7 @@ public class ConnectionOptions {
   static final boolean DEFAULT_ENABLE_END_TO_END_TRACING = false;
   static final boolean DEFAULT_AUTO_BATCH_DML = false;
   static final long DEFAULT_AUTO_BATCH_DML_UPDATE_COUNT = 1L;
+  static final long DEFAULT_BATCH_DML_UPDATE_COUNT = -1L;
   static final boolean DEFAULT_AUTO_BATCH_DML_UPDATE_COUNT_VERIFICATION = true;
   private static final String EXPERIMENTAL_HOST_PROJECT_ID = "default";
   private static final String DEFAULT_EXPERIMENTAL_HOST_INSTANCE_ID = "default";
@@ -314,6 +315,7 @@ public class ConnectionOptions {
       "auto_batch_dml_update_count";
   public static final String AUTO_BATCH_DML_UPDATE_COUNT_VERIFICATION_PROPERTY_NAME =
       "auto_batch_dml_update_count_verification";
+  public static final String BATCH_DML_UPDATE_COUNT_PROPERTY_NAME = "batch_dml_update_count";
 
   private static final String GUARDED_CONNECTION_PROPERTY_ERROR_MESSAGE =
       "%s can only be used if the system property %s has been set to true. "
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionProperties.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionProperties.java
@@ -21,6 +21,7 @@
 import static com.google.cloud.spanner.connection.ConnectionOptions.AUTO_BATCH_DML_UPDATE_COUNT_PROPERTY_NAME;
 import static com.google.cloud.spanner.connection.ConnectionOptions.AUTO_BATCH_DML_UPDATE_COUNT_VERIFICATION_PROPERTY_NAME;
 import static com.google.cloud.spanner.connection.ConnectionOptions.AUTO_PARTITION_MODE_PROPERTY_NAME;
+import static com.google.cloud.spanner.connection.ConnectionOptions.BATCH_DML_UPDATE_COUNT_PROPERTY_NAME;
 import static com.google.cloud.spanner.connection.ConnectionOptions.CHANNEL_PROVIDER_PROPERTY_NAME;
 import static com.google.cloud.spanner.connection.ConnectionOptions.CLIENT_CERTIFICATE_PROPERTY_NAME;
 import static com.google.cloud.spanner.connection.ConnectionOptions.CLIENT_KEY_PROPERTY_NAME;
@@ -34,6 +35,7 @@
 import static com.google.cloud.spanner.connection.ConnectionOptions.DEFAULT_AUTO_BATCH_DML_UPDATE_COUNT;
 import static com.google.cloud.spanner.connection.ConnectionOptions.DEFAULT_AUTO_BATCH_DML_UPDATE_COUNT_VERIFICATION;
 import static com.google.cloud.spanner.connection.ConnectionOptions.DEFAULT_AUTO_PARTITION_MODE;
+import static com.google.cloud.spanner.connection.ConnectionOptions.DEFAULT_BATCH_DML_UPDATE_COUNT;
 import static com.google.cloud.spanner.connection.ConnectionOptions.DEFAULT_CHANNEL_PROVIDER;
 import static com.google.cloud.spanner.connection.ConnectionOptions.DEFAULT_CLIENT_CERTIFICATE;
 import static com.google.cloud.spanner.connection.ConnectionOptions.DEFAULT_CLIENT_KEY;
@@ -747,6 +749,15 @@ public class ConnectionProperties {
           BOOLEANS,
           BooleanConverter.INSTANCE,
           Context.USER);
+  static final ConnectionProperty<Long> BATCH_DML_UPDATE_COUNT =
+      create(
+          BATCH_DML_UPDATE_COUNT_PROPERTY_NAME,
+          "The update count that is returned for DML statements that are executed in an "
+              + "explicit DML batch. The default is "
+              + DEFAULT_BATCH_DML_UPDATE_COUNT,
+          DEFAULT_BATCH_DML_UPDATE_COUNT,
+          LongConverter.INSTANCE,
+          Context.USER);
 
   static final ImmutableMap<String, ConnectionProperty<?>> CONNECTION_PROPERTIES =
       CONNECTION_PROPERTIES_BUILDER.build();
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionStatementExecutor.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionStatementExecutor.java
@@ -184,6 +184,8 @@ StatementResult statementSetPgSessionCharacteristicsTransactionMode(
 
   StatementResult statementRunPartitionedQuery(Statement statement);
 
+  StatementResult statementSetBatchDmlUpdateCount(Long updateCount, Boolean local);
+
   StatementResult statementSetAutoBatchDml(Boolean autoBatchDml);
 
   StatementResult statementShowAutoBatchDml();
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionStatementExecutorImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionStatementExecutorImpl.java
@@ -29,6 +29,7 @@
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SET_AUTO_BATCH_DML_UPDATE_COUNT;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SET_AUTO_BATCH_DML_UPDATE_COUNT_VERIFICATION;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SET_AUTO_PARTITION_MODE;
+import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SET_BATCH_DML_UPDATE_COUNT;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SET_DATA_BOOST_ENABLED;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SET_DEFAULT_TRANSACTION_ISOLATION;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SET_DELAY_TRANSACTION_START_UNTIL_FIRST_WRITE;
@@ -726,6 +727,12 @@ public StatementResult statementRunPartitionedQuery(Statement statement) {
         ClientSideStatementType.RUN_PARTITIONED_QUERY);
   }
 
+  @Override
+  public StatementResult statementSetBatchDmlUpdateCount(Long updateCount, Boolean local) {
+    getConnection().setBatchDmlUpdateCount(updateCount, local);
+    return noResult(SET_BATCH_DML_UPDATE_COUNT);
+  }
+
   @Override
   public StatementResult statementSetProtoDescriptors(byte[] protoDescriptors) {
     Preconditions.checkNotNull(protoDescriptors);
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
@@ -51,6 +51,7 @@ class DmlBatch extends AbstractBaseUnitOfWork {
   private final boolean autoBatch;
   private final Supplier<Long> autoBatchUpdateCountSupplier;
   private final Supplier<Boolean> verifyUpdateCountsSupplier;
+  private final Supplier<Long> dmlbatchUpdateCountSupplier;
   private final UnitOfWork transaction;
   private final String statementTag;
   private final List<ParsedStatement> statements = new ArrayList<>();
@@ -61,6 +62,7 @@ static class Builder extends AbstractBaseUnitOfWork.Builder<Builder, DmlBatch> {
     private boolean autoBatch;
     private Supplier<Long> autoBatchUpdateCountSupplier = Suppliers.ofInstance(1L);
     private Supplier<Boolean> verifyUpdateCountsSupplier = Suppliers.ofInstance(Boolean.FALSE);
+    private Supplier<Long> dmlbatchUpdateCountSupplier = Suppliers.ofInstance(-1L);
     private UnitOfWork transaction;
     private String statementTag;
 
@@ -81,6 +83,12 @@ Builder setAutoBatchUpdateCountVerificationSupplier(Supplier<Boolean> verificati
       return this;
     }
 
+    Builder setDmlBatchUpdateCountSupplier(Supplier<Long> dmlbatchUpdateCountSupplier) {
+      Preconditions.checkNotNull(dmlbatchUpdateCountSupplier);
+      this.dmlbatchUpdateCountSupplier = dmlbatchUpdateCountSupplier;
+      return this;
+    }
+
     Builder setTransaction(UnitOfWork transaction) {
       Preconditions.checkNotNull(transaction);
       this.transaction = transaction;
@@ -108,6 +116,7 @@ private DmlBatch(Builder builder) {
     this.autoBatch = builder.autoBatch;
     this.autoBatchUpdateCountSupplier = builder.autoBatchUpdateCountSupplier;
     this.verifyUpdateCountsSupplier = builder.verifyUpdateCountsSupplier;
+    this.dmlbatchUpdateCountSupplier = builder.dmlbatchUpdateCountSupplier;
     this.transaction = Preconditions.checkNotNull(builder.transaction);
     this.statementTag = builder.statementTag;
   }
@@ -193,7 +202,7 @@ public ApiFuture<Void> executeDdlAsync(CallType callType, ParsedStatement ddl) {
   long getUpdateCount() {
     // Auto-batching returns update count 1 by default, as this is what ORMs normally expect.
     // Standard batches return -1 by default, to indicate that the update count is unknown.
-    return isAutoBatch() ? autoBatchUpdateCountSupplier.get() : -1L;
+    return isAutoBatch() ? autoBatchUpdateCountSupplier.get() : dmlbatchUpdateCountSupplier.get();
   }
 
   @Override
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/StatementResult.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/StatementResult.java
@@ -120,6 +120,7 @@ enum ClientSideStatementType {
     SHOW_AUTO_BATCH_DML,
     SET_AUTO_BATCH_DML_UPDATE_COUNT,
     SHOW_AUTO_BATCH_DML_UPDATE_COUNT,
+    SET_BATCH_DML_UPDATE_COUNT,
     SET_AUTO_BATCH_DML_UPDATE_COUNT_VERIFICATION,
     SHOW_AUTO_BATCH_DML_UPDATE_COUNT_VERIFICATION,
     SHOW_READ_LOCK_MODE,
diff --git a/google-cloud-spanner/src/main/resources/com/google/cloud/spanner/connection/ClientSideStatements.json b/google-cloud-spanner/src/main/resources/com/google/cloud/spanner/connection/ClientSideStatements.json
@@ -753,6 +753,27 @@
 	    "converterName": "ClientSideStatementValueConverters$BooleanConverter"
 	  }
 	},
+  {
+    "name": "SET [LOCAL] BATCH_DML_UPDATE_COUNT = <INT64>",
+    "executorName": "ClientSideStatementSetExecutor",
+    "resultType": "NO_RESULT",
+    "statementType": "SET_BATCH_DML_UPDATE_COUNT",
+    "regex": "(?is)\\A\\s*set\\s+(local\\s+)?batch_dml_update_count\\s*(?:=)\\s*(.*)\\z",
+    "method": "statementSetBatchDmlUpdateCount",
+    "exampleStatements": [
+      "set local batch_dml_update_count = 0",
+      "set local batch_dml_update_count = 100",
+      "set batch_dml_update_count = 1",
+      "set batch_dml_update_count = 100"
+    ],
+    "examplePrerequisiteStatements": ["set readonly = false", "set autocommit = false"],
+    "setStatement": {
+      "propertyName": "BATCH_DML_UPDATE_COUNT",
+      "separator": "=",
+      "allowedValues": "(\\d{1,19})",
+      "converterName": "ClientSideStatementValueConverters$LongConverter"
+    }
+  },
 	{
 		"name": "SHOW VARIABLE READ_LOCK_MODE",
 		"executorName": "ClientSideStatementNoParamExecutor",
diff --git a/google-cloud-spanner/src/main/resources/com/google/cloud/spanner/connection/PG_ClientSideStatements.json b/google-cloud-spanner/src/main/resources/com/google/cloud/spanner/connection/PG_ClientSideStatements.json
@@ -966,6 +966,28 @@
         "converterName": "ClientSideStatementValueConverters$LongConverter"
       }
     },
+    {
+      "name": "SET [LOCAL] SPANNER.BATCH_DML_UPDATE_COUNT =|TO <bigint>",
+      "executorName": "ClientSideStatementSetExecutor",
+      "resultType": "NO_RESULT",
+      "statementType": "SET_BATCH_DML_UPDATE_COUNT",
+      "regex": "(?is)\\A\\s*set\\s+((?:session|local)\\s+)?spanner\\.batch_dml_update_count(?:\\s*=\\s*|\\s+to\\s+)(.*)\\z",
+      "method": "statementSetBatchDmlUpdateCount",
+      "exampleStatements": [
+        "set local spanner.batch_dml_update_count = 0",
+        "set local spanner.batch_dml_update_count = 100",
+        "set local spanner.batch_dml_update_count to 1",
+        "set spanner.batch_dml_update_count to 1",
+        "set spanner.batch_dml_update_count = 1"
+      ],
+      "examplePrerequisiteStatements": ["set spanner.readonly = false", "set autocommit = false"],
+      "setStatement": {
+        "propertyName": "SPANNER.BATCH_DML_UPDATE_COUNT",
+        "separator": "(?:=|\\s+TO\\s+)",
+        "allowedValues": "(\\d{1,19})",
+        "converterName": "ClientSideStatementValueConverters$LongConverter"
+      }
+    },
     {
       "name": "SET SPANNER.AUTO_BATCH_DML_UPDATE_COUNT_VERIFICATION = TRUE|FALSE",
       "executorName": "ClientSideStatementSetExecutor",
diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionAsyncApiTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionAsyncApiTest.java
@@ -28,8 +28,10 @@
 import com.google.cloud.spanner.AsyncResultSet;
 import com.google.cloud.spanner.AsyncResultSet.CallbackResponse;
 import com.google.cloud.spanner.AsyncResultSet.ReadyCallback;
+import com.google.cloud.spanner.Dialect;
 import com.google.cloud.spanner.ErrorCode;
 import com.google.cloud.spanner.ForceCloseSpannerFunction;
+import com.google.cloud.spanner.MockSpannerServiceImpl;
 import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime;
 import com.google.cloud.spanner.Mutation;
 import com.google.cloud.spanner.ResultSet;
@@ -48,6 +50,7 @@
 import com.google.spanner.v1.CommitRequest;
 import com.google.spanner.v1.ExecuteBatchDmlRequest;
 import com.google.spanner.v1.ExecuteSqlRequest;
+import java.util.Arrays;
 import java.util.List;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
@@ -336,6 +339,42 @@ public void testAutocommitRunBatch() {
     }
   }
 
+  @Test
+  public void testDmlBatchUpdateCount() {
+    Arrays.asList(Dialect.POSTGRESQL, Dialect.GOOGLE_STANDARD_SQL)
+        .forEach(
+            dialect -> {
+              String prefix = dialect == Dialect.POSTGRESQL ? "spanner." : "";
+              mockSpanner.putStatementResult(
+                  MockSpannerServiceImpl.StatementResult.detectDialectResult(dialect));
+              SpannerPool.closeSpannerPool();
+              try {
+                try (Connection connection = createConnection()) {
+                  connection.execute(
+                      Statement.of("set local " + prefix + "batch_dml_update_count = 1"));
+                  connection.execute(Statement.of("START BATCH DML"));
+                  List<Statement> statements = Arrays.asList(INSERT_STATEMENT, INSERT_STATEMENT);
+                  long[] updateCounts = connection.executeBatchUpdate(statements);
+                  assertThat(updateCounts).asList().containsExactly(1L, 1L);
+                  connection.execute(Statement.of("RUN BATCH"));
+                  connection.commit();
+
+                  connection.execute(Statement.of("START BATCH DML"));
+                  statements = Arrays.asList(INSERT_STATEMENT, INSERT_STATEMENT);
+                  updateCounts = connection.executeBatchUpdate(statements);
+                  assertThat(updateCounts).asList().containsExactly(-1L, -1L);
+                  connection.execute(Statement.of("RUN BATCH"));
+                  connection.commit();
+                }
+              } finally {
+                SpannerPool.closeSpannerPool();
+                mockSpanner.putStatementResult(
+                    MockSpannerServiceImpl.StatementResult.detectDialectResult(
+                        Dialect.GOOGLE_STANDARD_SQL));
+              }
+            });
+  }
+
   @Test
   public void testAutocommitRunBatchAsync() {
     try (Connection connection = createConnection()) {
diff --git a/google-cloud-spanner/src/test/resources/com/google/cloud/spanner/connection/ClientSideStatementsTest.sql b/google-cloud-spanner/src/test/resources/com/google/cloud/spanner/connection/ClientSideStatementsTest.sql
diff --git a/google-cloud-spanner/src/test/resources/com/google/cloud/spanner/connection/ConnectionImplGeneratedSqlScriptTest.sql b/google-cloud-spanner/src/test/resources/com/google/cloud/spanner/connection/ConnectionImplGeneratedSqlScriptTest.sql
diff --git a/google-cloud-spanner/src/test/resources/com/google/cloud/spanner/connection/postgresql/ClientSideStatementsTest.sql b/google-cloud-spanner/src/test/resources/com/google/cloud/spanner/connection/postgresql/ClientSideStatementsTest.sql
diff --git a/google-cloud-spanner/src/test/resources/com/google/cloud/spanner/connection/postgresql/ConnectionImplGeneratedSqlScriptTest.sql b/google-cloud-spanner/src/test/resources/com/google/cloud/spanner/connection/postgresql/ConnectionImplGeneratedSqlScriptTest.sql
