Skip to content

chore: Support dml_batch_update_count option in PG dialect #4142

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 7 commits into from
Oct 10, 2025
Merged

chore: Support dml_batch_update_count option in PG dialect #4142

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
abf6609
chore: Support dml_batch_update_count option in PG dialect
sakthivelmanii Oct 7, 2025
1b7f911
chore: generate libraries at Wed Oct 8 21:30:02 UTC 2025
cloud-java-bot Oct 8, 2025
1000fdc
Support in GoogleSQL dialect
sakthivelmanii Oct 9, 2025
c33c277
Address comments
sakthivelmanii Oct 9, 2025
e272679
Fix tests
sakthivelmanii Oct 9, 2025
aac9e9f
Update connection property name
sakthivelmanii Oct 10, 2025
71a3360
Address comments
sakthivelmanii Oct 10, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .github/workflows/update_generation_config.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: undo unrelated change? (Maybe in a follow-up PR)

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it's automatically added by bot.

with:
fetch-depth: 0
token: ${{ secrets.CLOUD_JAVA_BOT_TOKEN }}
Expand Down
10 changes: 10 additions & 0 deletions google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
Expand All @@ -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);
Expand Down Expand Up @@ -2340,6 +2349,7 @@ UnitOfWork createNewUnitOfWork(
.setAutoBatchUpdateCountSupplier(this::getAutoBatchDmlUpdateCount)
.setAutoBatchUpdateCountVerificationSupplier(
this::isAutoBatchDmlUpdateCountVerification)
.setDmlBatchUpdateCountSupplier(this::getDmlBatchUpdateCount)
.setTransaction(currentUnitOfWork)
.setStatementTimeout(statementTimeout)
.withStatementExecutor(statementExecutor)
Expand Down
2 changes: 2 additions & 0 deletions ...le-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionOptions.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down Expand Up @@ -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. "
Expand Down
11 changes: 11 additions & 0 deletions ...cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionProperties.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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();
Expand Down
2 changes: 2 additions & 0 deletions ...panner/src/main/java/com/google/cloud/spanner/connection/ConnectionStatementExecutor.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -184,6 +184,8 @@ StatementResult statementSetPgSessionCharacteristicsTransactionMode(

StatementResult statementRunPartitionedQuery(Statement statement);

StatementResult statementSetBatchDmlUpdateCount(Long updateCount, Boolean local);

StatementResult statementSetAutoBatchDml(Boolean autoBatchDml);

StatementResult statementShowAutoBatchDml();
Expand Down
7 changes: 7 additions & 0 deletions ...er/src/main/java/com/google/cloud/spanner/connection/ConnectionStatementExecutorImpl.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
Expand Down
11 changes: 10 additions & 1 deletion google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DmlBatch.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<>();
Expand All @@ -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;

Expand All @@ -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;
Expand Down Expand Up @@ -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;
}
Expand Down Expand Up @@ -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
Expand Down
1 change: 1 addition & 0 deletions google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/StatementResult.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
21 changes: 21 additions & 0 deletions ...-spanner/src/main/resources/com/google/cloud/spanner/connection/ClientSideStatements.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
22 changes: 22 additions & 0 deletions ...anner/src/main/resources/com/google/cloud/spanner/connection/PG_ClientSideStatements.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -966,6 +966,28 @@
"converterName": "ClientSideStatementValueConverters$LongConverter"
}
},
{
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Run the following command in the google-cloud-spanner folder:

mvn -Ddo_log_statements=true exec:java -Dexec.mainClass=com.google.cloud.spanner.connection.SqlTestScriptsGenerator -Dexec.classpathScope="test"

This will generate tests for the new statements that you have added to the client-side statement file(s).

See ConnectionImplGeneratedSqlScriptTest for more information.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done

"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",
Expand Down
39 changes: 39 additions & 0 deletions ...oud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionAsyncApiTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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()) {
Expand Down
Loading