From d132138fefaf2fb0a54f15124ab91e5763548694 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= <koloite@gmail.com>
Date: Tue, 10 Jun 2025 14:45:38 +0200
Subject: [PATCH] test: autocommit_dml_mode with PG

---
 .../ext/spanner/AbstractMockServerTest.java   | 11 +++++++-
 .../ext/spanner/SimpleMockServerTest.java     | 28 +++++++++++++++++++
 2 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/src/test/java/liquibase/ext/spanner/AbstractMockServerTest.java b/src/test/java/liquibase/ext/spanner/AbstractMockServerTest.java
index 850460fb..20ac26ad 100644
--- a/src/test/java/liquibase/ext/spanner/AbstractMockServerTest.java
+++ b/src/test/java/liquibase/ext/spanner/AbstractMockServerTest.java
@@ -452,10 +452,19 @@ protected static Connection createConnection() throws SQLException {
     return DriverManager.getConnection(createConnectionUrl());
   }
 
+  protected static Connection createConnection(String database) throws SQLException {
+    return DriverManager.getConnection(createConnectionUrl(database));
+  }
+
   protected static String createConnectionUrl() {
+    return createConnectionUrl("/projects/p/instances/i/databases/d");
+  }
+
+  protected static String createConnectionUrl(String database) {
     return new StringBuilder("jdbc:cloudspanner://localhost:")
         .append(server.getPort())
-        .append("/projects/p/instances/i/databases/d;usePlainText=true;minSessions=0")
+        .append(database)
+        .append(";usePlainText=true;minSessions=0")
         .toString();
   }
 }
diff --git a/src/test/java/liquibase/ext/spanner/SimpleMockServerTest.java b/src/test/java/liquibase/ext/spanner/SimpleMockServerTest.java
index de8e8f50..d91b3584 100644
--- a/src/test/java/liquibase/ext/spanner/SimpleMockServerTest.java
+++ b/src/test/java/liquibase/ext/spanner/SimpleMockServerTest.java
@@ -2,11 +2,16 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
+import com.google.cloud.spanner.Dialect;
 import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult;
+import com.google.cloud.spanner.connection.SpannerPool;
 import com.google.common.collect.ImmutableList;
 import com.google.spanner.admin.database.v1.UpdateDatabaseDdlRequest;
+import com.google.spanner.v1.BeginTransactionRequest;
+import com.google.spanner.v1.ExecuteSqlRequest;
 import java.sql.Connection;
 import java.sql.SQLException;
+import java.sql.Statement;
 import liquibase.Liquibase;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.parallel.Execution;
@@ -34,6 +39,29 @@ void testCreateTableStatement() throws SQLException {
     receivedRequestWithNonLiquibaseToken.set(false);
   }
 
+
+  @Test
+  void testAutocommitDmlMode() throws SQLException {
+    mockSpanner.putStatementResult(StatementResult.update(com.google.cloud.spanner.Statement.of("update foo set bar=1 where true"), 10000L));
+
+    // Configure the mock server to return POSTGRESQL as the dialect that is used by the database.
+    mockSpanner.putStatementResult(StatementResult.detectDialectResult(Dialect.POSTGRESQL));
+
+    try (Connection con = createConnection("/projects/p/instances/i/databases/pg_db"); Statement stmt = con.createStatement()) {
+      stmt.execute("set spanner.autocommit_dml_mode='partitioned_non_atomic'");
+      stmt.execute("update foo set bar=1 where true");
+    }
+    assertThat(mockSpanner.getRequestsOfType(BeginTransactionRequest.class)).hasSize(1);
+    BeginTransactionRequest request = mockSpanner.getRequestsOfType(BeginTransactionRequest.class).get(0);
+    assertThat(request.getOptions().hasPartitionedDml()).isTrue();
+
+    // This test does not use Liquibase but JDBC directly, so it is expected to send requests that
+    // do not contain a Liquibase client lib token.
+    assertThat(receivedRequestWithNonLiquibaseToken.get()).isTrue();
+    // Clear the flag to prevent the check after each test to fail.
+    receivedRequestWithNonLiquibaseToken.set(false);
+  }
+
   @Test
   void testInitLiquibase() throws Exception {
     mockSpanner.putStatementResult(