chore: Support dml_batch_update_count option in PG dialect

Author: sakthivelmanii

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

@@ -31,6 +31,7 @@
 import static com.google.cloud.spanner.connection.ConnectionProperties.DEFAULT_SEQUENCE_KIND;
 import static com.google.cloud.spanner.connection.ConnectionProperties.DELAY_TRANSACTION_START_UNTIL_FIRST_WRITE;
 import static com.google.cloud.spanner.connection.ConnectionProperties.DIRECTED_READ;
+import static com.google.cloud.spanner.connection.ConnectionProperties.DML_BATCH_UPDATE_COUNT;
 import static com.google.cloud.spanner.connection.ConnectionProperties.KEEP_TRANSACTION_ALIVE;
 import static com.google.cloud.spanner.connection.ConnectionProperties.MAX_COMMIT_DELAY;
 import static com.google.cloud.spanner.connection.ConnectionProperties.MAX_PARTITIONED_PARALLELISM;
@@ -1607,6 +1608,10 @@ public long getAutoBatchDmlUpdateCount() {
     return getConnectionPropertyValue(AUTO_BATCH_DML_UPDATE_COUNT);
   }
 
+  long getDmlBatchUpdateCount() {
+    return getConnectionPropertyValue(DML_BATCH_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 setDmlBatchUpdateCount(long updateCount, boolean local) {
+    setConnectionPropertyValue(DML_BATCH_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)

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

@@ -178,6 +178,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_DML_BATCH_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";
@@ -308,6 +309,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 DML_BATCH_UPDATE_COUNT_PROPERTY_NAME = "dml_batch_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. "

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

@@ -43,6 +43,7 @@
 import static com.google.cloud.spanner.connection.ConnectionOptions.DEFAULT_DDL_IN_TRANSACTION_MODE;
 import static com.google.cloud.spanner.connection.ConnectionOptions.DEFAULT_DEFAULT_SEQUENCE_KIND;
 import static com.google.cloud.spanner.connection.ConnectionOptions.DEFAULT_DELAY_TRANSACTION_START_UNTIL_FIRST_WRITE;
+import static com.google.cloud.spanner.connection.ConnectionOptions.DEFAULT_DML_BATCH_UPDATE_COUNT;
 import static com.google.cloud.spanner.connection.ConnectionOptions.DEFAULT_ENABLE_API_TRACING;
 import static com.google.cloud.spanner.connection.ConnectionOptions.DEFAULT_ENABLE_END_TO_END_TRACING;
 import static com.google.cloud.spanner.connection.ConnectionOptions.DEFAULT_ENABLE_EXTENDED_TRACING;
@@ -72,6 +73,7 @@
 import static com.google.cloud.spanner.connection.ConnectionOptions.DEFAULT_USE_VIRTUAL_THREADS;
 import static com.google.cloud.spanner.connection.ConnectionOptions.DELAY_TRANSACTION_START_UNTIL_FIRST_WRITE_NAME;
 import static com.google.cloud.spanner.connection.ConnectionOptions.DIALECT_PROPERTY_NAME;
+import static com.google.cloud.spanner.connection.ConnectionOptions.DML_BATCH_UPDATE_COUNT_PROPERTY_NAME;
 import static com.google.cloud.spanner.connection.ConnectionOptions.ENABLE_API_TRACING_PROPERTY_NAME;
 import static com.google.cloud.spanner.connection.ConnectionOptions.ENABLE_END_TO_END_TRACING_PROPERTY_NAME;
 import static com.google.cloud.spanner.connection.ConnectionOptions.ENABLE_EXTENDED_TRACING_PROPERTY_NAME;
@@ -728,6 +730,16 @@ public class ConnectionProperties {
           BOOLEANS,
           BooleanConverter.INSTANCE,
           Context.USER);
+  static final ConnectionProperty<Long> DML_BATCH_UPDATE_COUNT =
+      create(
+          DML_BATCH_UPDATE_COUNT_PROPERTY_NAME,
+          "DML statements that are executed in batch when "
+              + " 'start batch dml' statement is executed. This property determines"
+              + " the update count that is returned for these DML statements. The default is "
+              + DEFAULT_DML_BATCH_UPDATE_COUNT,
+          DEFAULT_DML_BATCH_UPDATE_COUNT,
+          LongConverter.INSTANCE,
+          Context.USER);
 
   static final ImmutableMap<String, ConnectionProperty<?>> CONNECTION_PROPERTIES =
       CONNECTION_PROPERTIES_BUILDER.build();

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

@@ -184,6 +184,8 @@ StatementResult statementSetPgSessionCharacteristicsTransactionMode(
 
   StatementResult statementRunPartitionedQuery(Statement statement);
 
+  StatementResult statementSetDmlBatchUpdateCount(Long updateCount, Boolean local);
+
   StatementResult statementSetAutoBatchDml(Boolean autoBatchDml);
 
   StatementResult statementShowAutoBatchDml();

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

@@ -33,6 +33,7 @@
 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;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SET_DIRECTED_READ;
+import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SET_DML_BATCH_UPDATE_COUNT;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SET_EXCLUDE_TXN_FROM_CHANGE_STREAMS;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SET_KEEP_TRANSACTION_ALIVE;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SET_MAX_COMMIT_DELAY;
@@ -726,6 +727,12 @@ public StatementResult statementRunPartitionedQuery(Statement statement) {
         ClientSideStatementType.RUN_PARTITIONED_QUERY);
   }
 
+  @Override
+  public StatementResult statementSetDmlBatchUpdateCount(Long updateCount, Boolean local) {
+    getConnection().setDmlBatchUpdateCount(updateCount, local);
+    return noResult(SET_DML_BATCH_UPDATE_COUNT);
+  }
+
   @Override
   public StatementResult statementSetProtoDescriptors(byte[] protoDescriptors) {
     Preconditions.checkNotNull(protoDescriptors);

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,11 @@ Builder setAutoBatchUpdateCountVerificationSupplier(Supplier<Boolean> verificati
       return this;
     }
 
+    Builder setDmlBatchUpdateCountSupplier(Supplier<Long> dmlbatchUpdateCountSupplier) {
+      this.dmlbatchUpdateCountSupplier = dmlbatchUpdateCountSupplier;
+      return this;
+    }
+
     Builder setTransaction(UnitOfWork transaction) {
       Preconditions.checkNotNull(transaction);
       this.transaction = transaction;
@@ -108,6 +115,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 +201,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

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_DML_BATCH_UPDATE_COUNT,
     SET_AUTO_BATCH_DML_UPDATE_COUNT_VERIFICATION,
     SHOW_AUTO_BATCH_DML_UPDATE_COUNT_VERIFICATION,
     SHOW_READ_LOCK_MODE,

google-cloud-spanner/src/main/resources/com/google/cloud/spanner/connection/PG_ClientSideStatements.json

@@ -966,6 +966,25 @@
         "converterName": "ClientSideStatementValueConverters$LongConverter"
       }
     },
+    {
+      "name": "SET LOCAL SPANNER.DML_BATCH_UPDATE_COUNT =|TO <bigint>",
+      "executorName": "ClientSideStatementSetExecutor",
+      "resultType": "NO_RESULT",
+      "statementType": "SET_DML_BATCH_UPDATE_COUNT",
+      "regex": "(?is)\\A\\s*set\\s+((?:session|local)\\s+)?spanner\\.dml_batch_update_count(?:\\s*=\\s*|\\s+to\\s+)(.*)\\z",
+      "method": "statementSetDmlBatchUpdateCount",
+      "exampleStatements": [
+        "set local spanner.dml_batch_update_count = 0",
+        "set local spanner.dml_batch_update_count = 100",
+        "set local spanner.dml_batch_update_count to 1"
+      ],
+      "setStatement": {
+        "propertyName": "SPANNER.DML_BATCH_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",

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,32 @@ public void testAutocommitRunBatch() {
     }
   }
 
+  @Test
+  public void testDmlBatchUpdateCount() {
+    SpannerPool.closeSpannerPool();
+    mockSpanner.putStatementResult(
+        MockSpannerServiceImpl.StatementResult.detectDialectResult(Dialect.POSTGRESQL));
+    try {
+      try (Connection connection = createConnection()) {
+        connection.execute(Statement.of("set local spanner.dml_batch_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);
+      }
+      try (Connection connection = createConnection()) {
+        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);
+      }
+    } finally {
+      SpannerPool.closeSpannerPool();
+      mockSpanner.putStatementResult(
+          MockSpannerServiceImpl.StatementResult.detectDialectResult(Dialect.GOOGLE_STANDARD_SQL));
+    }
+  }
+
   @Test
   public void testAutocommitRunBatchAsync() {
     try (Connection connection = createConnection()) {