chore: method for adding returning clause to statements (#1311)

* chore: method for adding returning clause to statements

Adds a method to JdbcStatement for appending a THEN RETURN/RETURNING
clause to the statement. This will be used to modify statements that
request generated keys to be returned.

* feat: support return all columns

* fix: only add THEN RETURN * from DML

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

@@ -23,19 +23,25 @@
 import com.google.cloud.spanner.Statement;
 import com.google.cloud.spanner.Type;
 import com.google.cloud.spanner.Type.StructField;
+import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStatement;
 import com.google.cloud.spanner.connection.StatementResult;
 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.ResultSet;
 import java.sql.SQLException;
 import java.util.ArrayList;
 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 {
+  static final ImmutableList<String> ALL_COLUMNS = ImmutableList.of("*");
+
   enum BatchType {
     NONE,
     DML,
@@ -98,6 +104,81 @@ public long executeLargeUpdate(String sql) throws SQLException {
     }
   }
 
+  /**
+   * Adds a THEN RETURN/RETURNING clause to the given statement if the following conditions are all
+   * met:
+   *
+   * <ol>
+   *   <li>The generatedKeysColumns is not null or 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)
+      throws SQLException {
+    if (generatedKeysColumns == null || generatedKeysColumns.isEmpty()) {
+      return statement;
+    }
+    // Check if the statement is a DML statement or not.
+    ParsedStatement parsedStatement = getConnection().getParser().parse(statement);
+    if (parsedStatement.isUpdate() && !parsedStatement.hasReturningClause()) {
+      if (generatedKeysColumns.size() == 1
+          && ALL_COLUMNS.get(0).equals(generatedKeysColumns.get(0))) {
+        // Add a 'THEN RETURN/RETURNING *' clause to the statement.
+        return statement
+            .toBuilder()
+            .replace(statement.getSql() + getReturningAllColumnsClause())
+            .build();
+      }
+      // Add a 'THEN RETURN/RETURNING col1, col2, ...' to the statement.
+      // The column names will be quoted using the dialect-specific identifier quoting character.
+      return statement
+          .toBuilder()
+          .replace(
+              generatedKeysColumns.stream()
+                  .map(this::quoteColumn)
+                  .collect(
+                      Collectors.joining(
+                          ", ", statement.getSql() + getReturningClause() + " ", "")))
+          .build();
+    }
+    return statement;
+  }
+
+  /** Returns the dialect-specific clause for returning values from a DML statement. */
+  String getReturningAllColumnsClause() {
+    switch (getConnection().getDialect()) {
+      case POSTGRESQL:
+        return "\nRETURNING *";
+      case GOOGLE_STANDARD_SQL:
+      default:
+        return "\nTHEN RETURN *";
+    }
+  }
+
+  /** Returns the dialect-specific clause for returning values from a DML statement. */
+  String getReturningClause() {
+    switch (getConnection().getDialect()) {
+      case POSTGRESQL:
+        return "\nRETURNING";
+      case GOOGLE_STANDARD_SQL:
+      default:
+        return "\nTHEN RETURN";
+    }
+  }
+
+  /** Adds dialect-specific quotes to the given column name. */
+  String quoteColumn(String column) {
+    switch (getConnection().getDialect()) {
+      case POSTGRESQL:
+        return "\"" + column + "\"";
+      case GOOGLE_STANDARD_SQL:
+      default:
+        return "`" + column + "`";
+    }
+  }
+
   @Override
   public boolean execute(String sql) throws SQLException {
     checkClosed();

src/test/java/com/google/cloud/spanner/jdbc/JdbcStatementTest.java

@@ -21,6 +21,7 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
@@ -39,6 +40,7 @@
 import com.google.cloud.spanner.connection.StatementResult;
 import com.google.cloud.spanner.connection.StatementResult.ResultType;
 import com.google.cloud.spanner.jdbc.JdbcSqlExceptionFactory.JdbcSqlExceptionImpl;
+import com.google.common.collect.ImmutableList;
 import com.google.rpc.Code;
 import java.sql.ResultSet;
 import java.sql.SQLException;
@@ -47,6 +49,7 @@
 import java.util.Arrays;
 import java.util.List;
 import java.util.concurrent.TimeUnit;
+import javax.annotation.Nullable;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.Parameterized;
@@ -599,4 +602,157 @@ public void testConvertUpdateCountsToSuccessNoInfo() throws SQLException {
               (long) Statement.SUCCESS_NO_INFO);
     }
   }
+
+  @Test
+  public void testAddReturningToStatement() throws SQLException {
+    JdbcConnection connection = mock(JdbcConnection.class);
+    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')", ImmutableList.of());
+      assertAddReturningEquals(
+          statement,
+          dialect == Dialect.POSTGRESQL
+              ? "insert into test (id, value) values (1, 'One')\nRETURNING \"id\""
+              : "insert into test (id, value) values (1, 'One')\nTHEN RETURN `id`",
+          "insert into test (id, value) values (1, 'One')",
+          ImmutableList.of("id"));
+      assertAddReturningEquals(
+          statement,
+          dialect == Dialect.POSTGRESQL
+              ? "insert into test (id, value) values (1, 'One')\nRETURNING \"id\", \"value\""
+              : "insert into test (id, value) values (1, 'One')\nTHEN RETURN `id`, `value`",
+          "insert into test (id, value) values (1, 'One')",
+          ImmutableList.of("id", "value"));
+      assertAddReturningEquals(
+          statement,
+          dialect == Dialect.POSTGRESQL
+              ? "insert into test (id, value) values (1, 'One')\nRETURNING *"
+              : "insert into test (id, value) values (1, 'One')\nTHEN RETURN *",
+          "insert into test (id, value) values (1, 'One')",
+          ImmutableList.of("*"));
+      // Requesting generated keys for a DML statement that already contains a returning clause is a
+      // no-op.
+      assertAddReturningSame(
+          statement,
+          "insert into test (id, value) values (1, 'One') "
+              + statement.getReturningClause()
+              + " value",
+          ImmutableList.of("id"));
+      // Requesting generated keys for a query is a no-op.
+      for (ImmutableList<String> keys :
+          ImmutableList.of(
+              ImmutableList.of("id"), ImmutableList.of("id", "value"), ImmutableList.of("*"))) {
+        assertAddReturningSame(statement, "select id, value from test", keys);
+      }
+
+      // 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", ImmutableList.of());
+      assertAddReturningEquals(
+          statement,
+          dialect == Dialect.POSTGRESQL
+              ? "update test set value='Two' where id=1\nRETURNING \"value\""
+              : "update test set value='Two' where id=1\nTHEN RETURN `value`",
+          "update test set value='Two' where id=1",
+          ImmutableList.of("value"));
+      assertAddReturningEquals(
+          statement,
+          dialect == Dialect.POSTGRESQL
+              ? "update test set value='Two' where id=1\nRETURNING \"value\", \"id\""
+              : "update test set value='Two' where id=1\nTHEN RETURN `value`, `id`",
+          "update test set value='Two' where id=1",
+          ImmutableList.of("value", "id"));
+      assertAddReturningEquals(
+          statement,
+          dialect == Dialect.POSTGRESQL
+              ? "update test set value='Two' where id=1\nRETURNING *"
+              : "update test set value='Two' where id=1\nTHEN RETURN *",
+          "update test set value='Two' where id=1",
+          ImmutableList.of("*"));
+      // Requesting generated keys for a DML statement that already contains a returning clause is a
+      // no-op.
+      assertAddReturningSame(
+          statement,
+          "update test set value='Two' where id=1 " + statement.getReturningClause() + " value",
+          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", ImmutableList.of());
+      assertAddReturningEquals(
+          statement,
+          dialect == Dialect.POSTGRESQL
+              ? "delete test where id=1\nRETURNING \"value\""
+              : "delete test where id=1\nTHEN RETURN `value`",
+          "delete test where id=1",
+          ImmutableList.of("value"));
+      assertAddReturningEquals(
+          statement,
+          dialect == Dialect.POSTGRESQL
+              ? "delete test where id=1\nRETURNING \"id\", \"value\""
+              : "delete test where id=1\nTHEN RETURN `id`, `value`",
+          "delete test where id=1",
+          ImmutableList.of("id", "value"));
+      assertAddReturningEquals(
+          statement,
+          dialect == Dialect.POSTGRESQL
+              ? "delete test where id=1\nRETURNING *"
+              : "delete test where id=1\nTHEN RETURN *",
+          "delete test where id=1",
+          ImmutableList.of("*"));
+      // Requesting generated keys for a DML statement that already contains a returning clause is a
+      // no-op.
+      for (ImmutableList<String> keys :
+          ImmutableList.of(
+              ImmutableList.of("id"), ImmutableList.of("id", "value"), ImmutableList.of("*"))) {
+        assertAddReturningSame(
+            statement,
+            "delete test where id=1 "
+                + (dialect == Dialect.POSTGRESQL
+                    ? "delete test where id=1\nRETURNING"
+                    : "delete test where id=1\nTHEN RETURN")
+                + " value",
+            keys);
+      }
+
+      // Requesting generated keys for DDL is a no-op.
+      for (ImmutableList<String> keys :
+          ImmutableList.of(
+              ImmutableList.of("id"), ImmutableList.of("id", "value"), ImmutableList.of("*"))) {
+        assertAddReturningSame(
+            statement,
+            dialect == Dialect.POSTGRESQL
+                ? "create table test (id bigint primary key, value varchar)"
+                : "create table test (id int64, value string(max)) primary key (id)",
+            keys);
+      }
+    }
+  }
+
+  private void assertAddReturningSame(
+      JdbcStatement statement, String sql, @Nullable ImmutableList<String> generatedKeysColumns)
+      throws SQLException {
+    com.google.cloud.spanner.Statement spannerStatement =
+        com.google.cloud.spanner.Statement.of(sql);
+    assertSame(
+        spannerStatement,
+        statement.addReturningToStatement(spannerStatement, generatedKeysColumns));
+  }
+
+  private void assertAddReturningEquals(
+      JdbcStatement statement,
+      String expectedSql,
+      String sql,
+      @Nullable ImmutableList<String> generatedKeysColumns)
+      throws SQLException {
+    com.google.cloud.spanner.Statement spannerStatement =
+        com.google.cloud.spanner.Statement.of(sql);
+    assertEquals(
+        com.google.cloud.spanner.Statement.of(expectedSql),
+        statement.addReturningToStatement(spannerStatement, generatedKeysColumns));
+  }
 }