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 f8ab4526..94b46042 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.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 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 { + 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/test/java/com/google/cloud/spanner/jdbc/JdbcResultSetTest.java b/src/test/java/com/google/cloud/spanner/jdbc/JdbcResultSetTest.java index 024f6c24..476c4049 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.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()); + } }