feat: return generated keys (#1310)

Adds support for returning generated keys.

Author: olavloite

src/main/java/com/google/cloud/spanner/jdbc/AbstractJdbcStatement.java

@@ -24,7 +24,6 @@
 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 @@ private ResultSet executeQuery(
     }
   }
 
-  /**
-   * 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.

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.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 @@ public Statement createStatement() throws SQLException {
 
   @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 @@ public PreparedStatement prepareStatement(
 
   @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

src/main/java/com/google/cloud/spanner/jdbc/JdbcPreparedStatement.java

@@ -24,7 +24,9 @@
 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 @@ ResultSet executeQueryWithOptions(QueryOption... options) throws SQLException {
 
   @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