chore: add copyOf method for result sets (#1316)

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.

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.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 @@ static JdbcResultSet of(Statement statement, com.google.cloud.spanner.ResultSet
     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;

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

@@ -37,6 +37,7 @@
 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.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 void testGetOffsetDateTime() throws SQLException {
         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());
+  }
 }