Repository: googleapis/java-spanner-jdbc
Pull request #1310 (2023-08-11 17:22:15)

feat: return generated keys
Adds support for returning generated keys.


Patch:
diff --git a/src/main/java/com/google/cloud/spanner/jdbc/AbstractJdbcStatement.java b/src/main/java/com/google/cloud/spanner/jdbc/AbstractJdbcStatement.java
index 2ee68af..e7b1a47 100644
--- a/src/main/java/com/google/cloud/spanner/jdbc/AbstractJdbcStatement.java
+++ b/src/main/java/com/google/cloud/spanner/jdbc/AbstractJdbcStatement.java
@@ -24,7 +24,6 @@ import com.google.cloud.spanner.connection.AbstractStatementParser;
 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.rpc.Code;
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.SQLWarning;
@@ -201,43 +200,6 @@ abstract class AbstractJdbcStatement extends AbstractJdbcWrapper implements Stat
     }
   }
 
-  /**
-   * Executes a SQL statement on the connection of this {@link Statement} as an update (DML)
-   * statement.
-   *
-   * @param statement The SQL statement to execute
-   * @return the number of rows that was inserted/updated/deleted
-   * @throws SQLException if a database error occurs, or if the number of rows affected is larger
-   *     than {@link Integer#MAX_VALUE}
-   */
-  int executeUpdate(com.google.cloud.spanner.Statement statement) throws SQLException {
-    long count = executeLargeUpdate(statement);
-    if (count > Integer.MAX_VALUE) {
-      throw JdbcSqlExceptionFactory.of(
-          "update count too large for executeUpdate: " + count, Code.OUT_OF_RANGE);
-    }
-    return (int) count;
-  }
-
-  /**
-   * Executes a SQL statement on the connection of this {@link Statement} as an update (DML)
-   * statement.
-   *
-   * @param statement The SQL statement to execute
-   * @return the number of rows that was inserted/updated/deleted
-   * @throws SQLException if a database error occurs
-   */
-  long executeLargeUpdate(com.google.cloud.spanner.Statement statement) throws SQLException {
-    StatementTimeout originalTimeout = setTemporaryStatementTimeout();
-    try {
-      return connection.getSpannerConnection().executeUpdate(statement);
-    } catch (SpannerException e) {
-      throw JdbcSqlExceptionFactory.of(e);
-    } finally {
-      resetStatementTimeout(originalTimeout);
-    }
-  }
-
   /**
    * Executes a SQL statement on the connection of this {@link Statement}. The SQL statement can be
    * any supported SQL statement, including client side statements such as SET AUTOCOMMIT ON|OFF.
diff --git a/src/main/java/com/google/cloud/spanner/jdbc/JdbcConnection.java b/src/main/java/com/google/cloud/spanner/jdbc/JdbcConnection.java
index a710ac8..6d6bd80 100644
--- a/src/main/java/com/google/cloud/spanner/jdbc/JdbcConnection.java
+++ b/src/main/java/com/google/cloud/spanner/jdbc/JdbcConnection.java
@@ -16,6 +16,9 @@
 
 package com.google.cloud.spanner.jdbc;
 
+import static com.google.cloud.spanner.jdbc.JdbcStatement.ALL_COLUMNS;
+import static com.google.cloud.spanner.jdbc.JdbcStatement.isNullOrEmpty;
+
 import com.google.api.client.util.Preconditions;
 import com.google.cloud.spanner.CommitResponse;
 import com.google.cloud.spanner.Mutation;
@@ -25,6 +28,7 @@ import com.google.cloud.spanner.connection.AutocommitDmlMode;
 import com.google.cloud.spanner.connection.ConnectionOptions;
 import com.google.cloud.spanner.connection.SavepointSupport;
 import com.google.cloud.spanner.connection.TransactionMode;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterators;
 import java.sql.Array;
 import java.sql.Blob;
@@ -49,9 +53,10 @@ class JdbcConnection extends AbstractJdbcConnection {
       "Only result sets with concurrency CONCUR_READ_ONLY are supported";
   private static final String ONLY_CLOSE_CURSORS_AT_COMMIT =
       "Only result sets with holdability CLOSE_CURSORS_AT_COMMIT are supported";
-  static final String ONLY_NO_GENERATED_KEYS = "Only NO_GENERATED_KEYS are supported";
   static final String IS_VALID_QUERY = "SELECT 1";
 
+  static final ImmutableList<String> NO_GENERATED_KEY_COLUMNS = ImmutableList.of();
+
   private Map<String, Class<?>> typeMap = new HashMap<>();
 
   JdbcConnection(String connectionUrl, ConnectionOptions options) throws SQLException {
@@ -66,8 +71,13 @@ class JdbcConnection extends AbstractJdbcConnection {
 
   @Override
   public JdbcPreparedStatement prepareStatement(String sql) throws SQLException {
+    return prepareStatement(sql, NO_GENERATED_KEY_COLUMNS);
+  }
+
+  private JdbcPreparedStatement prepareStatement(
+      String sql, ImmutableList<String> generatedKeyColumns) throws SQLException {
     checkClosed();
-    return new JdbcPreparedStatement(this, sql);
+    return new JdbcPreparedStatement(this, sql, generatedKeyColumns);
   }
 
   @Override
@@ -299,22 +309,26 @@ class JdbcConnection extends AbstractJdbcConnection {
 
   @Override
   public PreparedStatement prepareStatement(String sql, int autoGeneratedKeys) throws SQLException {
-    checkClosed();
-    JdbcPreconditions.checkSqlFeatureSupported(
-        autoGeneratedKeys == Statement.NO_GENERATED_KEYS, ONLY_NO_GENERATED_KEYS);
-    return prepareStatement(sql);
+    return prepareStatement(
+        sql,
+        autoGeneratedKeys == Statement.RETURN_GENERATED_KEYS
+            ? ALL_COLUMNS
+            : NO_GENERATED_KEY_COLUMNS);
   }
 
   @Override
   public PreparedStatement prepareStatement(String sql, int[] columnIndexes) throws SQLException {
-    checkClosed();
-    return prepareStatement(sql);
+    // This should preferably have returned an error, but the initial version of the driver just
+    // accepted and ignored this. Starting to throw an error now would be a breaking change.
+    // TODO: Consider throwing an Unsupported error for the next major version bump.
+    return prepareStatement(sql, NO_GENERATED_KEY_COLUMNS);
   }
 
   @Override
   public PreparedStatement prepareStatement(String sql, String[] columnNames) throws SQLException {
-    checkClosed();
-    return prepareStatement(sql);
+    return prepareStatement(
+        sql,
+        isNullOrEmpty(columnNames) ? NO_GENERATED_KEY_COLUMNS : ImmutableList.copyOf(columnNames));
   }
 
   @Override
diff --git a/src/main/java/com/google/cloud/spanner/jdbc/JdbcPreparedStatement.java b/src/main/java/com/google/cloud/spanner/jdbc/JdbcPreparedStatement.java
index a1b327b..3ab302c 100644
--- a/src/main/java/com/google/cloud/spanner/jdbc/JdbcPreparedStatement.java
+++ b/src/main/java/com/google/cloud/spanner/jdbc/JdbcPreparedStatement.java
@@ -24,7 +24,9 @@ import com.google.cloud.spanner.Statement;
 import com.google.cloud.spanner.Type;
 import com.google.cloud.spanner.connection.AbstractStatementParser.ParametersInfo;
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
+import com.google.rpc.Code;
 import java.sql.PreparedStatement;
 import java.sql.ResultSet;
 import java.sql.ResultSetMetaData;
@@ -36,8 +38,11 @@ class JdbcPreparedStatement extends AbstractJdbcPreparedStatement {
   private final String sql;
   private final String sqlWithoutComments;
   private final ParametersInfo parameters;
+  private final ImmutableList<String> generatedKeysColumns;
 
-  JdbcPreparedStatement(JdbcConnection connection, String sql) throws SQLException {
+  JdbcPreparedStatement(
+      JdbcConnection connection, String sql, ImmutableList<String> generatedKeysColumns)
+      throws SQLException {
     super(connection);
     this.sql = sql;
     try {
@@ -47,6 +52,7 @@ class JdbcPreparedStatement extends AbstractJdbcPreparedStatement {
     } catch (SpannerException e) {
       throw JdbcSqlExceptionFactory.of(e);
     }
+    this.generatedKeysColumns = Preconditions.checkNotNull(generatedKeysColumns);
   }
 
   ParametersInfo getParametersInfo() {
@@ -76,19 +82,22 @@ class JdbcPreparedStatement extends AbstractJdbcPreparedStatement {
 
   @Override
   public int executeUpdate() throws SQLException {
-    checkClosed();
-    return executeUpdate(createStatement());
+    long count = executeLargeUpdate(createStatement(), generatedKeysColumns);
+    if (count > Integer.MAX_VALUE) {
+      throw JdbcSqlExceptionFactory.of(
+          "update count too large for executeUpdate: " + count, Code.OUT_OF_RANGE);
+    }
+    return (int) count;
   }
 
+  @Override
   public long executeLargeUpdate() throws SQLException {
-    checkClosed();
-    return executeLargeUpdate(createStatement());
+    return executeLargeUpdate(createStatement(), generatedKeysColumns);
   }
 
   @Override
   public boolean execute() throws SQLException {
-    checkClosed();
-    return executeStatement(createStatement());
+    return executeStatement(createStatement(), generatedKeysColumns);
   }
 
   @Override
diff --git a/src/main/java/com/google/cloud/spanner/jdbc/JdbcStatement.java b/src/main/java/com/google/cloud/spanner/jdbc/JdbcStatement.java
index 3385ef0..6c6a45a 100644
--- a/src/main/java/com/google/cloud/spanner/jdbc/JdbcStatement.java
+++ b/src/main/java/com/google/cloud/spanner/jdbc/JdbcStatement.java
@@ -16,6 +16,8 @@
 
 package com.google.cloud.spanner.jdbc;
 
+import static com.google.cloud.spanner.jdbc.JdbcConnection.NO_GENERATED_KEY_COLUMNS;
+
 import com.google.cloud.spanner.Options;
 import com.google.cloud.spanner.ResultSets;
 import com.google.cloud.spanner.SpannerBatchUpdateException;
@@ -36,7 +38,6 @@ import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import java.util.stream.Collectors;
-import javax.annotation.Nullable;
 
 /** Implementation of {@link java.sql.Statement} for Google Cloud Spanner. */
 class JdbcStatement extends AbstractJdbcStatement {
@@ -49,6 +50,7 @@ class JdbcStatement extends AbstractJdbcStatement {
   }
 
   private ResultSet currentResultSet;
+  private ResultSet currentGeneratedKeys;
   private long currentUpdateCount;
   private int fetchSize;
   private BatchType currentBatchType = BatchType.NONE;
@@ -72,7 +74,12 @@ class JdbcStatement extends AbstractJdbcStatement {
    */
   @Override
   public int executeUpdate(String sql) throws SQLException {
-    long result = executeLargeUpdate(sql);
+    return executeUpdate(sql, NO_GENERATED_KEY_COLUMNS);
+  }
+
+  private int executeUpdate(String sql, ImmutableList<String> generatedKeysColumns)
+      throws SQLException {
+    long result = executeLargeUpdate(sql, generatedKeysColumns);
     if (result > Integer.MAX_VALUE) {
       throw JdbcSqlExceptionFactory.of("update count too large: " + result, Code.OUT_OF_RANGE);
     }
@@ -87,14 +94,33 @@ class JdbcStatement extends AbstractJdbcStatement {
    */
   @Override
   public long executeLargeUpdate(String sql) throws SQLException {
+    return executeLargeUpdate(sql, NO_GENERATED_KEY_COLUMNS);
+  }
+
+  private long executeLargeUpdate(String sql, ImmutableList<String> generatedKeysColumns)
+      throws SQLException {
+    return executeLargeUpdate(Statement.of(sql), generatedKeysColumns);
+  }
+
+  protected long executeLargeUpdate(Statement statement, ImmutableList<String> generatedKeysColumns)
+      throws SQLException {
+    Preconditions.checkNotNull(generatedKeysColumns);
     checkClosed();
-    Statement statement = Statement.of(sql);
-    StatementResult result = execute(statement);
+    Statement statementWithReturningClause =
+        addReturningToStatement(statement, generatedKeysColumns);
+    StatementResult result = execute(statementWithReturningClause);
     switch (result.getResultType()) {
       case RESULT_SET:
-        // Close the result set as we are not going to return it to the user. This prevents the
-        // underlying session from potentially being leaked.
-        throw closeResultSetAndCreateInvalidQueryException(result);
+        if (generatedKeysColumns.isEmpty()) {
+          // Close the result set as we are not going to return it to the user. This prevents the
+          // underlying session from potentially being leaked.
+          throw closeResultSetAndCreateInvalidQueryException(result);
+        }
+        // Make a copy of the result set as it does not matter if the user does not close the result
+        // set. This also consumes all rows of the result set, which again means that it is safe to
+        // extract the update count.
+        this.currentGeneratedKeys = JdbcResultSet.copyOf(result.getResultSet());
+        return extractUpdateCountAndClose(result.getResultSet());
       case UPDATE_COUNT:
         return result.getUpdateCount();
       case NO_RESULT:
@@ -105,6 +131,35 @@ class JdbcStatement extends AbstractJdbcStatement {
     }
   }
 
+  /**
+   * Extracts the update count from the given result set and then closes the result set. This method
+   * may only be called for a {@link com.google.cloud.spanner.ResultSet} where all rows have been
+   * fetched. That is; {@link com.google.cloud.spanner.ResultSet#next()} must have returned false.
+   */
+  private long extractUpdateCountAndClose(com.google.cloud.spanner.ResultSet resultSet)
+      throws SQLException {
+    try {
+      if (resultSet.getStats() == null) {
+        throw JdbcSqlExceptionFactory.of(
+            "Result does not contain any stats", Code.FAILED_PRECONDITION);
+      }
+      long updateCount;
+      if (resultSet.getStats().hasRowCountExact()) {
+        updateCount = resultSet.getStats().getRowCountExact();
+      } else if (resultSet.getStats().hasRowCountLowerBound()) {
+        // This is returned by Cloud Spanner if the user had set the autocommit_dml_mode to
+        // 'partitioned_non_atomic' (i.e. PDML).
+        updateCount = resultSet.getStats().getRowCountLowerBound();
+      } else {
+        throw JdbcSqlExceptionFactory.of(
+            "Result does not contain an update count", Code.FAILED_PRECONDITION);
+      }
+      return updateCount;
+    } finally {
+      resultSet.close();
+    }
+  }
+
   private SQLException closeResultSetAndCreateInvalidQueryException(StatementResult result) {
     //noinspection finally
     try {
@@ -121,15 +176,14 @@ class JdbcStatement extends AbstractJdbcStatement {
    * met:
    *
    * <ol>
-   *   <li>The generatedKeysColumns is not null or empty
+   *   <li>The generatedKeysColumns is not empty
    *   <li>The statement is a DML statement
    *   <li>The DML statement does not already contain a THEN RETURN/RETURNING clause
    * </ol>
    */
-  Statement addReturningToStatement(
-      Statement statement, @Nullable ImmutableList<String> generatedKeysColumns)
+  Statement addReturningToStatement(Statement statement, ImmutableList<String> generatedKeysColumns)
       throws SQLException {
-    if (generatedKeysColumns == null || generatedKeysColumns.isEmpty()) {
+    if (Preconditions.checkNotNull(generatedKeysColumns).isEmpty()) {
       return statement;
     }
     // Check if the statement is a DML statement or not.
@@ -193,17 +247,30 @@ class JdbcStatement extends AbstractJdbcStatement {
 
   @Override
   public boolean execute(String sql) throws SQLException {
-    checkClosed();
-    return executeStatement(Statement.of(sql));
+    return executeStatement(Statement.of(sql), NO_GENERATED_KEY_COLUMNS);
   }
 
-  boolean executeStatement(Statement statement) throws SQLException {
-    StatementResult result = execute(statement);
+  boolean executeStatement(Statement statement, ImmutableList<String> generatedKeysColumns)
+      throws SQLException {
+    checkClosed();
+    // This will return the same Statement instance if no THEN RETURN clause is added to the
+    // statement.
+    Statement statementWithReturning = addReturningToStatement(statement, generatedKeysColumns);
+    StatementResult result = execute(statementWithReturning);
     switch (result.getResultType()) {
       case RESULT_SET:
-        currentResultSet = JdbcResultSet.of(this, result.getResultSet());
-        currentUpdateCount = JdbcConstants.STATEMENT_RESULT_SET;
-        return true;
+        // Check whether the statement was modified to include a RETURNING clause for generated
+        // keys. If so, then we return the result as an update count and the rows as the generated
+        // keys. We can safely use '==', as the addReturningToStatement(..) method returns the same
+        // instance if no generated keys were requested.
+        if (statementWithReturning == statement) {
+          currentResultSet = JdbcResultSet.of(this, result.getResultSet());
+          currentUpdateCount = JdbcConstants.STATEMENT_RESULT_SET;
+          return true;
+        }
+        this.currentGeneratedKeys = JdbcResultSet.copyOf(result.getResultSet());
+        this.currentUpdateCount = extractUpdateCountAndClose(result.getResultSet());
+        return false;
       case UPDATE_COUNT:
         currentResultSet = null;
         currentUpdateCount = result.getUpdateCount();
@@ -441,77 +508,94 @@ class JdbcStatement extends AbstractJdbcStatement {
   @Override
   public ResultSet getGeneratedKeys() throws SQLException {
     checkClosed();
-    // Return an empty result set instead of throwing an exception, to facilitate any application
-    // that might not check on beforehand whether the driver supports any generated keys.
-    com.google.cloud.spanner.ResultSet rs =
-        ResultSets.forRows(
-            Type.struct(
-                StructField.of("COLUMN_NAME", Type.string()),
-                StructField.of("VALUE", Type.int64())),
-            Collections.emptyList());
-    return JdbcResultSet.of(rs);
+    if (this.currentGeneratedKeys == null) {
+      // Return an empty result set instead of throwing an exception, as that is what the JDBC spec
+      // says we should do. Note that we need to create a new instance every time, as users could in
+      // theory call close() on the returned result set.
+      this.currentGeneratedKeys =
+          JdbcResultSet.of(
+              ResultSets.forRows(
+                  Type.struct(
+                      StructField.of("COLUMN_NAME", Type.string()),
+                      StructField.of("VALUE", Type.int64())),
+                  Collections.emptyList()));
+    }
+    return this.currentGeneratedKeys;
   }
 
   @Override
   public int executeUpdate(String sql, int autoGeneratedKeys) throws SQLException {
-    checkClosed();
-    JdbcPreconditions.checkSqlFeatureSupported(
-        autoGeneratedKeys == java.sql.Statement.NO_GENERATED_KEYS,
-        JdbcConnection.ONLY_NO_GENERATED_KEYS);
-    return executeUpdate(sql);
+    return executeUpdate(
+        sql,
+        autoGeneratedKeys == java.sql.Statement.RETURN_GENERATED_KEYS
+            ? ALL_COLUMNS
+            : NO_GENERATED_KEY_COLUMNS);
   }
 
   @Override
   public int executeUpdate(String sql, int[] columnIndexes) throws SQLException {
-    checkClosed();
-    return executeUpdate(sql);
+    // This should preferably have returned an error, but the initial version of the driver just
+    // accepted and ignored this. Starting to throw an error now would be a breaking change.
+    // TODO: Consider throwing an Unsupported error for the next major version bump.
+    return executeUpdate(sql, NO_GENERATED_KEY_COLUMNS);
   }
 
   @Override
   public int executeUpdate(String sql, String[] columnNames) throws SQLException {
-    checkClosed();
-    return executeUpdate(sql);
+    return executeUpdate(
+        sql,
+        isNullOrEmpty(columnNames) ? NO_GENERATED_KEY_COLUMNS : ImmutableList.copyOf(columnNames));
   }
 
   @Override
   public long executeLargeUpdate(String sql, int autoGeneratedKeys) throws SQLException {
-    checkClosed();
-    JdbcPreconditions.checkSqlFeatureSupported(
-        autoGeneratedKeys == java.sql.Statement.NO_GENERATED_KEYS,
-        JdbcConnection.ONLY_NO_GENERATED_KEYS);
-    return executeLargeUpdate(sql);
+    return executeLargeUpdate(
+        sql,
+        autoGeneratedKeys == java.sql.Statement.RETURN_GENERATED_KEYS
+            ? ALL_COLUMNS
+            : NO_GENERATED_KEY_COLUMNS);
   }
 
   @Override
   public long executeLargeUpdate(String sql, int[] columnIndexes) throws SQLException {
-    checkClosed();
-    return executeLargeUpdate(sql);
+    // This should preferably have returned an error, but the initial version of the driver just
+    // accepted and ignored this. Starting to throw an error now would be a breaking change.
+    // TODO: Consider throwing an Unsupported error for the next major version bump.
+    return executeLargeUpdate(sql, NO_GENERATED_KEY_COLUMNS);
   }
 
   @Override
   public long executeLargeUpdate(String sql, String[] columnNames) throws SQLException {
-    checkClosed();
-    return executeLargeUpdate(sql);
+    return executeLargeUpdate(
+        sql,
+        isNullOrEmpty(columnNames) ? NO_GENERATED_KEY_COLUMNS : ImmutableList.copyOf(columnNames));
   }
 
   @Override
   public boolean execute(String sql, int autoGeneratedKeys) throws SQLException {
-    checkClosed();
-    JdbcPreconditions.checkSqlFeatureSupported(
-        autoGeneratedKeys == java.sql.Statement.NO_GENERATED_KEYS,
-        JdbcConnection.ONLY_NO_GENERATED_KEYS);
-    return execute(sql);
+    return executeStatement(
+        Statement.of(sql),
+        autoGeneratedKeys == java.sql.Statement.RETURN_GENERATED_KEYS
+            ? ALL_COLUMNS
+            : NO_GENERATED_KEY_COLUMNS);
   }
 
   @Override
   public boolean execute(String sql, int[] columnIndexes) throws SQLException {
-    checkClosed();
-    return execute(sql);
+    // This should preferably have returned an error, but the initial version of the driver just
+    // accepted and ignored this. Starting to throw an error now would be a breaking change.
+    // TODO: Consider throwing an Unsupported error for the next major version bump.
+    return executeStatement(Statement.of(sql), NO_GENERATED_KEY_COLUMNS);
   }
 
   @Override
   public boolean execute(String sql, String[] columnNames) throws SQLException {
-    checkClosed();
-    return execute(sql);
+    return executeStatement(
+        Statement.of(sql),
+        isNullOrEmpty(columnNames) ? NO_GENERATED_KEY_COLUMNS : ImmutableList.copyOf(columnNames));
+  }
+
+  static boolean isNullOrEmpty(String[] columnNames) {
+    return columnNames == null || columnNames.length == 0;
   }
 }


Test patch:
diff --git a/src/test/java/com/google/cloud/spanner/jdbc/JdbcConnectionTest.java b/src/test/java/com/google/cloud/spanner/jdbc/JdbcConnectionTest.java
index 8010c21..592d404 100644
--- a/src/test/java/com/google/cloud/spanner/jdbc/JdbcConnectionTest.java
+++ b/src/test/java/com/google/cloud/spanner/jdbc/JdbcConnectionTest.java
@@ -686,14 +686,12 @@ public class JdbcConnectionTest {
     try (JdbcConnection connection = createConnection(mockOptions())) {
       PreparedStatement statement =
           connection.prepareStatement(sql, java.sql.Statement.NO_GENERATED_KEYS);
-      ResultSet rs = statement.getGeneratedKeys();
-      assertThat(rs.next()).isFalse();
-      try {
-        statement = connection.prepareStatement(sql, java.sql.Statement.RETURN_GENERATED_KEYS);
-        fail("missing expected SQLFeatureNotSupportedException");
-      } catch (SQLFeatureNotSupportedException e) {
-        // ignore, this is the expected exception.
-      }
+      ResultSet generatedKeys = statement.getGeneratedKeys();
+      assertFalse(generatedKeys.next());
+
+      statement = connection.prepareStatement(sql, java.sql.Statement.RETURN_GENERATED_KEYS);
+      generatedKeys = statement.getGeneratedKeys();
+      assertFalse(generatedKeys.next());
     }
   }
 
diff --git a/src/test/java/com/google/cloud/spanner/jdbc/JdbcPreparedStatementTest.java b/src/test/java/com/google/cloud/spanner/jdbc/JdbcPreparedStatementTest.java
index 8f21e68..310d154 100644
--- a/src/test/java/com/google/cloud/spanner/jdbc/JdbcPreparedStatementTest.java
+++ b/src/test/java/com/google/cloud/spanner/jdbc/JdbcPreparedStatementTest.java
@@ -16,6 +16,7 @@
 
 package com.google.cloud.spanner.jdbc;
 
+import static com.google.cloud.spanner.jdbc.JdbcConnection.NO_GENERATED_KEY_COLUMNS;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertThrows;
@@ -142,7 +143,8 @@ public class JdbcPreparedStatementTest {
               Collections.singleton(com.google.cloud.Timestamp.ofTimeSecondsAndNanos(999L, 99))),
         }) {
 
-      try (JdbcPreparedStatement ps = new JdbcPreparedStatement(connection, sql)) {
+      try (JdbcPreparedStatement ps =
+          new JdbcPreparedStatement(connection, sql, NO_GENERATED_KEY_COLUMNS)) {
         ps.setObject(1, value);
         Statement statement = ps.createStatement();
         assertEquals(statement.getParameters().get("p1"), value);
@@ -157,7 +159,8 @@ public class JdbcPreparedStatementTest {
     String sql = generateSqlWithParameters(numberOfParams);
 
     JdbcConnection connection = createMockConnection();
-    try (JdbcPreparedStatement ps = new JdbcPreparedStatement(connection, sql)) {
+    try (JdbcPreparedStatement ps =
+        new JdbcPreparedStatement(connection, sql, NO_GENERATED_KEY_COLUMNS)) {
       ps.setArray(1, connection.createArrayOf("INT64", new Long[] {1L, 2L, 3L}));
       ps.setAsciiStream(2, new ByteArrayInputStream("TEST".getBytes()));
       ps.setAsciiStream(3, new ByteArrayInputStream("TEST".getBytes()), 4);
@@ -276,7 +279,8 @@ public class JdbcPreparedStatementTest {
   public void testSetNullValues() throws SQLException {
     final int numberOfParameters = 31;
     String sql = generateSqlWithParameters(numberOfParameters);
-    try (JdbcPreparedStatement ps = new JdbcPreparedStatement(createMockConnection(), sql)) {
+    try (JdbcPreparedStatement ps =
+        new JdbcPreparedStatement(createMockConnection(), sql, NO_GENERATED_KEY_COLUMNS)) {
       int index = 0;
       ps.setNull(++index, Types.BLOB);
       ps.setNull(++index, Types.NVARCHAR);
@@ -350,7 +354,8 @@ public class JdbcPreparedStatementTest {
                     .build()));
     when(connection.analyzeQuery(Statement.of(sql), QueryAnalyzeMode.PLAN)).thenReturn(rs);
     try (JdbcPreparedStatement ps =
-        new JdbcPreparedStatement(createMockConnection(connection), sql)) {
+        new JdbcPreparedStatement(
+            createMockConnection(connection), sql, NO_GENERATED_KEY_COLUMNS)) {
       ResultSetMetaData metadata = ps.getMetaData();
       assertEquals(4, metadata.getColumnCount());
       assertEquals("ID", metadata.getColumnLabel(1));
@@ -369,7 +374,9 @@ public class JdbcPreparedStatementTest {
     Connection connection = mock(Connection.class);
     try (JdbcPreparedStatement ps =
         new JdbcPreparedStatement(
-            createMockConnection(connection), "UPDATE FOO SET BAR=1 WHERE TRUE")) {
+            createMockConnection(connection),
+            "UPDATE FOO SET BAR=1 WHERE TRUE",
+            NO_GENERATED_KEY_COLUMNS)) {
       ResultSetMetaData metadata = ps.getMetaData();
       assertEquals(0, metadata.getColumnCount());
     }
@@ -381,7 +388,9 @@ public class JdbcPreparedStatementTest {
     SQLException sqlException =
         assertThrows(
             SQLException.class,
-            () -> new JdbcPreparedStatement(createMockConnection(mock(Connection.class)), sql));
+            () ->
+                new JdbcPreparedStatement(
+                    createMockConnection(mock(Connection.class)), sql, NO_GENERATED_KEY_COLUMNS));
     assertTrue(sqlException instanceof JdbcSqlException);
     JdbcSqlException jdbcSqlException = (JdbcSqlException) sqlException;
     assertEquals(
diff --git a/src/test/java/com/google/cloud/spanner/jdbc/JdbcStatementTest.java b/src/test/java/com/google/cloud/spanner/jdbc/JdbcStatementTest.java
index d445835..727e63a 100644
--- a/src/test/java/com/google/cloud/spanner/jdbc/JdbcStatementTest.java
+++ b/src/test/java/com/google/cloud/spanner/jdbc/JdbcStatementTest.java
@@ -16,6 +16,7 @@
 
 package com.google.cloud.spanner.jdbc;
 
+import static com.google.cloud.spanner.jdbc.JdbcConnection.NO_GENERATED_KEY_COLUMNS;
 import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
@@ -44,7 +45,6 @@ import com.google.common.collect.ImmutableList;
 import com.google.rpc.Code;
 import java.sql.ResultSet;
 import java.sql.SQLException;
-import java.sql.SQLFeatureNotSupportedException;
 import java.sql.Statement;
 import java.util.Arrays;
 import java.util.List;
@@ -101,6 +101,10 @@ public class JdbcStatementTest {
     when(updateResult.getResultType()).thenReturn(ResultType.UPDATE_COUNT);
     when(updateResult.getUpdateCount()).thenReturn(1L);
     when(spanner.execute(com.google.cloud.spanner.Statement.of(UPDATE))).thenReturn(updateResult);
+    when(spanner.execute(com.google.cloud.spanner.Statement.of(UPDATE + "\nTHEN RETURN *")))
+        .thenReturn(updateResult);
+    when(spanner.execute(com.google.cloud.spanner.Statement.of(UPDATE + "\nRETURNING *")))
+        .thenReturn(updateResult);
 
     StatementResult largeUpdateResult = mock(StatementResult.class);
     when(largeUpdateResult.getResultType()).thenReturn(ResultType.UPDATE_COUNT);
@@ -274,15 +278,14 @@ public class JdbcStatementTest {
 
   @Test
   public void testExecuteWithGeneratedKeys() throws SQLException {
-    Statement statement = createStatement();
-    assertThat(statement.execute(UPDATE, Statement.NO_GENERATED_KEYS)).isFalse();
-    ResultSet keys = statement.getGeneratedKeys();
-    assertThat(keys.next()).isFalse();
-    try {
+    try (Statement statement = createStatement()) {
+      assertFalse(statement.execute(UPDATE, Statement.NO_GENERATED_KEYS));
+      ResultSet keys = statement.getGeneratedKeys();
+      assertFalse(keys.next());
+
       statement.execute(UPDATE, Statement.RETURN_GENERATED_KEYS);
-      fail("missing expected exception");
-    } catch (SQLFeatureNotSupportedException e) {
-      // Ignore, this is the expected exception.
+      keys = statement.getGeneratedKeys();
+      assertFalse(keys.next());
     }
   }
 
@@ -357,12 +360,18 @@ public class JdbcStatementTest {
         com.google.cloud.spanner.Statement.of(UPDATE);
     com.google.cloud.spanner.Statement largeUpdateStatement =
         com.google.cloud.spanner.Statement.of(LARGE_UPDATE);
-    when(spannerConnection.executeUpdate(updateStatement)).thenReturn(1L);
-    when(spannerConnection.executeUpdate(largeUpdateStatement)).thenReturn(Integer.MAX_VALUE + 1L);
+    StatementResult updateResult = mock(StatementResult.class);
+    when(updateResult.getUpdateCount()).thenReturn(1L);
+    when(updateResult.getResultType()).thenReturn(ResultType.UPDATE_COUNT);
+    when(spannerConnection.execute(updateStatement)).thenReturn(updateResult);
+    StatementResult largeUpdateResult = mock(StatementResult.class);
+    when(largeUpdateResult.getUpdateCount()).thenReturn(Integer.MAX_VALUE + 1L);
+    when(largeUpdateResult.getResultType()).thenReturn(ResultType.UPDATE_COUNT);
+    when(spannerConnection.execute(largeUpdateStatement)).thenReturn(largeUpdateResult);
     try (JdbcStatement statement = new JdbcStatement(connection)) {
-      assertThat(statement.executeUpdate(updateStatement)).isEqualTo(1);
+      assertThat(statement.executeUpdate(UPDATE)).isEqualTo(1);
       try {
-        statement.executeUpdate(largeUpdateStatement);
+        statement.executeUpdate(LARGE_UPDATE);
         fail("missing expected exception");
       } catch (JdbcSqlExceptionImpl e) {
         assertThat(e.getCode()).isEqualTo(Code.OUT_OF_RANGE);
@@ -380,12 +389,17 @@ public class JdbcStatementTest {
         com.google.cloud.spanner.Statement.of(UPDATE);
     com.google.cloud.spanner.Statement largeUpdateStatement =
         com.google.cloud.spanner.Statement.of(LARGE_UPDATE);
-    when(spannerConnection.executeUpdate(updateStatement)).thenReturn(1L);
-    when(spannerConnection.executeUpdate(largeUpdateStatement)).thenReturn(Integer.MAX_VALUE + 1L);
+    StatementResult updateResult = mock(StatementResult.class);
+    when(updateResult.getUpdateCount()).thenReturn(1L);
+    when(updateResult.getResultType()).thenReturn(ResultType.UPDATE_COUNT);
+    when(spannerConnection.execute(updateStatement)).thenReturn(updateResult);
+    StatementResult largeUpdateResult = mock(StatementResult.class);
+    when(largeUpdateResult.getUpdateCount()).thenReturn(Integer.MAX_VALUE + 1L);
+    when(largeUpdateResult.getResultType()).thenReturn(ResultType.UPDATE_COUNT);
+    when(spannerConnection.execute(largeUpdateStatement)).thenReturn(largeUpdateResult);
     try (JdbcStatement statement = new JdbcStatement(connection)) {
-      assertThat(statement.executeLargeUpdate(updateStatement)).isEqualTo(1);
-      assertThat(statement.executeLargeUpdate(largeUpdateStatement))
-          .isEqualTo(Integer.MAX_VALUE + 1L);
+      assertThat(statement.executeLargeUpdate(UPDATE)).isEqualTo(1);
+      assertThat(statement.executeLargeUpdate(LARGE_UPDATE)).isEqualTo(Integer.MAX_VALUE + 1L);
     }
   }
 
@@ -448,15 +462,14 @@ public class JdbcStatementTest {
 
   @Test
   public void testExecuteUpdateWithGeneratedKeys() throws SQLException {
-    Statement statement = createStatement();
-    assertThat(statement.executeUpdate(UPDATE, Statement.NO_GENERATED_KEYS)).isEqualTo(1);
-    ResultSet keys = statement.getGeneratedKeys();
-    assertThat(keys.next()).isFalse();
-    try {
-      statement.executeUpdate(UPDATE, Statement.RETURN_GENERATED_KEYS);
-      fail("missing expected exception");
-    } catch (SQLFeatureNotSupportedException e) {
-      // Ignore, this is the expected exception.
+    try (Statement statement = createStatement()) {
+      assertEquals(1, statement.executeUpdate(UPDATE, Statement.NO_GENERATED_KEYS));
+      ResultSet keys = statement.getGeneratedKeys();
+      assertFalse(keys.next());
+
+      assertEquals(1, statement.executeUpdate(UPDATE, Statement.RETURN_GENERATED_KEYS));
+      keys = statement.getGeneratedKeys();
+      assertFalse(keys.next());
     }
   }
 
@@ -609,7 +622,8 @@ public class JdbcStatementTest {
     when(connection.getDialect()).thenReturn(dialect);
     when(connection.getParser()).thenReturn(AbstractStatementParser.getInstance(dialect));
     try (JdbcStatement statement = new JdbcStatement(connection)) {
-      assertAddReturningSame(statement, "insert into test (id, value) values (1, 'One')", null);
+      assertAddReturningSame(
+          statement, "insert into test (id, value) values (1, 'One')", NO_GENERATED_KEY_COLUMNS);
       assertAddReturningSame(
           statement, "insert into test (id, value) values (1, 'One')", ImmutableList.of());
       assertAddReturningEquals(
@@ -649,7 +663,8 @@ public class JdbcStatementTest {
       }
 
       // Update statements may also request generated keys.
-      assertAddReturningSame(statement, "update test set value='Two' where id=1", null);
+      assertAddReturningSame(
+          statement, "update test set value='Two' where id=1", NO_GENERATED_KEY_COLUMNS);
       assertAddReturningSame(
           statement, "update test set value='Two' where id=1", ImmutableList.of());
       assertAddReturningEquals(
@@ -681,7 +696,7 @@ public class JdbcStatementTest {
           ImmutableList.of("value"));
 
       // Delete statements may also request generated keys.
-      assertAddReturningSame(statement, "delete test where id=1", null);
+      assertAddReturningSame(statement, "delete test where id=1", NO_GENERATED_KEY_COLUMNS);
       assertAddReturningSame(statement, "delete test where id=1", ImmutableList.of());
       assertAddReturningEquals(
           statement,
diff --git a/src/test/java/com/google/cloud/spanner/jdbc/PreparedStatementReturnGeneratedKeysMockServerTest.java b/src/test/java/com/google/cloud/spanner/jdbc/PreparedStatementReturnGeneratedKeysMockServerTest.java
new file mode 100644
index 0000000..b88428b
--- /dev/null
+++ b/src/test/java/com/google/cloud/spanner/jdbc/PreparedStatementReturnGeneratedKeysMockServerTest.java
@@ -0,0 +1,322 @@
+/*
+ * Copyright 2023 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.spanner.jdbc;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+
+import com.google.cloud.spanner.Dialect;
+import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult;
+import com.google.cloud.spanner.Statement;
+import com.google.cloud.spanner.connection.AbstractMockServerTest;
+import com.google.cloud.spanner.connection.RandomResultSetGenerator;
+import com.google.cloud.spanner.connection.SpannerPool;
+import com.google.common.collect.ImmutableList;
+import com.google.longrunning.Operation;
+import com.google.protobuf.Any;
+import com.google.protobuf.Empty;
+import com.google.protobuf.ListValue;
+import com.google.protobuf.Value;
+import com.google.spanner.admin.database.v1.UpdateDatabaseDdlMetadata;
+import com.google.spanner.v1.ResultSetMetadata;
+import com.google.spanner.v1.ResultSetStats;
+import com.google.spanner.v1.StructType;
+import com.google.spanner.v1.StructType.Field;
+import com.google.spanner.v1.Type;
+import com.google.spanner.v1.TypeCode;
+import java.sql.Connection;
+import java.sql.DriverManager;
+import java.sql.PreparedStatement;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Types;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+import org.junit.runners.Parameterized.Parameter;
+import org.junit.runners.Parameterized.Parameters;
+
+@RunWith(Parameterized.class)
+public class PreparedStatementReturnGeneratedKeysMockServerTest extends AbstractMockServerTest {
+  @Parameter public Dialect dialect;
+
+  @Parameters(name = "dialect = {0}")
+  public static Object[] data() {
+    return Dialect.values();
+  }
+
+  @Before
+  public void setupDialect() {
+    mockSpanner.putStatementResult(StatementResult.detectDialectResult(this.dialect));
+  }
+
+  @After
+  public void clearRequests() {
+    mockSpanner.clearRequests();
+    // Close Spanner pool to force a refresh of the dialect of a given database.
+    SpannerPool.closeSpannerPool();
+  }
+
+  private String createUrl() {
+    return String.format(
+        "jdbc:cloudspanner://localhost:%d/projects/%s/instances/%s/databases/%s?usePlainText=true",
+        getPort(), "proj", "inst", "db");
+  }
+
+  private Connection createConnection() throws SQLException {
+    return DriverManager.getConnection(createUrl());
+  }
+
+  interface SqlRunnable {
+    void run() throws SQLException;
+  }
+
+  @Test
+  public void testReturnAutoGeneratedKeys() throws SQLException {
+    String sql = "insert into test (id, value) values (1, 'One')";
+    String returningClause = dialect == Dialect.POSTGRESQL ? "\nRETURNING *" : "\nTHEN RETURN *";
+    mockSpanner.putStatementResult(
+        StatementResult.query(
+            Statement.of(sql + returningClause),
+            com.google.spanner.v1.ResultSet.newBuilder()
+                .setMetadata(
+                    ResultSetMetadata.newBuilder()
+                        .setRowType(
+                            StructType.newBuilder()
+                                .addFields(
+                                    Field.newBuilder()
+                                        .setName("id")
+                                        .setType(Type.newBuilder().setCode(TypeCode.INT64).build())
+                                        .build())
+                                .addFields(
+                                    Field.newBuilder()
+                                        .setName("value")
+                                        .setType(Type.newBuilder().setCode(TypeCode.STRING).build())
+                                        .build())
+                                .build())
+                        .build())
+                .addRows(
+                    ListValue.newBuilder()
+                        .addValues(Value.newBuilder().setStringValue("1").build())
+                        .addValues(Value.newBuilder().setStringValue("One").build())
+                        .build())
+                .setStats(ResultSetStats.newBuilder().setRowCountExact(1L).build())
+                .build()));
+
+    try (Connection connection = createConnection();
+        PreparedStatement statement =
+            connection.prepareStatement(sql, java.sql.Statement.RETURN_GENERATED_KEYS)) {
+      for (SqlRunnable runnable :
+          ImmutableList.<SqlRunnable>of(
+              () -> assertEquals(1, statement.executeUpdate()),
+              () -> assertEquals(1L, statement.executeLargeUpdate()),
+              () -> {
+                assertFalse(statement.execute());
+                assertEquals(1, statement.getUpdateCount());
+                assertEquals(1L, statement.getLargeUpdateCount());
+              })) {
+        runnable.run();
+        assertNotNull(statement.getGeneratedKeys());
+        assertEquals(2, statement.getGeneratedKeys().getMetaData().getColumnCount());
+        assertEquals(Types.BIGINT, statement.getGeneratedKeys().getMetaData().getColumnType(1));
+        assertEquals(Types.NVARCHAR, statement.getGeneratedKeys().getMetaData().getColumnType(2));
+        assertEquals("id", statement.getGeneratedKeys().getMetaData().getColumnName(1));
+        assertEquals("value", statement.getGeneratedKeys().getMetaData().getColumnName(2));
+        try (ResultSet keys = statement.getGeneratedKeys()) {
+          assertTrue(keys.next());
+          assertEquals(1L, keys.getLong(1));
+          assertEquals("One", keys.getString(2));
+          assertFalse(keys.next());
+        }
+        // There should be no more results for this statement.
+        assertFalse(statement.getMoreResults());
+        assertEquals(-1, statement.getUpdateCount());
+      }
+    }
+  }
+
+  @Test
+  public void testReturnColumnNames() throws SQLException {
+    String sql = "insert into test (id, value) values (1, 'One')";
+    String returningClause =
+        dialect == Dialect.POSTGRESQL ? "\nRETURNING \"id\"" : "\nTHEN RETURN `id`";
+    mockSpanner.putStatementResult(
+        StatementResult.query(
+            Statement.of(sql + returningClause),
+            com.google.spanner.v1.ResultSet.newBuilder()
+                .setMetadata(
+                    ResultSetMetadata.newBuilder()
+                        .setRowType(
+                            StructType.newBuilder()
+                                .addFields(
+                                    Field.newBuilder()
+                                        .setName("id")
+                                        .setType(Type.newBuilder().setCode(TypeCode.INT64).build())
+                                        .build())
+                                .build())
+                        .build())
+                .addRows(
+                    ListValue.newBuilder()
+                        .addValues(Value.newBuilder().setStringValue("1").build())
+                        .build())
+                .setStats(ResultSetStats.newBuilder().setRowCountExact(1L).build())
+                .build()));
+
+    try (Connection connection = createConnection();
+        PreparedStatement statement = connection.prepareStatement(sql, new String[] {"id"})) {
+      for (SqlRunnable runnable :
+          ImmutableList.<SqlRunnable>of(
+              () -> assertEquals(1, statement.executeUpdate()),
+              () -> assertEquals(1L, statement.executeLargeUpdate()),
+              () -> {
+                assertFalse(statement.execute());
+                assertEquals(1, statement.getUpdateCount());
+                assertEquals(1L, statement.getLargeUpdateCount());
+              })) {
+        runnable.run();
+        assertNotNull(statement.getGeneratedKeys());
+        assertEquals(1, statement.getGeneratedKeys().getMetaData().getColumnCount());
+        assertEquals(Types.BIGINT, statement.getGeneratedKeys().getMetaData().getColumnType(1));
+        assertEquals("id", statement.getGeneratedKeys().getMetaData().getColumnName(1));
+        try (ResultSet keys = statement.getGeneratedKeys()) {
+          assertTrue(keys.next());
+          assertEquals(1L, keys.getLong(1));
+          assertFalse(keys.next());
+        }
+        // There should be no more results for this statement.
+        assertFalse(statement.getMoreResults());
+        assertEquals(-1, statement.getUpdateCount());
+      }
+    }
+  }
+
+  @Test
+  public void testReturnColumnIndices() throws SQLException {
+    String sql = "insert into test (id, value) values (1, 'One')";
+    mockSpanner.putStatementResult(StatementResult.update(Statement.of(sql), 1L));
+
+    try (Connection connection = createConnection();
+        PreparedStatement statement = connection.prepareStatement(sql, new int[] {1})) {
+      // The JDBC driver silently ignores the request for column indices.
+      for (SqlRunnable runnable :
+          ImmutableList.<SqlRunnable>of(
+              () -> assertEquals(1, statement.executeUpdate()),
+              () -> assertEquals(1L, statement.executeLargeUpdate()),
+              () -> {
+                assertFalse(statement.execute());
+                assertEquals(1, statement.getUpdateCount());
+                assertEquals(1L, statement.getLargeUpdateCount());
+              })) {
+        runnable.run();
+        assertNotNull(statement.getGeneratedKeys());
+        assertFalse(statement.getGeneratedKeys().next());
+      }
+    }
+  }
+
+  @Test
+  public void testReturnGeneratedKeysForQuery() throws SQLException {
+    String sql = "select * from test";
+    RandomResultSetGenerator generator = new RandomResultSetGenerator(1);
+    mockSpanner.putStatementResult(StatementResult.query(Statement.of(sql), generator.generate()));
+
+    try (Connection connection = createConnection();
+        PreparedStatement statement =
+            connection.prepareStatement(sql, java.sql.Statement.RETURN_GENERATED_KEYS)) {
+      assertTrue(statement.execute());
+      try (ResultSet resultSet = statement.getResultSet()) {
+        assertTrue(resultSet.next());
+        assertFalse(resultSet.next());
+      }
+      // Statements without generated keys return an empty result set.
+      assertFalse(statement.getGeneratedKeys().next());
+    }
+  }
+
+  @Test
+  public void testReturnGeneratedKeysForDdl() throws SQLException {
+    String sql = "create table test";
+    addDdlResponseToSpannerAdmin();
+
+    try (Connection connection = createConnection();
+        PreparedStatement statement =
+            connection.prepareStatement(sql, java.sql.Statement.RETURN_GENERATED_KEYS)) {
+      assertFalse(statement.execute());
+      assertEquals(java.sql.Statement.SUCCESS_NO_INFO, statement.getUpdateCount());
+      // Statements without generated keys return an empty result set.
+      assertFalse(statement.getGeneratedKeys().next());
+    }
+  }
+
+  @Test
+  public void testReturnAutoGeneratedKeysForStatementWithReturning() throws SQLException {
+    String sql =
+        "insert into test (id, value) values (1, 'One') "
+            + (dialect == Dialect.POSTGRESQL ? "RETURNING value" : "THEN RETURN value");
+    mockSpanner.putStatementResult(
+        StatementResult.query(
+            Statement.of(sql),
+            com.google.spanner.v1.ResultSet.newBuilder()
+                .setMetadata(
+                    ResultSetMetadata.newBuilder()
+                        .setRowType(
+                            StructType.newBuilder()
+                                .addFields(
+                                    Field.newBuilder()
+                                        .setName("value")
+                                        .setType(Type.newBuilder().setCode(TypeCode.STRING).build())
+                                        .build())
+                                .build())
+                        .build())
+                .addRows(
+                    ListValue.newBuilder()
+                        .addValues(Value.newBuilder().setStringValue("One").build())
+                        .build())
+                .setStats(ResultSetStats.newBuilder().setRowCountExact(1L).build())
+                .build()));
+
+    try (Connection connection = createConnection();
+        PreparedStatement statement =
+            connection.prepareStatement(sql, java.sql.Statement.RETURN_GENERATED_KEYS)) {
+      assertTrue(statement.execute());
+      try (ResultSet resultSet = statement.getResultSet()) {
+        assertTrue(resultSet.next());
+        assertEquals("One", resultSet.getString(1));
+        assertFalse(resultSet.next());
+      }
+
+      assertNotNull(statement.getGeneratedKeys());
+      assertFalse(statement.getGeneratedKeys().next());
+      // There should be no more results for this statement.
+      assertFalse(statement.getMoreResults());
+      assertEquals(-1, statement.getUpdateCount());
+    }
+  }
+
+  private void addDdlResponseToSpannerAdmin() {
+    mockDatabaseAdmin.addResponse(
+        Operation.newBuilder()
+            .setDone(true)
+            .setResponse(Any.pack(Empty.getDefaultInstance()))
+            .setMetadata(Any.pack(UpdateDatabaseDdlMetadata.getDefaultInstance()))
+            .build());
+  }
+}
diff --git a/src/test/java/com/google/cloud/spanner/jdbc/StatementReturnGeneratedKeysMockServerTest.java b/src/test/java/com/google/cloud/spanner/jdbc/StatementReturnGeneratedKeysMockServerTest.java
new file mode 100644
index 0000000..9cd9068
--- /dev/null
+++ b/src/test/java/com/google/cloud/spanner/jdbc/StatementReturnGeneratedKeysMockServerTest.java
@@ -0,0 +1,322 @@
+/*
+ * Copyright 2023 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.spanner.jdbc;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+
+import com.google.cloud.spanner.Dialect;
+import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult;
+import com.google.cloud.spanner.Statement;
+import com.google.cloud.spanner.connection.AbstractMockServerTest;
+import com.google.cloud.spanner.connection.RandomResultSetGenerator;
+import com.google.cloud.spanner.connection.SpannerPool;
+import com.google.common.collect.ImmutableList;
+import com.google.longrunning.Operation;
+import com.google.protobuf.Any;
+import com.google.protobuf.Empty;
+import com.google.protobuf.ListValue;
+import com.google.protobuf.Value;
+import com.google.spanner.admin.database.v1.UpdateDatabaseDdlMetadata;
+import com.google.spanner.v1.ResultSetMetadata;
+import com.google.spanner.v1.ResultSetStats;
+import com.google.spanner.v1.StructType;
+import com.google.spanner.v1.StructType.Field;
+import com.google.spanner.v1.Type;
+import com.google.spanner.v1.TypeCode;
+import java.sql.Connection;
+import java.sql.DriverManager;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Types;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+import org.junit.runners.Parameterized.Parameter;
+import org.junit.runners.Parameterized.Parameters;
+
+@RunWith(Parameterized.class)
+public class StatementReturnGeneratedKeysMockServerTest extends AbstractMockServerTest {
+  @Parameter public Dialect dialect;
+
+  @Parameters(name = "dialect = {0}")
+  public static Object[] data() {
+    return Dialect.values();
+  }
+
+  @Before
+  public void setupDialect() {
+    mockSpanner.putStatementResult(StatementResult.detectDialectResult(this.dialect));
+  }
+
+  @After
+  public void clearRequests() {
+    mockSpanner.clearRequests();
+    // Close Spanner pool to force a refresh of the dialect of a given database.
+    SpannerPool.closeSpannerPool();
+  }
+
+  private String createUrl() {
+    return String.format(
+        "jdbc:cloudspanner://localhost:%d/projects/%s/instances/%s/databases/%s?usePlainText=true",
+        getPort(), "proj", "inst", "db");
+  }
+
+  private Connection createConnection() throws SQLException {
+    return DriverManager.getConnection(createUrl());
+  }
+
+  interface SqlRunnable {
+    void run() throws SQLException;
+  }
+
+  @Test
+  public void testReturnAutoGeneratedKeys() throws SQLException {
+    String sql = "insert into test (id, value) values (1, 'One')";
+    String returningClause = dialect == Dialect.POSTGRESQL ? "\nRETURNING *" : "\nTHEN RETURN *";
+    mockSpanner.putStatementResult(
+        StatementResult.query(
+            Statement.of(sql + returningClause),
+            com.google.spanner.v1.ResultSet.newBuilder()
+                .setMetadata(
+                    ResultSetMetadata.newBuilder()
+                        .setRowType(
+                            StructType.newBuilder()
+                                .addFields(
+                                    Field.newBuilder()
+                                        .setName("id")
+                                        .setType(Type.newBuilder().setCode(TypeCode.INT64).build())
+                                        .build())
+                                .addFields(
+                                    Field.newBuilder()
+                                        .setName("value")
+                                        .setType(Type.newBuilder().setCode(TypeCode.STRING).build())
+                                        .build())
+                                .build())
+                        .build())
+                .addRows(
+                    ListValue.newBuilder()
+                        .addValues(Value.newBuilder().setStringValue("1").build())
+                        .addValues(Value.newBuilder().setStringValue("One").build())
+                        .build())
+                .setStats(ResultSetStats.newBuilder().setRowCountExact(1L).build())
+                .build()));
+
+    try (Connection connection = createConnection();
+        java.sql.Statement statement = connection.createStatement()) {
+      for (SqlRunnable runnable :
+          ImmutableList.<SqlRunnable>of(
+              () ->
+                  assertEquals(
+                      1, statement.executeUpdate(sql, java.sql.Statement.RETURN_GENERATED_KEYS)),
+              () ->
+                  assertEquals(
+                      1L,
+                      statement.executeLargeUpdate(sql, java.sql.Statement.RETURN_GENERATED_KEYS)),
+              () -> {
+                assertFalse(statement.execute(sql, java.sql.Statement.RETURN_GENERATED_KEYS));
+                assertEquals(1, statement.getUpdateCount());
+                assertEquals(1L, statement.getLargeUpdateCount());
+              })) {
+        runnable.run();
+        assertNotNull(statement.getGeneratedKeys());
+        assertEquals(2, statement.getGeneratedKeys().getMetaData().getColumnCount());
+        assertEquals(Types.BIGINT, statement.getGeneratedKeys().getMetaData().getColumnType(1));
+        assertEquals(Types.NVARCHAR, statement.getGeneratedKeys().getMetaData().getColumnType(2));
+        assertEquals("id", statement.getGeneratedKeys().getMetaData().getColumnName(1));
+        assertEquals("value", statement.getGeneratedKeys().getMetaData().getColumnName(2));
+        try (ResultSet keys = statement.getGeneratedKeys()) {
+          assertTrue(keys.next());
+          assertEquals(1L, keys.getLong(1));
+          assertEquals("One", keys.getString(2));
+          assertFalse(keys.next());
+        }
+        // There should be no more results for this statement.
+        assertFalse(statement.getMoreResults());
+        assertEquals(-1, statement.getUpdateCount());
+      }
+    }
+  }
+
+  @Test
+  public void testReturnColumnNames() throws SQLException {
+    String sql = "insert into test (id, value) values (1, 'One')";
+    String returningClause =
+        dialect == Dialect.POSTGRESQL ? "\nRETURNING \"id\"" : "\nTHEN RETURN `id`";
+    mockSpanner.putStatementResult(
+        StatementResult.query(
+            Statement.of(sql + returningClause),
+            com.google.spanner.v1.ResultSet.newBuilder()
+                .setMetadata(
+                    ResultSetMetadata.newBuilder()
+                        .setRowType(
+                            StructType.newBuilder()
+                                .addFields(
+                                    Field.newBuilder()
+                                        .setName("id")
+                                        .setType(Type.newBuilder().setCode(TypeCode.INT64).build())
+                                        .build())
+                                .build())
+                        .build())
+                .addRows(
+                    ListValue.newBuilder()
+                        .addValues(Value.newBuilder().setStringValue("1").build())
+                        .build())
+                .setStats(ResultSetStats.newBuilder().setRowCountExact(1L).build())
+                .build()));
+
+    try (Connection connection = createConnection();
+        java.sql.Statement statement = connection.createStatement()) {
+      for (SqlRunnable runnable :
+          ImmutableList.<SqlRunnable>of(
+              () -> assertEquals(1, statement.executeUpdate(sql, new String[] {"id"})),
+              () -> assertEquals(1L, statement.executeLargeUpdate(sql, new String[] {"id"})),
+              () -> {
+                assertFalse(statement.execute(sql, new String[] {"id"}));
+                assertEquals(1, statement.getUpdateCount());
+                assertEquals(1L, statement.getLargeUpdateCount());
+              })) {
+        runnable.run();
+        assertNotNull(statement.getGeneratedKeys());
+        assertEquals(1, statement.getGeneratedKeys().getMetaData().getColumnCount());
+        assertEquals(Types.BIGINT, statement.getGeneratedKeys().getMetaData().getColumnType(1));
+        assertEquals("id", statement.getGeneratedKeys().getMetaData().getColumnName(1));
+        try (ResultSet keys = statement.getGeneratedKeys()) {
+          assertTrue(keys.next());
+          assertEquals(1L, keys.getLong(1));
+          assertFalse(keys.next());
+        }
+        // There should be no more results for this statement.
+        assertFalse(statement.getMoreResults());
+        assertEquals(-1, statement.getUpdateCount());
+      }
+    }
+  }
+
+  @Test
+  public void testReturnColumnIndices() throws SQLException {
+    String sql = "insert into test (id, value) values (1, 'One')";
+    mockSpanner.putStatementResult(StatementResult.update(Statement.of(sql), 1L));
+
+    try (Connection connection = createConnection();
+        java.sql.Statement statement = connection.createStatement()) {
+      // The JDBC driver silently ignores the request for column indices.
+      for (SqlRunnable runnable :
+          ImmutableList.<SqlRunnable>of(
+              () -> assertEquals(1, statement.executeUpdate(sql, new int[] {1})),
+              () -> assertEquals(1L, statement.executeLargeUpdate(sql, new int[] {1})),
+              () -> {
+                assertFalse(statement.execute(sql, new int[] {1}));
+                assertEquals(1, statement.getUpdateCount());
+                assertEquals(1L, statement.getLargeUpdateCount());
+              })) {
+        runnable.run();
+        assertNotNull(statement.getGeneratedKeys());
+        assertFalse(statement.getGeneratedKeys().next());
+      }
+    }
+  }
+
+  @Test
+  public void testReturnGeneratedKeysForQuery() throws SQLException {
+    String sql = "select * from test";
+    RandomResultSetGenerator generator = new RandomResultSetGenerator(1);
+    mockSpanner.putStatementResult(StatementResult.query(Statement.of(sql), generator.generate()));
+
+    try (Connection connection = createConnection();
+        java.sql.Statement statement = connection.createStatement()) {
+      assertTrue(statement.execute(sql, java.sql.Statement.RETURN_GENERATED_KEYS));
+      try (ResultSet resultSet = statement.getResultSet()) {
+        assertTrue(resultSet.next());
+        assertFalse(resultSet.next());
+      }
+      // Statements without generated keys return an empty result set.
+      assertFalse(statement.getGeneratedKeys().next());
+    }
+  }
+
+  @Test
+  public void testReturnGeneratedKeysForDdl() throws SQLException {
+    String sql = "create table test";
+    addDdlResponseToSpannerAdmin();
+
+    try (Connection connection = createConnection();
+        java.sql.Statement statement = connection.createStatement()) {
+      assertFalse(statement.execute(sql, java.sql.Statement.RETURN_GENERATED_KEYS));
+      assertEquals(java.sql.Statement.SUCCESS_NO_INFO, statement.getUpdateCount());
+      // Statements without generated keys return an empty result set.
+      assertFalse(statement.getGeneratedKeys().next());
+    }
+  }
+
+  @Test
+  public void testReturnAutoGeneratedKeysForStatementWithReturning() throws SQLException {
+    String sql =
+        "insert into test (id, value) values (1, 'One') "
+            + (dialect == Dialect.POSTGRESQL ? "RETURNING value" : "THEN RETURN value");
+    mockSpanner.putStatementResult(
+        StatementResult.query(
+            Statement.of(sql),
+            com.google.spanner.v1.ResultSet.newBuilder()
+                .setMetadata(
+                    ResultSetMetadata.newBuilder()
+                        .setRowType(
+                            StructType.newBuilder()
+                                .addFields(
+                                    Field.newBuilder()
+                                        .setName("value")
+                                        .setType(Type.newBuilder().setCode(TypeCode.STRING).build())
+                                        .build())
+                                .build())
+                        .build())
+                .addRows(
+                    ListValue.newBuilder()
+                        .addValues(Value.newBuilder().setStringValue("One").build())
+                        .build())
+                .setStats(ResultSetStats.newBuilder().setRowCountExact(1L).build())
+                .build()));
+
+    try (Connection connection = createConnection();
+        java.sql.Statement statement = connection.createStatement()) {
+      assertTrue(statement.execute(sql, java.sql.Statement.RETURN_GENERATED_KEYS));
+      try (ResultSet resultSet = statement.getResultSet()) {
+        assertTrue(resultSet.next());
+        assertEquals("One", resultSet.getString(1));
+        assertFalse(resultSet.next());
+      }
+
+      assertNotNull(statement.getGeneratedKeys());
+      assertFalse(statement.getGeneratedKeys().next());
+      // There should be no more results for this statement.
+      assertFalse(statement.getMoreResults());
+      assertEquals(-1, statement.getUpdateCount());
+    }
+  }
+
+  private void addDdlResponseToSpannerAdmin() {
+    mockDatabaseAdmin.addResponse(
+        Operation.newBuilder()
+            .setDone(true)
+            .setResponse(Any.pack(Empty.getDefaultInstance()))
+            .setMetadata(Any.pack(UpdateDatabaseDdlMetadata.getDefaultInstance()))
+            .build());
+  }
+}