Repository: GoogleCloudPlatform/pgadapter
Pull request #498 (2022-11-20 20:36:55)

feat: support DML RETURNING clause
Add support for DML statements with a RETURNING clause.

codecov[bot]: # [Codecov](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/498?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) Report
> Merging [#498](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/498?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) (04b4d4f) into [remove-guess-types](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/commit/c955b810e845f3918b60941b86cea9e311eb9c1b?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) (c955b81) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

```diff
@@                  Coverage Diff                  @@
##             remove-guess-types     #498   +/-   ##
=====================================================
  Coverage                 87.66%   87.67%           
- Complexity                 2004     2008    +4     
=====================================================
  Files                       120      120           
  Lines                      6567     6572    +5     
  Branches                    929      932    +3     
=====================================================
+ Hits                       5757     5762    +5     
  Misses                      577      577           
  Partials                    233      233           
```

| Flag | Coverage Δ | |
|---|---|---|
| all_tests | `87.67% <100.00%> (+<0.01%)` | :arrow_up: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#carryforward-flags-in-the-pull-request-comment) to find out more.

| [Impacted Files](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/498?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) | Coverage Δ | |
|---|---|---|
| [...spanner/pgadapter/metadata/SendResultSetState.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/498/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL21ldGFkYXRhL1NlbmRSZXN1bHRTZXRTdGF0ZS5qYXZh) | `100.00% <100.00%> (ø)` | |
| [...panner/pgadapter/statements/BackendConnection.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/498/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3N0YXRlbWVudHMvQmFja2VuZENvbm5lY3Rpb24uamF2YQ==) | `92.57% <100.00%> (ø)` | |
| [...er/pgadapter/statements/IntermediateStatement.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/498/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3N0YXRlbWVudHMvSW50ZXJtZWRpYXRlU3RhdGVtZW50LmphdmE=) | `98.88% <100.00%> (+0.01%)` | :arrow_up: |
| [...spanner/pgadapter/wireprotocol/ControlMessage.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/498/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3dpcmVwcm90b2NvbC9Db250cm9sTWVzc2FnZS5qYXZh) | `86.66% <100.00%> (+0.12%)` | :arrow_up: |

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform)


Patch:
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/metadata/SendResultSetState.java b/src/main/java/com/google/cloud/spanner/pgadapter/metadata/SendResultSetState.java
index b5ef4570..ad796c51 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/metadata/SendResultSetState.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/metadata/SendResultSetState.java
@@ -31,7 +31,9 @@ public class SendResultSetState {
   }
 
   public String getCommandAndNumRows() {
-    return getCommandTag() + " " + getNumberOfRowsSent();
+    String command = getCommandTag();
+    command += ("INSERT".equals(command) ? " 0 " : " ") + getNumberOfRowsSent();
+    return command;
   }
 
   public String getCommandTag() {
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 665105e9..3dd97159 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
@@ -277,7 +277,7 @@ public class BackendConnection {
       statement = statementBinder.apply(statement);
       if (analyze) {
         ResultSet resultSet;
-        if (parsedStatement.isUpdate()) {
+        if (parsedStatement.isUpdate() && !parsedStatement.hasReturningClause()) {
           // TODO(#477): Single analyzeUpdate statements that are executed in an implicit
           //     transaction could use a single-use read/write transaction. Replays are not
           //     dangerous for those.
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ControlMessage.java b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ControlMessage.java
index 0a3b07a0..f75639cd 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ControlMessage.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ControlMessage.java
@@ -28,6 +28,7 @@ import com.google.cloud.spanner.SpannerOptions;
 import com.google.cloud.spanner.connection.Connection;
 import com.google.cloud.spanner.connection.ConnectionOptionsHelper;
 import com.google.cloud.spanner.connection.StatementResult;
+import com.google.cloud.spanner.connection.StatementResult.ResultType;
 import com.google.cloud.spanner.pgadapter.ConnectionHandler;
 import com.google.cloud.spanner.pgadapter.ConnectionHandler.ConnectionStatus;
 import com.google.cloud.spanner.pgadapter.ConnectionHandler.QueryMode;
@@ -253,7 +254,6 @@ public abstract class ControlMessage extends WireMessage {
     if (statement.getStatementResult() == null) {
       return;
     }
-
     switch (statement.getStatementType()) {
       case DDL:
       case CLIENT_SIDE:
@@ -261,23 +261,26 @@ public abstract class ControlMessage extends WireMessage {
         new CommandCompleteResponse(this.outputStream, command).send(false);
         break;
       case QUERY:
-        SendResultSetState state = sendResultSet(statement, mode, maxRows);
-        statement.setHasMoreData(state.hasMoreRows());
-        if (state.hasMoreRows()) {
-          new PortalSuspendedResponse(this.outputStream).send(false);
+      case UPDATE:
+        if (statement.getStatementResult().getResultType() == ResultType.RESULT_SET) {
+          SendResultSetState state = sendResultSet(statement, mode, maxRows);
+          statement.setHasMoreData(state.hasMoreRows());
+          if (state.hasMoreRows()) {
+            new PortalSuspendedResponse(this.outputStream).send(false);
+          } else {
+            statement.close();
+            new CommandCompleteResponse(this.outputStream, state.getCommandAndNumRows())
+                .send(false);
+          }
         } else {
-          statement.close();
-          new CommandCompleteResponse(this.outputStream, state.getCommandAndNumRows()).send(false);
+          // For an INSERT command, the tag is INSERT oid rows, where rows is the number of rows
+          // inserted. oid used to be the object ID of the inserted row if rows was 1 and the target
+          // table had OIDs, but OIDs system columns are not supported anymore; therefore oid is
+          // always 0.
+          command += ("INSERT".equals(command) ? " 0 " : " ") + statement.getUpdateCount();
+          new CommandCompleteResponse(this.outputStream, command).send(false);
         }
         break;
-      case UPDATE:
-        // For an INSERT command, the tag is INSERT oid rows, where rows is the number of rows
-        // inserted. oid used to be the object ID of the inserted row if rows was 1 and the target
-        // table had OIDs, but OIDs system columns are not supported anymore; therefore oid is
-        // always 0.
-        command += ("INSERT".equals(command) ? " 0 " : " ") + statement.getUpdateCount();
-        new CommandCompleteResponse(this.outputStream, command).send(false);
-        break;
       default:
         throw new IllegalStateException("Unknown statement type: " + statement.getStatement());
     }


Test patch:
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/statements/IntermediateStatement.java b/src/main/java/com/google/cloud/spanner/pgadapter/statements/IntermediateStatement.java
index 73bdfcf9..1cacc63e 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/statements/IntermediateStatement.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/statements/IntermediateStatement.java
@@ -137,7 +137,8 @@ public class IntermediateStatement {
 
   /** @return True if this is a select statement, false otherwise. */
   public boolean containsResultSet() {
-    return this.parsedStatement.isQuery();
+    return this.parsedStatement.isQuery()
+        || (this.parsedStatement.isUpdate() && this.parsedStatement.hasReturningClause());
   }
 
   /** @return True if this statement was executed, False otherwise. */
diff --git a/src/test/golang/pgadapter_pgx_tests/pgx.go b/src/test/golang/pgadapter_pgx_tests/pgx.go
index 48e56c60..cee4fcd7 100644
--- a/src/test/golang/pgadapter_pgx_tests/pgx.go
+++ b/src/test/golang/pgadapter_pgx_tests/pgx.go
@@ -254,6 +254,95 @@ func TestInsertNullsAllDataTypes(connString string) *C.char {
 	return nil
 }
 
+//export TestInsertAllDataTypesReturning
+func TestInsertAllDataTypesReturning(connString string) *C.char {
+	ctx := context.Background()
+	conn, err := pgx.Connect(ctx, connString)
+	if err != nil {
+		return C.CString(err.Error())
+	}
+	defer conn.Close(ctx)
+
+	sql := "INSERT INTO all_types (col_bigint, col_bool, col_bytea, col_float8, col_int, col_numeric, col_timestamptz, col_date, col_varchar, col_jsonb) " +
+		"values ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10) returning *"
+	numeric := pgtype.Numeric{}
+	_ = numeric.Set("6.626")
+	timestamptz, _ := time.Parse(time.RFC3339Nano, "2022-03-24T07:39:10.123456789+01:00")
+	date := pgtype.Date{}
+	_ = date.Set("2022-04-02")
+	var row pgx.Row
+	if strings.Contains(connString, "prefer_simple_protocol=true") {
+		// Simple mode will format the date as '2022-04-02 00:00:00Z', which is not supported by the
+		// backend yet.
+		row = conn.QueryRow(ctx, sql, 100, true, []byte("test_bytes"), 3.14, 1, numeric, timestamptz, "2022-04-02", "test_string", "{\"key\": \"value\"}")
+	} else {
+		row = conn.QueryRow(ctx, sql, 100, true, []byte("test_bytes"), 3.14, 1, numeric, timestamptz, date, "test_string", "{\"key\": \"value\"}")
+	}
+	var bigintValue int64
+	var boolValue bool
+	var byteaValue []byte
+	var float8Value float64
+	var intValue int
+	var numericValue pgtype.Numeric // pgx by default maps numeric to string
+	var timestamptzValue time.Time
+	var dateValue time.Time
+	var varcharValue string
+	var jsonbValue string
+
+	err = row.Scan(
+		&bigintValue,
+		&boolValue,
+		&byteaValue,
+		&float8Value,
+		&intValue,
+		&numericValue,
+		&timestamptzValue,
+		&dateValue,
+		&varcharValue,
+		&jsonbValue,
+	)
+	if err != nil {
+		return C.CString(fmt.Sprintf("Failed to execute insert: %v", err.Error()))
+	}
+	if g, w := bigintValue, int64(1); g != w {
+		return C.CString(fmt.Sprintf("value mismatch\n Got: %v\nWant: %v", g, w))
+	}
+	if g, w := boolValue, true; g != w {
+		return C.CString(fmt.Sprintf("value mismatch\n Got: %v\nWant: %v", g, w))
+	}
+	if g, w := byteaValue, []byte("test"); !reflect.DeepEqual(g, w) {
+		return C.CString(fmt.Sprintf("value mismatch\n Got: %v\nWant: %v", g, w))
+	}
+	if g, w := float8Value, 3.14; g != w {
+		return C.CString(fmt.Sprintf("value mismatch\n Got: %v\nWant: %v", g, w))
+	}
+	if g, w := intValue, 100; g != w {
+		return C.CString(fmt.Sprintf("value mismatch\n Got: %v\nWant: %v", g, w))
+	}
+	var wantNumericValue pgtype.Numeric
+	_ = wantNumericValue.Scan("6.626")
+	if g, w := numericValue, wantNumericValue; !reflect.DeepEqual(g, w) {
+		return C.CString(fmt.Sprintf("value mismatch\n Got: %v\nWant: %v", g, w))
+	}
+	wantDateValue, _ := time.Parse("2006-01-02", "2022-03-29")
+	if g, w := dateValue, wantDateValue; !reflect.DeepEqual(g, w) {
+		return C.CString(fmt.Sprintf("value mismatch\n Got: %v\nWant: %v", g, w))
+	}
+	// Encoding the timestamp values as a parameter will truncate it to microsecond precision.
+	wantTimestamptzValue, _ := time.Parse(time.RFC3339Nano, "2022-02-16T13:18:02.123456+00:00")
+	if g, w := timestamptzValue.UTC().String(), wantTimestamptzValue.UTC().String(); g != w {
+		return C.CString(fmt.Sprintf("value mismatch\n Got: %v\nWant: %v", g, w))
+	}
+	if g, w := varcharValue, "test"; g != w {
+		return C.CString(fmt.Sprintf("value mismatch\n Got: %v\nWant: %v", g, w))
+	}
+	if g, w := jsonbValue, "{\"key\": \"value\"}"; g != w {
+		return C.CString(fmt.Sprintf("value mismatch\n Got: %v\nWant: %v", g, w))
+	}
+
+	return nil
+}
+
 //export TestUpdateAllDataTypes
 func TestUpdateAllDataTypes(connString string) *C.char {
 	ctx := context.Background()
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java
index 60f6fc25..0676b60c 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java
@@ -42,6 +42,7 @@ import com.google.cloud.spanner.pgadapter.wireprotocol.DescribeMessage;
 import com.google.cloud.spanner.pgadapter.wireprotocol.ExecuteMessage;
 import com.google.cloud.spanner.pgadapter.wireprotocol.ParseMessage;
 import com.google.common.collect.ImmutableList;
+import com.google.protobuf.ListValue;
 import com.google.protobuf.Value;
 import com.google.spanner.admin.database.v1.UpdateDatabaseDdlRequest;
 import com.google.spanner.v1.BeginTransactionRequest;
@@ -66,6 +67,7 @@ import java.sql.DriverManager;
 import java.sql.ParameterMetaData;
 import java.sql.PreparedStatement;
 import java.sql.ResultSet;
+import java.sql.ResultSetMetaData;
 import java.sql.SQLException;
 import java.sql.Types;
 import java.time.LocalDate;
@@ -1620,6 +1622,123 @@ public class JdbcMockServerTest extends AbstractMockServerTest {
     }
   }
 
+  @Test
+  public void testPreparedStatementReturning() throws SQLException {
+    String pgSql =
+        "insert into all_types "
+            + "(col_bigint, col_bool, col_bytea, col_float8, col_int, col_numeric, col_timestamptz, col_date, col_varchar, col_jsonb) "
+            + "values ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10) "
+            + "returning *";
+    String sql =
+        "insert into all_types "
+            + "(col_bigint, col_bool, col_bytea, col_float8, col_int, col_numeric, col_timestamptz, col_date, col_varchar, col_jsonb) "
+            + "values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) "
+            + "returning *";
+    mockSpanner.putStatementResult(
+        StatementResult.query(
+            Statement.of(pgSql),
+            com.google.spanner.v1.ResultSet.newBuilder()
+                .setMetadata(
+                    ALL_TYPES_METADATA
+                        .toBuilder()
+                        .setUndeclaredParameters(
+                            createParameterTypesMetadata(
+                                    ImmutableList.of(
+                                        TypeCode.INT64,
+                                        TypeCode.BOOL,
+                                        TypeCode.BYTES,
+                                        TypeCode.FLOAT64,
+                                        TypeCode.INT64,
+                                        TypeCode.NUMERIC,
+                                        TypeCode.TIMESTAMP,
+                                        TypeCode.DATE,
+                                        TypeCode.STRING,
+                                        TypeCode.JSON))
+                                .getUndeclaredParameters())
+                        .build())
+                .setStats(ResultSetStats.newBuilder().build())
+                .build()));
+    mockSpanner.putStatementResult(
+        StatementResult.query(
+            Statement.newBuilder(pgSql)
+                .bind("p1")
+                .to(1L)
+                .bind("p2")
+                .to(true)
+                .bind("p3")
+                .to(ByteArray.copyFrom("test"))
+                .bind("p4")
+                .to(3.14d)
+                .bind("p5")
+                .to(100L)
+                .bind("p6")
+                .to(com.google.cloud.spanner.Value.pgNumeric("6.626"))
+                .bind("p7")
+                .to(Timestamp.parseTimestamp("2022-02-16T13:18:02.123457000Z"))
+                .bind("p8")
+                .to(Date.parseDate("2022-03-29"))
+                .bind("p9")
+                .to("test")
+                .bind("p10")
+                .to("{\"key\": \"value\"}")
+                .build(),
+            com.google.spanner.v1.ResultSet.newBuilder()
+                .setMetadata(ALL_TYPES_METADATA)
+                .setStats(ResultSetStats.newBuilder().setRowCountExact(1L).build())
+                .addRows(ALL_TYPES_RESULTSET.getRows(0))
+                .build()));
+
+    OffsetDateTime zonedDateTime =
+        LocalDateTime.of(2022, 2, 16, 13, 18, 2, 123456789).atOffset(ZoneOffset.UTC);
+    try (Connection connection = DriverManager.getConnection(createUrl())) {
+      try (PreparedStatement statement = connection.prepareStatement(sql)) {
+        ParameterMetaData parameterMetaData = statement.getParameterMetaData();
+        assertEquals(10, parameterMetaData.getParameterCount());
+        assertEquals(Types.BIGINT, parameterMetaData.getParameterType(1));
+        assertEquals(Types.BIT, parameterMetaData.getParameterType(2));
+        assertEquals(Types.BINARY, parameterMetaData.getParameterType(3));
+        assertEquals(Types.DOUBLE, parameterMetaData.getParameterType(4));
+        assertEquals(Types.BIGINT, parameterMetaData.getParameterType(5));
+        assertEquals(Types.NUMERIC, parameterMetaData.getParameterType(6));
+        assertEquals(Types.TIMESTAMP, parameterMetaData.getParameterType(7));
+        assertEquals(Types.DATE, parameterMetaData.getParameterType(8));
+        assertEquals(Types.VARCHAR, parameterMetaData.getParameterType(9));
+        // TODO: Enable when support for JSONB has been enabled.
+        // assertEquals(Types.OTHER, parameterMetaData.getParameterType(10));
+        ResultSetMetaData metadata = statement.getMetaData();
+        assertEquals(10, metadata.getColumnCount());
+        assertEquals(Types.BIGINT, metadata.getColumnType(1));
+        assertEquals(Types.BIT, metadata.getColumnType(2));
+        assertEquals(Types.BINARY, metadata.getColumnType(3));
+        assertEquals(Types.DOUBLE, metadata.getColumnType(4));
+        assertEquals(Types.BIGINT, metadata.getColumnType(5));
+        assertEquals(Types.NUMERIC, metadata.getColumnType(6));
+        assertEquals(Types.TIMESTAMP, metadata.getColumnType(7));
+        assertEquals(Types.DATE, metadata.getColumnType(8));
+        assertEquals(Types.VARCHAR, metadata.getColumnType(9));
+        // TODO: Enable when support for JSONB has been enabled.
+        // assertEquals(Types.OTHER, metadata.getColumnType(10));
+
+        int index = 0;
+        statement.setLong(++index, 1L);
+        statement.setBoolean(++index, true);
+        statement.setBytes(++index, "test".getBytes(StandardCharsets.UTF_8));
+        statement.setDouble(++index, 3.14d);
+        statement.setInt(++index, 100);
+        statement.setBigDecimal(++index, new BigDecimal("6.626"));
+        statement.setObject(++index, zonedDateTime);
+        statement.setObject(++index, LocalDate.of(2022, 3, 29));
+        statement.setString(++index, "test");
+        statement.setString(++index, "{\"key\": \"value\"}");
+
+        try (ResultSet resultSet = statement.executeQuery()) {
+          assertTrue(resultSet.next());
+          assertFalse(resultSet.next());
+        }
+      }
+    }
+  }
+
   @Test
   public void testCursorSuccess() throws SQLException {
     try (Connection connection = DriverManager.getConnection(createUrl())) {
@@ -2939,6 +3058,86 @@ public class JdbcMockServerTest extends AbstractMockServerTest {
     }
   }
 
+  @Test
+  public void testDmlReturning() throws SQLException {
+    String sql = "INSERT INTO test (value) values ('test') RETURNING id";
+    mockSpanner.putStatementResult(
+        StatementResult.query(
+            Statement.of(sql),
+            com.google.spanner.v1.ResultSet.newBuilder()
+                .setMetadata(
+                    createMetadata(ImmutableList.of(TypeCode.INT64), ImmutableList.of("id")))
+                .setStats(ResultSetStats.newBuilder().setRowCountExact(1L).build())
+                .addRows(
+                    ListValue.newBuilder()
+                        .addValues(Value.newBuilder().setStringValue("9999").build())
+                        .build())
+                .build()));
+
+    try (Connection connection = DriverManager.getConnection(createUrl())) {
+      try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) {
+        assertTrue(resultSet.next());
+        assertEquals(9999L, resultSet.getLong(1));
+        assertFalse(resultSet.next());
+      }
+      try (java.sql.Statement statement = connection.createStatement()) {
+        assertTrue(statement.execute(sql));
+        try (ResultSet resultSet = statement.getResultSet()) {
+          assertTrue(resultSet.next());
+          assertEquals(9999L, resultSet.getLong(1));
+          assertFalse(resultSet.next());
+        }
+        assertFalse(statement.getMoreResults());
+        assertEquals(-1, statement.getUpdateCount());
+      }
+    }
+  }
+
+  @Test
+  public void testDmlReturningMultipleRows() throws SQLException {
+    String sql = "UPDATE test SET value='new_value' WHERE value='old_value' RETURNING id";
+    mockSpanner.putStatementResult(
+        StatementResult.query(
+            Statement.of(sql),
+            com.google.spanner.v1.ResultSet.newBuilder()
+                .setMetadata(
+                    createMetadata(ImmutableList.of(TypeCode.INT64), ImmutableList.of("id")))
+                .setStats(ResultSetStats.newBuilder().setRowCountExact(3L).build())
+                .addRows(
+                    ListValue.newBuilder()
+                        .addValues(Value.newBuilder().setStringValue("1").build())
+                        .build())
+                .addRows(
+                    ListValue.newBuilder()
+                        .addValues(Value.newBuilder().setStringValue("2").build())
+                        .build())
+                .addRows(
+                    ListValue.newBuilder()
+                        .addValues(Value.newBuilder().setStringValue("3").build())
+                        .build())
+                .build()));
+
+    try (Connection connection = DriverManager.getConnection(createUrl())) {
+      try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) {
+        assertTrue(resultSet.next());
+        assertEquals(1L, resultSet.getLong(1));
+        assertTrue(resultSet.next());
+        assertEquals(2L, resultSet.getLong(1));
+        assertTrue(resultSet.next());
+        assertEquals(3L, resultSet.getLong(1));
+        assertFalse(resultSet.next());
+      }
+    }
+
+    assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
+    ExecuteSqlRequest executeRequest =
+        mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0);
+    assertEquals(QueryMode.NORMAL, executeRequest.getQueryMode());
+    assertTrue(executeRequest.getTransaction().hasBegin());
+    assertTrue(executeRequest.getTransaction().getBegin().hasReadWrite());
+    assertEquals(1, mockSpanner.countRequestsOfType(CommitRequest.class));
+  }
+
   private void verifySettingIsNull(Connection connection, String setting) throws SQLException {
     try (ResultSet resultSet =
         connection.createStatement().executeQuery(String.format("show %s", setting))) {
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/golang/PgxMockServerTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/golang/PgxMockServerTest.java
index bb3c157a..410fd207 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/golang/PgxMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/golang/PgxMockServerTest.java
@@ -401,6 +401,98 @@ public class PgxMockServerTest extends AbstractMockServerTest {
     assertEquals(QueryMode.NORMAL, executeRequest.getQueryMode());
   }
 
+  @Test
+  public void testInsertAllDataTypesReturning() {
+    String sql =
+        "INSERT INTO all_types "
+            + "(col_bigint, col_bool, col_bytea, col_float8, col_int, col_numeric, col_timestamptz, col_date, col_varchar, col_jsonb) "
+            + "values ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10) returning *";
+    mockSpanner.putStatementResult(
+        StatementResult.query(
+            Statement.of(sql),
+            ResultSet.newBuilder()
+                .setMetadata(
+                    ALL_TYPES_METADATA
+                        .toBuilder()
+                        .setUndeclaredParameters(
+                            createParameterTypesMetadata(
+                                    ImmutableList.of(
+                                        TypeCode.INT64,
+                                        TypeCode.BOOL,
+                                        TypeCode.BYTES,
+                                        TypeCode.FLOAT64,
+                                        TypeCode.INT64,
+                                        TypeCode.NUMERIC,
+                                        TypeCode.TIMESTAMP,
+                                        TypeCode.DATE,
+                                        TypeCode.STRING,
+                                        TypeCode.STRING))
+                                .getUndeclaredParameters()))
+                .setStats(ResultSetStats.newBuilder().build())
+                .build()));
+    mockSpanner.putStatementResult(
+        StatementResult.query(
+            Statement.newBuilder(sql)
+                .bind("p1")
+                .to(100L)
+                .bind("p2")
+                .to(true)
+                .bind("p3")
+                .to(ByteArray.copyFrom("test_bytes"))
+                .bind("p4")
+                .to(3.14d)
+                .bind("p5")
+                .to(1L)
+                .bind("p6")
+                .to(com.google.cloud.spanner.Value.pgNumeric("6.626"))
+                .bind("p7")
+                .to(Timestamp.parseTimestamp("2022-03-24T06:39:10.123456000Z"))
+                .bind("p8")
+                .to(Date.parseDate("2022-04-02"))
+                .bind("p9")
+                .to("test_string")
+                .bind("p10")
+                .to("{\"key\": \"value\"}")
+                .build(),
+            ResultSet.newBuilder()
+                .setMetadata(
+                    ALL_TYPES_METADATA
+                        .toBuilder()
+                        .setUndeclaredParameters(
+                            createParameterTypesMetadata(
+                                    ImmutableList.of(
+                                        TypeCode.INT64,
+                                        TypeCode.BOOL,
+                                        TypeCode.BYTES,
+                                        TypeCode.FLOAT64,
+                                        TypeCode.INT64,
+                                        TypeCode.NUMERIC,
+                                        TypeCode.TIMESTAMP,
+                                        TypeCode.DATE,
+                                        TypeCode.STRING,
+                                        TypeCode.STRING))
+                                .getUndeclaredParameters()))
+                .setStats(ResultSetStats.newBuilder().setRowCountExact(1L).build())
+                .addRows(ALL_TYPES_RESULTSET.getRows(0))
+                .build()));
+
+    String res = pgxTest.TestInsertAllDataTypesReturning(createConnString());
+
+    assertNull(res);
+    List<ExecuteSqlRequest> requests = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class);
+    // pgx by default always uses prepared statements. That means that each request is sent two
+    // times to the backend the first time it is executed:
+    // 1. DescribeStatement (parameters)
+    // 2. Execute
+    assertEquals(2, requests.size());
+    ExecuteSqlRequest describeParamsRequest = requests.get(0);
+    assertEquals(sql, describeParamsRequest.getSql());
+    assertEquals(QueryMode.PLAN, describeParamsRequest.getQueryMode());
+    ExecuteSqlRequest executeRequest = requests.get(1);
+    assertEquals(sql, executeRequest.getSql());
+    assertEquals(QueryMode.NORMAL, executeRequest.getQueryMode());
+  }
+
   @Test
   public void testInsertBatch() {
     String sql =
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/golang/PgxTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/golang/PgxTest.java
index efd9f97a..3521d6cd 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/golang/PgxTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/golang/PgxTest.java
@@ -31,6 +31,8 @@ public interface PgxTest extends Library {
 
   String TestInsertNullsAllDataTypes(GoString connString);
 
+  String TestInsertAllDataTypesReturning(GoString connString);
+
   String TestUpdateAllDataTypes(GoString connString);
 
   String TestPrepareStatement(GoString connString);