From ce7d74c01e1e180dfe9eb7db3b1bdfd157c03aba Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= Date: Fri, 25 Aug 2023 14:04:57 +0200 Subject: [PATCH] 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. --- .../cloud/spanner/jdbc/JdbcResultSet.java | 59 ++++++++++++++++++- .../cloud/spanner/jdbc/JdbcResultSetTest.java | 56 ++++++++++++++++++ 2 files changed, 113 insertions(+), 2 deletions(-) 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()); + } }