Repository: GoogleCloudPlatform/pgadapter (java)

Issue:
Connection ends up in ABORTED state after a failed COMMIT
If a `COMMIT` or a `ROLLBACK` fails, the connection is put into the `ABORTED` state. This requires the client application to execute a `ROLLBACK` before a new transaction can be executed. This is not consistent with real PostgreSQL, which puts the connection in the `IDLE` state if a  `COMMIT` or  `ROLLBACK` fails.

Pull request description:
fix: connection should be idle after failed commit
A connection would be put in the ABORTED state if a COMMIT statement failed with an error. This is not consistent with real PostgreSQL, which puts the connection in the IDLE state after a COMMIT fails. The same applies to ROLLBACK.

Fixes #1836

Patch:
diff --git a/benchmarks/tpcc/pom.xml b/benchmarks/tpcc/pom.xml
index 501f2257..564a5caf 100644
--- a/benchmarks/tpcc/pom.xml
+++ b/benchmarks/tpcc/pom.xml
@@ -28,7 +28,7 @@
       <dependency>
         <groupId>io.opentelemetry</groupId>
         <artifactId>opentelemetry-bom</artifactId>
-        <version>1.37.0</version>
+        <version>1.38.0</version>
         <type>pom</type>
         <scope>import</scope>
       </dependency>
@@ -59,7 +59,7 @@
     <dependency>
           <groupId>io.opentelemetry</groupId>
           <artifactId>opentelemetry-sdk</artifactId>
-          <version>1.35.0</version>
+          <version>1.38.0</version>
     </dependency>
     <dependency>
           <groupId>io.opentelemetry</groupId>
@@ -69,7 +69,7 @@
     <dependency>
         <groupId>io.opentelemetry</groupId>
         <artifactId>opentelemetry-exporter-otlp</artifactId>
-        <version>1.35.0</version>
+        <version>1.38.0</version>
     </dependency>
   </dependencies>
 
diff --git a/samples/golang/gorm/go.mod b/samples/golang/gorm/go.mod
index f9de1437..0e8af400 100644
--- a/samples/golang/gorm/go.mod
+++ b/samples/golang/gorm/go.mod
@@ -3,7 +3,7 @@ module gorm
 go 1.21
 
 require (
-	github.com/GoogleCloudPlatform/pgadapter/wrappers/golang v0.0.0-20240516111254-4142f0a3ddee
+	github.com/GoogleCloudPlatform/pgadapter/wrappers/golang v0.0.0-20240522081835-cfb53f31e39f
 	github.com/google/uuid v1.6.0
 	github.com/jackc/pgtype v1.14.3
 	github.com/shopspring/decimal v1.4.0
@@ -26,7 +26,7 @@ require (
 	github.com/containerd/log v0.1.0 // indirect
 	github.com/cpuguy83/dockercfg v0.3.1 // indirect
 	github.com/distribution/reference v0.5.0 // indirect
-	github.com/docker/docker v26.1.2+incompatible // indirect
+	github.com/docker/docker v26.1.3+incompatible // indirect
 	github.com/docker/go-connections v0.5.0 // indirect
 	github.com/docker/go-units v0.5.0 // indirect
 	github.com/felixge/httpsnoop v1.0.4 // indirect
diff --git a/samples/golang/gorm/go.sum b/samples/golang/gorm/go.sum
index d37bc2bf..0ba72fed 100644
--- a/samples/golang/gorm/go.sum
+++ b/samples/golang/gorm/go.sum
@@ -15,6 +15,8 @@ github.com/GoogleCloudPlatform/pgadapter/wrappers/golang v0.0.0-20240124134511-6
 github.com/GoogleCloudPlatform/pgadapter/wrappers/golang v0.0.0-20240124134511-6f84a1722435/go.mod h1:RDjSohojBNouglltL1D+Z1E5l35DiSA+bQrwqKt+nEU=
 github.com/GoogleCloudPlatform/pgadapter/wrappers/golang v0.0.0-20240516111254-4142f0a3ddee h1:7mu4ylrYLCPYRPQ7BaYk0MJuFZVSSr0URqrh1o050m0=
 github.com/GoogleCloudPlatform/pgadapter/wrappers/golang v0.0.0-20240516111254-4142f0a3ddee/go.mod h1:ukAIL5yxFOr/l8R1uDAlBVwJAcQoYZGLcKJfhnifaH4=
+github.com/GoogleCloudPlatform/pgadapter/wrappers/golang v0.0.0-20240522081835-cfb53f31e39f h1:T4YWsLWNb9OPxsmf8m1PyYKwM4STgNiyfQVNMbLhk8k=
+github.com/GoogleCloudPlatform/pgadapter/wrappers/golang v0.0.0-20240522081835-cfb53f31e39f/go.mod h1:DhQKwQgWo9DILxKwvNu93rJ1GDZYUu7m87V+KZCwgoc=
 github.com/Masterminds/semver/v3 v3.1.1/go.mod h1:VPu/7SZ7ePZ3QOrcuXROw5FAcLl4a0cBrbBpGY/8hQs=
 github.com/Microsoft/go-winio v0.6.1 h1:9/kr64B9VUZrLm5YYwbGtUJnMgqWVOdUAXu6Migciow=
 github.com/Microsoft/go-winio v0.6.1/go.mod h1:LRdKpFKfdobln8UmuiYcKPot9D2v6svN5+sAH+4kjUM=
@@ -43,6 +45,8 @@ github.com/docker/docker v25.0.5+incompatible h1:UmQydMduGkrD5nQde1mecF/YnSbTOaP
 github.com/docker/docker v25.0.5+incompatible/go.mod h1:eEKB0N0r5NX/I1kEveEz05bcu8tLC/8azJZsviup8Sk=
 github.com/docker/docker v26.1.2+incompatible h1:UVX5ZOrrfTGZZYEP+ZDq3Xn9PdHNXaSYMFPDumMqG2k=
 github.com/docker/docker v26.1.2+incompatible/go.mod h1:eEKB0N0r5NX/I1kEveEz05bcu8tLC/8azJZsviup8Sk=
+github.com/docker/docker v26.1.3+incompatible h1:lLCzRbrVZrljpVNobJu1J2FHk8V0s4BawoZippkc+xo=
+github.com/docker/docker v26.1.3+incompatible/go.mod h1:eEKB0N0r5NX/I1kEveEz05bcu8tLC/8azJZsviup8Sk=
 github.com/docker/go-connections v0.5.0 h1:USnMq7hx7gwdVZq1L49hLXaFtUdTADjXGp+uj1Br63c=
 github.com/docker/go-connections v0.5.0/go.mod h1:ov60Kzw0kKElRwhNs9UlUHAE/F9Fe6GLaXnqyDdmEXc=
 github.com/docker/go-units v0.5.0 h1:69rxXcBk27SvSaaxTtLh/8llcHD8vYHT7WSdRZ/jvr4=
diff --git a/samples/golang/pgx/go.mod b/samples/golang/pgx/go.mod
index 30af25bc..94b584b1 100644
--- a/samples/golang/pgx/go.mod
+++ b/samples/golang/pgx/go.mod
@@ -5,7 +5,7 @@ go 1.21
 toolchain go1.22.3
 
 require (
-	github.com/GoogleCloudPlatform/pgadapter/wrappers/golang v0.0.0-20240516111254-4142f0a3ddee
+	github.com/GoogleCloudPlatform/pgadapter/wrappers/golang v0.0.0-20240522081835-cfb53f31e39f
 	github.com/jackc/pgx/v5 v5.5.5
 )
 
@@ -22,7 +22,7 @@ require (
 	github.com/cpuguy83/dockercfg v0.3.1 // indirect
 	github.com/distribution/reference v0.5.0 // indirect
 	github.com/docker/distribution v2.8.2+incompatible // indirect
-	github.com/docker/docker v26.1.2+incompatible // indirect
+	github.com/docker/docker v26.1.3+incompatible // indirect
 	github.com/docker/go-connections v0.5.0 // indirect
 	github.com/docker/go-units v0.5.0 // indirect
 	github.com/felixge/httpsnoop v1.0.4 // indirect
diff --git a/samples/golang/pgx/go.sum b/samples/golang/pgx/go.sum
index 39122aa4..01b6aebf 100644
--- a/samples/golang/pgx/go.sum
+++ b/samples/golang/pgx/go.sum
@@ -13,6 +13,8 @@ github.com/GoogleCloudPlatform/pgadapter/wrappers/golang v0.0.0-20240124134511-6
 github.com/GoogleCloudPlatform/pgadapter/wrappers/golang v0.0.0-20240124134511-6f84a1722435/go.mod h1:RDjSohojBNouglltL1D+Z1E5l35DiSA+bQrwqKt+nEU=
 github.com/GoogleCloudPlatform/pgadapter/wrappers/golang v0.0.0-20240516111254-4142f0a3ddee h1:7mu4ylrYLCPYRPQ7BaYk0MJuFZVSSr0URqrh1o050m0=
 github.com/GoogleCloudPlatform/pgadapter/wrappers/golang v0.0.0-20240516111254-4142f0a3ddee/go.mod h1:ukAIL5yxFOr/l8R1uDAlBVwJAcQoYZGLcKJfhnifaH4=
+github.com/GoogleCloudPlatform/pgadapter/wrappers/golang v0.0.0-20240522081835-cfb53f31e39f h1:T4YWsLWNb9OPxsmf8m1PyYKwM4STgNiyfQVNMbLhk8k=
+github.com/GoogleCloudPlatform/pgadapter/wrappers/golang v0.0.0-20240522081835-cfb53f31e39f/go.mod h1:DhQKwQgWo9DILxKwvNu93rJ1GDZYUu7m87V+KZCwgoc=
 github.com/Microsoft/go-winio v0.6.1 h1:9/kr64B9VUZrLm5YYwbGtUJnMgqWVOdUAXu6Migciow=
 github.com/Microsoft/go-winio v0.6.1/go.mod h1:LRdKpFKfdobln8UmuiYcKPot9D2v6svN5+sAH+4kjUM=
 github.com/Microsoft/hcsshim v0.11.4 h1:68vKo2VN8DE9AdN4tnkWnmdhqdbpUFM8OF3Airm7fz8=
@@ -39,6 +41,8 @@ github.com/docker/docker v24.0.9+incompatible h1:HPGzNmwfLZWdxHqK9/II92pyi1EpYKs
 github.com/docker/docker v24.0.9+incompatible/go.mod h1:eEKB0N0r5NX/I1kEveEz05bcu8tLC/8azJZsviup8Sk=
 github.com/docker/docker v26.1.2+incompatible h1:UVX5ZOrrfTGZZYEP+ZDq3Xn9PdHNXaSYMFPDumMqG2k=
 github.com/docker/docker v26.1.2+incompatible/go.mod h1:eEKB0N0r5NX/I1kEveEz05bcu8tLC/8azJZsviup8Sk=
+github.com/docker/docker v26.1.3+incompatible h1:lLCzRbrVZrljpVNobJu1J2FHk8V0s4BawoZippkc+xo=
+github.com/docker/docker v26.1.3+incompatible/go.mod h1:eEKB0N0r5NX/I1kEveEz05bcu8tLC/8azJZsviup8Sk=
 github.com/docker/go-connections v0.4.0 h1:El9xVISelRB7BuFusrZozjnkIM5YnzCViNKohAFqRJQ=
 github.com/docker/go-connections v0.4.0/go.mod h1:Gbd7IOopHjR8Iph03tsViu4nIes5XhDvyHbTtUxmeec=
 github.com/docker/go-connections v0.5.0 h1:USnMq7hx7gwdVZq1L49hLXaFtUdTADjXGp+uj1Br63c=
diff --git a/samples/java/hibernate/pom.xml b/samples/java/hibernate/pom.xml
index 5e9be88a..a913c88f 100644
--- a/samples/java/hibernate/pom.xml
+++ b/samples/java/hibernate/pom.xml
@@ -39,7 +39,7 @@
     <dependency>
       <groupId>org.hibernate.orm</groupId>
       <artifactId>hibernate-core</artifactId>
-      <version>6.5.1.Final</version>
+      <version>6.5.2.Final</version>
     </dependency>
     <!-- Postgresql JDBC driver dependency -->
     <dependency>
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/statements/BackendConnection.java b/src/main/java/com/google/cloud/spanner/pgadapter/statements/BackendConnection.java
index 1c5709d6..9970892e 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/statements/BackendConnection.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/statements/BackendConnection.java
@@ -1168,15 +1168,18 @@ public class BackendConnection {
               sessionState.rollback();
             }
             closeAllPortals.run();
-            transactionMode = TransactionMode.IMPLICIT;
-            connectionState = ConnectionState.IDLE;
-            currentTransactionId = null;
+            clearCurrentTransaction();
           }
           index++;
         }
       }
     } catch (Exception exception) {
-      connectionState = ConnectionState.ABORTED;
+      // The connection should not transition to the ABORTED state if a COMMIT or ROLLBACK fails.
+      if (isCommit(index) || isRollback(index)) {
+        clearCurrentTransaction();
+      } else {
+        connectionState = ConnectionState.ABORTED;
+      }
       closeAllPortals.run();
       sessionState.rollback();
       if (spannerConnection.isInTransaction()) {
@@ -1193,6 +1196,12 @@ public class BackendConnection {
     }
   }
 
+  private void clearCurrentTransaction() {
+    transactionMode = TransactionMode.IMPLICIT;
+    connectionState = ConnectionState.IDLE;
+    currentTransactionId = null;
+  }
+
   /** Starts an implicit transaction if that is necessary. */
   private void maybeBeginImplicitTransaction(int index, boolean isSync) {
     if (connectionState != ConnectionState.IDLE) {


Test patch:
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/AbortedMockServerTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/AbortedMockServerTest.java
index 7c952996..e2817887 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/AbortedMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/AbortedMockServerTest.java
@@ -37,6 +37,7 @@ import com.google.cloud.spanner.Spanner;
 import com.google.cloud.spanner.SpannerOptions;
 import com.google.cloud.spanner.Statement;
 import com.google.cloud.spanner.connection.RandomResultSetGenerator;
+import com.google.cloud.spanner.pgadapter.error.SQLState;
 import com.google.common.collect.ImmutableList;
 import com.google.spanner.v1.ExecuteSqlRequest;
 import com.google.spanner.v1.ResultSetStats;
@@ -68,6 +69,7 @@ import org.postgresql.PGConnection;
 import org.postgresql.PGStatement;
 import org.postgresql.core.Oid;
 import org.postgresql.jdbc.PgStatement;
+import org.postgresql.util.PSQLException;
 
 @RunWith(JUnit4.class)
 public class AbortedMockServerTest extends AbstractMockServerTest {
@@ -937,6 +939,32 @@ public class AbortedMockServerTest extends AbstractMockServerTest {
     }
   }
 
+  @Test
+  public void testCommitFailsDueToConcurrentModification() throws SQLException {
+    String sql = "insert into foo (id, value) values (1, 'One') on conflict do nothing";
+    mockSpanner.putStatementResult(StatementResult.update(Statement.of(sql), 1L));
+
+    try (Connection connection = createConnection()) {
+      connection.setAutoCommit(false);
+      assertEquals(1, connection.createStatement().executeUpdate(sql));
+
+      // Change the update count that is returned and abort the transaction. This will cause the
+      // commit() to fail due to a concurrent modification.
+      mockSpanner.putStatementResult(StatementResult.update(Statement.of(sql), 0));
+      mockSpanner.abortNextStatement();
+
+      PSQLException psqlException = assertThrows(PSQLException.class, connection::commit);
+      assertNotNull(psqlException.getServerErrorMessage());
+      assertEquals(
+          SQLState.SerializationFailure.toString(),
+          psqlException.getServerErrorMessage().getSQLState());
+
+      // Verify that the connection is usable for a new transaction after a failed commit.
+      assertEquals(0, connection.createStatement().executeUpdate(sql));
+      connection.commit();
+    }
+  }
+
   private void assertEqual(com.google.cloud.spanner.ResultSet spannerResult, ResultSet pgResult)
       throws SQLException {
     assertEquals(spannerResult.getColumnCount(), pgResult.getMetaData().getColumnCount());
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/statements/BackendConnectionTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/statements/BackendConnectionTest.java
index f56f53e2..58a06cd4 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/statements/BackendConnectionTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/statements/BackendConnectionTest.java
@@ -44,6 +44,7 @@ import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStateme
 import com.google.cloud.spanner.connection.AbstractStatementParser.StatementType;
 import com.google.cloud.spanner.connection.Connection;
 import com.google.cloud.spanner.connection.StatementResult;
+import com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType;
 import com.google.cloud.spanner.connection.StatementResult.ResultType;
 import com.google.cloud.spanner.connection.TransactionMode;
 import com.google.cloud.spanner.pgadapter.AbstractMockServerTest;
@@ -53,6 +54,7 @@ import com.google.cloud.spanner.pgadapter.error.PGExceptionFactory;
 import com.google.cloud.spanner.pgadapter.error.SQLState;
 import com.google.cloud.spanner.pgadapter.metadata.ConnectionMetadata;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
+import com.google.cloud.spanner.pgadapter.statements.BackendConnection.ConnectionState;
 import com.google.cloud.spanner.pgadapter.statements.BackendConnection.NoResult;
 import com.google.cloud.spanner.pgadapter.statements.BackendConnection.QueryResult;
 import com.google.cloud.spanner.pgadapter.statements.local.ListDatabasesStatement;
@@ -494,6 +496,133 @@ public class BackendConnectionTest {
     assertEquals(executionException.getCause(), PGExceptionFactory.toPGException(error));
   }
 
+  @Test
+  public void testUpdateException_leavesConnectionInAbortedState() {
+    Connection connection = mock(Connection.class);
+    when(connection.isInTransaction()).thenReturn(true);
+
+    Statement statement = Statement.of("insert into foo (id) values (1)");
+    ParsedStatement parsedStatement = mock(ParsedStatement.class);
+    when(parsedStatement.getSqlWithoutComments()).thenReturn(statement.getSql());
+    RuntimeException error = new RuntimeException("test error");
+    when(connection.execute(statement)).thenThrow(error);
+
+    BackendConnection backendConnection =
+        new BackendConnection(
+            NOOP_OTEL,
+            NOOP_OTEL_METER,
+            METRIC_ATTRIBUTES,
+            UUID.randomUUID().toString(),
+            DO_NOTHING,
+            DATABASE_ID,
+            connection,
+            () -> WellKnownClient.UNSPECIFIED,
+            mock(OptionsMetadata.class),
+            () -> EMPTY_LOCAL_STATEMENTS);
+    Future<StatementResult> resultFuture =
+        backendConnection.execute("INSERT", parsedStatement, statement, Function.identity());
+    backendConnection.flush();
+
+    ExecutionException executionException =
+        assertThrows(ExecutionException.class, resultFuture::get);
+    assertEquals(PGException.class, executionException.getCause().getClass());
+    PGException pgException = (PGException) executionException.getCause();
+    assertTrue(pgException.getMessage().contains("test error"));
+
+    // Verify that the connection is now in the Aborted state.
+    assertEquals(ConnectionState.ABORTED, backendConnection.getConnectionState());
+  }
+
+  @Test
+  public void testCommitException_leavesConnectionInIdleState() {
+    Connection connection = mock(Connection.class);
+    when(connection.isInTransaction()).thenReturn(true);
+
+    Statement commitStatement = Statement.of("commit");
+    ParsedStatement parsedCommitStatement = mock(ParsedStatement.class);
+    when(parsedCommitStatement.getSqlWithoutComments()).thenReturn(commitStatement.getSql());
+    when(parsedCommitStatement.getType()).thenReturn(StatementType.CLIENT_SIDE);
+    when(parsedCommitStatement.getClientSideStatementType())
+        .thenReturn(ClientSideStatementType.COMMIT);
+    StatementResult commitResult = mock(StatementResult.class);
+    when(commitResult.getResultType()).thenReturn(ResultType.NO_RESULT);
+    when(connection.execute(commitStatement))
+        .thenThrow(
+            SpannerExceptionFactory.newSpannerException(
+                ErrorCode.FAILED_PRECONDITION, "internal error"));
+
+    BackendConnection backendConnection =
+        new BackendConnection(
+            NOOP_OTEL,
+            NOOP_OTEL_METER,
+            METRIC_ATTRIBUTES,
+            UUID.randomUUID().toString(),
+            DO_NOTHING,
+            DATABASE_ID,
+            connection,
+            () -> WellKnownClient.UNSPECIFIED,
+            mock(OptionsMetadata.class),
+            () -> EMPTY_LOCAL_STATEMENTS);
+    Future<StatementResult> commitFuture =
+        backendConnection.execute(
+            "COMMIT", parsedCommitStatement, commitStatement, Function.identity());
+    backendConnection.flush();
+
+    ExecutionException executionException =
+        assertThrows(ExecutionException.class, commitFuture::get);
+    assertEquals(PGException.class, executionException.getCause().getClass());
+    PGException pgException = (PGException) executionException.getCause();
+    assertTrue(pgException.getMessage().contains("internal error"));
+
+    // Verify that the connection is now in the idle state.
+    assertEquals(ConnectionState.IDLE, backendConnection.getConnectionState());
+  }
+
+  @Test
+  public void testRollbackException_leavesConnectionInIdleState() {
+    Connection connection = mock(Connection.class);
+    when(connection.isInTransaction()).thenReturn(true);
+
+    Statement rollbackStatement = Statement.of("rollback");
+    ParsedStatement parsedRollbackStatement = mock(ParsedStatement.class);
+    when(parsedRollbackStatement.getSqlWithoutComments()).thenReturn(rollbackStatement.getSql());
+    when(parsedRollbackStatement.getType()).thenReturn(StatementType.CLIENT_SIDE);
+    when(parsedRollbackStatement.getClientSideStatementType())
+        .thenReturn(ClientSideStatementType.ROLLBACK);
+    StatementResult rollbackResult = mock(StatementResult.class);
+    when(rollbackResult.getResultType()).thenReturn(ResultType.NO_RESULT);
+    when(connection.execute(rollbackStatement))
+        .thenThrow(
+            SpannerExceptionFactory.newSpannerException(
+                ErrorCode.FAILED_PRECONDITION, "internal error"));
+
+    BackendConnection backendConnection =
+        new BackendConnection(
+            NOOP_OTEL,
+            NOOP_OTEL_METER,
+            METRIC_ATTRIBUTES,
+            UUID.randomUUID().toString(),
+            DO_NOTHING,
+            DATABASE_ID,
+            connection,
+            () -> WellKnownClient.UNSPECIFIED,
+            mock(OptionsMetadata.class),
+            () -> EMPTY_LOCAL_STATEMENTS);
+    Future<StatementResult> rollbackFuture =
+        backendConnection.execute(
+            "ROLLBACK", parsedRollbackStatement, rollbackStatement, Function.identity());
+    backendConnection.flush();
+
+    ExecutionException executionException =
+        assertThrows(ExecutionException.class, rollbackFuture::get);
+    assertEquals(PGException.class, executionException.getCause().getClass());
+    PGException pgException = (PGException) executionException.getCause();
+    assertTrue(pgException.getMessage().contains("internal error"));
+
+    // Verify that the connection is now in the idle state.
+    assertEquals(ConnectionState.IDLE, backendConnection.getConnectionState());
+  }
+
   @Test
   public void testCancelledException() {
     Connection connection = mock(Connection.class);