Repository: googleapis/java-spanner-jdbc
Pull request #1316 (2023-08-13 15:20:24)

chore: add copyOf method for result sets
Adds a method for creating a JDBC ResultSet that creates a copy of a Cloud Spanner ResultSet. This type of ResultSet can be used for small result sets that are unlikely to be closed by the client application, such as a ResultSet that holds the generated keys of an update statement. If an application requests the JDBC driver to return the generated keys, and then never gets and closes the generated keys ResultSet, the underlying session and/or result stream could be leaked.



Patch:
diff --git a/.github/sync-repo-settings.yaml b/.github/sync-repo-settings.yaml
index 4879f02..a259dc0 100644
--- a/.github/sync-repo-settings.yaml
+++ b/.github/sync-repo-settings.yaml
@@ -39,6 +39,7 @@ branchProtectionRules:
     - OwlBot Post Processor
     - 'Kokoro - Test: Java GraalVM Native Image'
     - 'Kokoro - Test: Java 17 GraalVM Native Image'
+    - javadoc
 
 # Identifies the protection rule pattern. Name of the branch to be protected.
 # Defaults to `main`
diff --git a/CHANGELOG.md b/CHANGELOG.md
index af3d292..3992709 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,19 @@
 # Changelog
 
+## [2.11.9](https://github.com/googleapis/java-spanner-jdbc/compare/v2.11.8...v2.11.9) (2023-08-21)
+
+
+### Dependencies
+
+* Update dependency com.google.cloud:google-cloud-spanner-bom to v6.45.3 ([#1325](https://github.com/googleapis/java-spanner-jdbc/issues/1325)) ([d96d278](https://github.com/googleapis/java-spanner-jdbc/commit/d96d278f58942c38b30dc0a4e1f9f92aefdc760b))
+
+## [2.11.8](https://github.com/googleapis/java-spanner-jdbc/compare/v2.11.7...v2.11.8) (2023-08-15)
+
+
+### Dependencies
+
+* Update dependency com.google.cloud:google-cloud-spanner-bom to v6.45.2 ([#1318](https://github.com/googleapis/java-spanner-jdbc/issues/1318)) ([e924178](https://github.com/googleapis/java-spanner-jdbc/commit/e9241787b94cb614f658f5e6c977ffc008fd3397))
+
 ## [2.11.7](https://github.com/googleapis/java-spanner-jdbc/compare/v2.11.6...v2.11.7) (2023-08-13)
 
 
diff --git a/pom.xml b/pom.xml
index 3378a23..4a33dea 100644
--- a/pom.xml
+++ b/pom.xml
@@ -4,7 +4,7 @@
   xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
   <modelVersion>4.0.0</modelVersion>
   <artifactId>google-cloud-spanner-jdbc</artifactId>
-  <version>2.11.8-SNAPSHOT</version><!-- {x-version-update:google-cloud-spanner-jdbc:current} -->
+  <version>2.11.10-SNAPSHOT</version><!-- {x-version-update:google-cloud-spanner-jdbc:current} -->
   <packaging>jar</packaging>
   <name>Google Cloud Spanner JDBC</name>
   <url>https://github.com/googleapis/java-spanner-jdbc</url>
@@ -62,7 +62,7 @@
       <dependency>
         <groupId>com.google.cloud</groupId>
         <artifactId>google-cloud-spanner-bom</artifactId>
-        <version>6.45.1</version>
+        <version>6.45.3</version>
         <type>pom</type>
         <scope>import</scope>
       </dependency>
diff --git a/samples/install-without-bom/pom.xml b/samples/install-without-bom/pom.xml
index e5c2fd9..be63235 100644
--- a/samples/install-without-bom/pom.xml
+++ b/samples/install-without-bom/pom.xml
@@ -29,7 +29,7 @@
     <dependency>
       <groupId>com.google.cloud</groupId>
       <artifactId>google-cloud-spanner-jdbc</artifactId>
-      <version>2.11.6</version>
+      <version>2.11.9</version>
     </dependency>
     <!-- [END spanner-jdbc_install_without_bom] -->
 
diff --git a/samples/snapshot/pom.xml b/samples/snapshot/pom.xml
index 99c5094..5cc1db0 100644
--- a/samples/snapshot/pom.xml
+++ b/samples/snapshot/pom.xml
@@ -28,7 +28,7 @@
     <dependency>
       <groupId>com.google.cloud</groupId>
       <artifactId>google-cloud-spanner-jdbc</artifactId>
-      <version>2.11.8-SNAPSHOT</version>
+      <version>2.11.10-SNAPSHOT</version>
     </dependency>
     <!-- {x-version-update-end} -->
 
diff --git a/src/main/java/com/google/cloud/spanner/jdbc/JdbcResultSet.java b/src/main/java/com/google/cloud/spanner/jdbc/JdbcResultSet.java
index f8ab452..94b4604 100644
--- a/src/main/java/com/google/cloud/spanner/jdbc/JdbcResultSet.java
+++ b/src/main/java/com/google/cloud/spanner/jdbc/JdbcResultSet.java
@@ -16,12 +16,13 @@
 
 package com.google.cloud.spanner.jdbc;
 
-import static com.google.cloud.spanner.Type.Code.PG_NUMERIC;
-
+import com.google.cloud.spanner.ResultSets;
+import com.google.cloud.spanner.Struct;
 import com.google.cloud.spanner.Type;
 import com.google.cloud.spanner.Type.Code;
 import com.google.cloud.spanner.Value;
 import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableList;
 import java.io.ByteArrayInputStream;
 import java.io.InputStream;
 import java.io.Reader;
@@ -43,8 +44,11 @@ import java.sql.Statement;
 import java.sql.Time;
 import java.sql.Timestamp;
 import java.util.Calendar;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.NoSuchElementException;
+import javax.annotation.Nonnull;
 
 /** Implementation of {@link ResultSet} for Cloud Spanner */
 class JdbcResultSet extends AbstractJdbcResultSet {
@@ -60,6 +64,57 @@ class JdbcResultSet extends AbstractJdbcResultSet {
     return new JdbcResultSet(statement, resultSet);
   }
 
+  /**
+   * Creates a JDBC result set by copying the given Spanner {@link
+   * com.google.cloud.spanner.ResultSet}. This can be used for result sets that are known not to be
+   * too large. This type of result set should be preferred for results that are unlikely to be
+   * closed by the client application, such as the returned generated keys of an update statement.
+   * The copy will not hold on to a reference to a Cloud Spanner session or result stream. All the
+   * data in the given Spanner {@link com.google.cloud.spanner.ResultSet} have been consumed after
+   * calling this method. The {@link com.google.cloud.spanner.ResultSet} is not closed by this
+   * method.
+   */
+  static JdbcResultSet copyOf(@Nonnull com.google.cloud.spanner.ResultSet resultSet) {
+    Preconditions.checkNotNull(resultSet);
+    // Make the copy first. This ensures that ResultSet#next() has been called at least once, which
+    // is necessary to get the type of the result set.
+    ImmutableList<Struct> rows = ImmutableList.copyOf(new ResultSetIterator(resultSet));
+    return of(ResultSets.forRows(resultSet.getType(), rows));
+  }
+
+  /**
+   * {@link Iterator} implementation for {@link com.google.cloud.spanner.ResultSet}. This is used to
+   * create a copy of an existing result set without the need to iterate the rows more than once.
+   */
+  private static class ResultSetIterator implements Iterator<Struct> {
+    private final com.google.cloud.spanner.ResultSet resultSet;
+    private boolean calculatedHasNext = false;
+    private boolean hasNext = false;
+
+    ResultSetIterator(com.google.cloud.spanner.ResultSet resultSet) {
+      this.resultSet = resultSet;
+    }
+
+    @Override
+    public boolean hasNext() {
+      if (!calculatedHasNext) {
+        calculatedHasNext = true;
+        hasNext = resultSet.next();
+      }
+      return hasNext;
+    }
+
+    @Override
+    public Struct next() {
+      if (hasNext()) {
+        // Indicate that the next call to hasNext() must re-check whether there are more results.
+        calculatedHasNext = false;
+        return resultSet.getCurrentRowAsStruct();
+      }
+      throw new NoSuchElementException();
+    }
+  }
+
   private boolean closed = false;
   private final Statement statement;
   private boolean wasNull = false;
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 051ed9d..3e7c86f 100644
--- a/src/main/java/com/google/cloud/spanner/jdbc/JdbcStatement.java
+++ b/src/main/java/com/google/cloud/spanner/jdbc/JdbcStatement.java
@@ -23,9 +23,11 @@ import com.google.cloud.spanner.SpannerException;
 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;
@@ -33,9 +35,13 @@ 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 @@ class JdbcStatement extends AbstractJdbcStatement {
     }
   }
 
+  /**
+   * 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();
diff --git a/versions.txt b/versions.txt
index 5459f1d..4a69a3a 100644
--- a/versions.txt
+++ b/versions.txt
@@ -1,4 +1,4 @@
 # Format:
 # module:released-version:current-version
 
-google-cloud-spanner-jdbc:2.11.7:2.11.8-SNAPSHOT
+google-cloud-spanner-jdbc:2.11.9:2.11.10-SNAPSHOT


Test patch:
diff --git a/src/test/java/com/google/cloud/spanner/jdbc/JdbcResultSetTest.java b/src/test/java/com/google/cloud/spanner/jdbc/JdbcResultSetTest.java
index 024f6c2..476c404 100644
--- a/src/test/java/com/google/cloud/spanner/jdbc/JdbcResultSetTest.java
+++ b/src/test/java/com/google/cloud/spanner/jdbc/JdbcResultSetTest.java
@@ -37,6 +37,7 @@ import com.google.cloud.spanner.Type;
 import com.google.cloud.spanner.Type.StructField;
 import com.google.cloud.spanner.Value;
 import com.google.cloud.spanner.jdbc.JdbcSqlExceptionFactory.JdbcSqlExceptionImpl;
+import com.google.common.collect.ImmutableList;
 import com.google.rpc.Code;
 import java.io.IOException;
 import java.io.InputStream;
@@ -46,9 +47,11 @@ import java.math.RoundingMode;
 import java.net.MalformedURLException;
 import java.net.URL;
 import java.nio.charset.StandardCharsets;
+import java.sql.ResultSetMetaData;
 import java.sql.SQLException;
 import java.sql.Statement;
 import java.sql.Time;
+import java.sql.Types;
 import java.time.Instant;
 import java.time.LocalDate;
 import java.time.OffsetDateTime;
@@ -1753,4 +1756,57 @@ public class JdbcResultSetTest {
         offsetDateTime);
     assertFalse(subject.wasNull());
   }
+
+  @Test
+  public void testCopyOf() throws SQLException {
+    ResultSet original =
+        ResultSets.forRows(
+            Type.struct(StructField.of("id", Type.int64()), StructField.of("value", Type.string())),
+            ImmutableList.of(
+                Struct.newBuilder().set("id").to(1L).set("value").to("One").build(),
+                Struct.newBuilder().set("id").to(2L).set("value").to("Two").build()));
+    java.sql.ResultSet copy = JdbcResultSet.copyOf(original);
+    // The original result set has been fully consumed.
+    assertFalse(original.next());
+    // We can safely close the original result set and still use the copy.
+    original.close();
+
+    ResultSetMetaData metadata = copy.getMetaData();
+    assertEquals(2, metadata.getColumnCount());
+    assertEquals("id", metadata.getColumnName(1));
+    assertEquals("value", metadata.getColumnName(2));
+    assertEquals(Types.BIGINT, metadata.getColumnType(1));
+    assertEquals(Types.NVARCHAR, metadata.getColumnType(2));
+
+    assertTrue(copy.next());
+    assertEquals(1L, copy.getLong(1));
+    assertEquals("One", copy.getString(2));
+    assertTrue(copy.next());
+    assertEquals(2L, copy.getLong("id"));
+    assertEquals("Two", copy.getString("value"));
+    assertFalse(copy.next());
+  }
+
+  @Test
+  public void testCopyOfEmpty() throws SQLException {
+    ResultSet original =
+        ResultSets.forRows(
+            Type.struct(StructField.of("id", Type.int64()), StructField.of("value", Type.string())),
+            ImmutableList.of());
+    java.sql.ResultSet copy = JdbcResultSet.copyOf(original);
+    // The original result set has been fully consumed.
+    assertFalse(original.next());
+    // We can safely close the original result set and still use the copy.
+    original.close();
+
+    ResultSetMetaData metadata = copy.getMetaData();
+    assertEquals(2, metadata.getColumnCount());
+    assertEquals("id", metadata.getColumnName(1));
+    assertEquals("value", metadata.getColumnName(2));
+    assertEquals(Types.BIGINT, metadata.getColumnType(1));
+    assertEquals(Types.NVARCHAR, metadata.getColumnType(2));
+
+    // The copy should not contain any rows.
+    assertFalse(copy.next());
+  }
 }
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 c7b6a6c..d445835 100644
--- a/src/test/java/com/google/cloud/spanner/jdbc/JdbcStatementTest.java
+++ b/src/test/java/com/google/cloud/spanner/jdbc/JdbcStatementTest.java
@@ -21,6 +21,7 @@ import static org.junit.Assert.assertArrayEquals;
 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.Connection;
 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.sql.Statement;
 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 class JdbcStatementTest {
               (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));
+  }
 }