SafeSql.queryLazily()

Author: fluentfuture

mug-guava/pom.xml

@@ -70,6 +70,12 @@
       <groupId>com.google.guava</groupId>
       <artifactId>guava-testlib</artifactId>
     </dependency>
+
+    <dependency>
+        <groupId>org.mockito</groupId>
+        <artifactId>mockito-core</artifactId>
+        <scope>test</scope>
+    </dependency>
     <dependency>
       <groupId>com.h2database</groupId>
       <artifactId>h2</artifactId>

mug-guava/src/.DS_Store

@@ -0,0 +1,66 @@
+/*****************************************************************************
+ * ------------------------------------------------------------------------- *
+ * Licensed under the Apache License, Version 2.0 (the "License");           *
+ * you may not use this file except in compliance with the License.          *
+ * You may obtain a copy of the License at                                   *
+ *                                                                           *
+ * http://www.apache.org/licenses/LICENSE-2.0                                *
+ *                                                                           *
+ * Unless required by applicable law or agreed to in writing, software       *
+ * distributed under the License is distributed on an "AS IS" BASIS,         *
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  *
+ * See the License for the specific language governing permissions and       *
+ * limitations under the License.                                            *
+ *****************************************************************************/
+package com.google.mu.safesql;
+
+import static com.google.common.base.Preconditions.checkState;
+
+import java.io.IOException;
+import java.sql.SQLException;
+import java.util.stream.Stream;
+
+import com.google.common.base.VerifyException;
+import com.google.common.io.Closer;
+import com.google.errorprone.annotations.MustBeClosed;
+
+/** Helps manage JDBC resources that need to be closed now or lazily. */
+final class JdbcCloser implements AutoCloseable {
+  interface JdbcCloseable {
+    void run() throws SQLException;
+  }
+
+  private final Closer closer = Closer.create();
+  private boolean owned = true;
+
+  void register(JdbcCloseable closeable) {
+    closer.register(() -> {
+      try {
+        closeable.run();
+      } catch (SQLException e) {
+        throw new UncheckedSqlException(e);
+      }
+    });
+  }
+
+  @MustBeClosed
+  <T> Stream<T> attachTo(Stream<T> stream) {
+    checkState(owned);
+    owned = false;
+    return stream.onClose(this::closeNow);
+  }
+
+  @Override public void close() {
+    if (owned) {
+      closeNow();
+    }
+  }
+
+  private void closeNow() {
+    try {
+      closer.close();
+    } catch (IOException e) {
+      throw new VerifyException(e);
+    }
+  }
+}

mug-guava/src/main/java/com/google/mu/safesql/JdbcCloser.java

@@ -29,6 +29,7 @@
 import static com.google.mu.util.Substring.suffix;
 import static com.google.mu.util.Substring.word;
 import static com.google.mu.util.stream.MoreStreams.indexesFrom;
+import static com.google.mu.util.stream.MoreStreams.whileNotNull;
 import static java.util.Collections.emptyList;
 import static java.util.Collections.unmodifiableList;
 import static java.util.stream.Collectors.collectingAndThen;
@@ -47,6 +48,7 @@
 import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Supplier;
 import java.util.stream.Collector;
 import java.util.stream.Collectors;
@@ -702,7 +704,8 @@ public SafeSql orElse(Supplier<SafeSql> fallback) {
    *     .query(connection, row -> row.getLong("id"));
    * }</pre>
    *
-   * <p>Internally it delegates to {@link PreparedStatement#executeQuery}.
+   * <p>Internally it delegates to {@link PreparedStatement#executeQuery} or {@link
+   * Statement#executeQuery} if this sql contains no JDBC binding parameters.
    *
    * @throws UncheckedSqlException wraps {@link SQLException} if failed
    */
@@ -725,6 +728,64 @@ public <T> List<T> query(
     }
   }
 
+  /**
+   * Executes the encapsulated SQL as a query against {@code connection}, sets {@code fetchSize}
+   * using {@link Statement#setFetchSize}, and then fetches the results lazily in a stream.
+   *
+   * <p>The returned {@code Stream} includes results transformed by {@code rowMapper}.
+   * The caller must close it using try-with-resources idiom, which will close the associated
+   * statement and {@link ResultSet}.
+   *
+   * <p>For example: <pre>{@code
+   * SafeSql sql = SafeSql.of("SELECT name FROM Users WHERE name LIKE '%{name}%'", name);
+   * try (Stream<String> names = sql.query(connection, row -> row.getLong("id"))) {
+   *   return names.findFirst();
+   * }
+   * }</pre>
+   *
+   * <p>Internally it delegates to {@link PreparedStatement#executeQuery} or {@link
+   * Statement#executeQuery} if this sql contains no JDBC binding parameters.
+   *
+   * @throws UncheckedSqlException wraps {@link SQLException} if failed
+   * @since 8.4
+   */
+  @MustBeClosed
+  @SuppressWarnings("MustBeClosedChecker")
+  public <T> Stream<T> queryLazily(
+      Connection connection, int fetchSize, SqlFunction<? super ResultSet, ? extends T> rowMapper) {
+    checkNotNull(rowMapper);
+    if (paramValues.isEmpty()) {
+      return lazy(connection::createStatement, fetchSize, stmt -> stmt.executeQuery(sql), rowMapper);
+    }
+    return lazy(() -> prepareStatement(connection), fetchSize, PreparedStatement::executeQuery, rowMapper);
+  }
+
+  @MustBeClosed
+  private static <S extends Statement, T> Stream<T> lazy(
+      SqlSupplier<? extends S> createStatement,
+      int fetchSize,
+      SqlFunction<? super S, ResultSet> execute,
+      SqlFunction<? super ResultSet, ? extends T> rowMapper) {
+    try (JdbcCloser closer = new JdbcCloser()) {
+      S stmt = createStatement.get();
+      closer.register(stmt::close);
+      stmt.setFetchSize(fetchSize);
+      ResultSet resultSet = execute.apply(stmt);
+      closer.register(resultSet::close);
+      return closer.attachTo(
+          whileNotNull(() -> {
+            try {
+              return resultSet.next() ? new AtomicReference<T>(rowMapper.apply(resultSet)) : null;
+            } catch (SQLException e) {
+              throw new UncheckedSqlException(e);
+            }
+          })
+          .map(AtomicReference::get));
+    } catch (SQLException e) {
+      throw new UncheckedSqlException(e);
+    }
+  }
+
   /**
    * Executes the encapsulated DML (create, update, delete statements) against {@code connection}
    * and returns the number of affected rows.

mug-guava/src/main/java/com/google/mu/safesql/SafeSql.java

@@ -0,0 +1,9 @@
+package com.google.mu.safesql;
+
+import java.sql.SQLException;
+
+/** For lambdas that can throw {@link SQLException}. */
+@FunctionalInterface
+public interface SqlSupplier<T> {
+  T get() throws SQLException;
+}

mug-guava/src/main/java/com/google/mu/safesql/SqlSupplier.java

@@ -0,0 +1,72 @@
+package com.google.mu.safesql;
+
+import static org.junit.Assert.assertThrows;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+
+import java.sql.Statement;
+import java.util.stream.Stream;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+import com.google.testing.junit.testparameterinjector.TestParameterInjector;
+
+@RunWith(TestParameterInjector.class)
+public class JdbcCloserTest {
+  @Mock private Statement statement;
+
+  @Before public void setUpMocks() {
+    MockitoAnnotations.initMocks(this);
+  }
+
+  @Test public void nothingRegistered() {
+    try (JdbcCloser closer = new JdbcCloser()) {}
+  }
+
+  @SuppressWarnings("MustBeClosedChecker")
+  @Test public void attachTo_nothingRegistered() {
+    Stream<String> stream = Stream.of("foo", "bar");
+    try (JdbcCloser closer = new JdbcCloser()) {
+      stream = closer.attachTo(stream);
+    }
+    try (Stream<String> closeMe = stream) {}
+  }
+
+  @Test public void resourceClosed() throws Exception {
+    try (JdbcCloser closer = new JdbcCloser()) {
+      closer.register(statement::close);
+    }
+    verify(statement).close();
+  }
+
+  @SuppressWarnings("MustBeClosedChecker")
+  @Test public void resourceAttached() throws Exception {
+    Stream<String> stream = Stream.of("foo", "bar");
+    try (JdbcCloser closer = new JdbcCloser()) {
+      closer.register(statement::close);
+      stream = closer.attachTo(stream);
+    }
+    verify(statement, never()).close();
+
+    // Now close the stream
+    try (Stream<?> closeMe = stream) {}
+    verify(statement).close();
+  }
+
+  @SuppressWarnings("MustBeClosedChecker")
+  @Test public void cannotAttachTwice() throws Exception {
+    Stream<String> stream = Stream.of("foo", "bar");
+    try (JdbcCloser closer = new JdbcCloser()) {
+      closer.register(statement::close);
+      stream = closer.attachTo(stream);
+      assertThrows(IllegalStateException.class, () -> closer.attachTo(Stream.empty()));
+    }
+    verify(statement, never()).close();
+    try (Stream<?> closeMe = stream) {}
+    verify(statement).close();
+  }
+}

mug-guava/src/test/java/com/google/mu/safesql/JdbcCloserTest.java

@@ -9,6 +9,8 @@
 import java.time.ZoneId;
 import java.time.ZonedDateTime;
 import java.util.List;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 import javax.sql.DataSource;
 
@@ -72,6 +74,23 @@ protected DatabaseOperation getTearDownOperation() {
         .containsExactly("bar");
   }
 
+  @Test public void queryLazily() throws Exception {
+    ZonedDateTime barTime = ZonedDateTime.of(2024, 11, 1, 10, 20, 30, 40, ZoneId.of("UTC"));
+    assertThat(
+            SafeSql.of("insert into ITEMS(id, title) VALUES({id}, {title})", testId(), "foo")
+                .update(connection()))
+        .isEqualTo(1);
+    assertThat(
+            SafeSql.of("insert into ITEMS(id, title, time) VALUES({id}, {title}, {time})", testId() + 1, "bar", barTime)
+                .update(connection()))
+        .isEqualTo(1);
+    assertThat(queryColumnStream(
+            SafeSql.of("select title from {tbl} where id = {id}", /* tbl */ SafeSql.of("ITEMS"), testId()), "title"))
+        .containsExactly("foo");
+    assertThat(queryColumnStream(SafeSql.of("select title from ITEMS where id = {id}", testId() + 1), "title"))
+        .containsExactly("bar");
+  }
+
   @Test public void likeExpressionWithWildcardInArg() throws Exception {
     assertThat(
             SafeSql.of("insert into ITEMS(id, title) VALUES({id}, {title})", testId(), "foo")
@@ -312,6 +331,12 @@ private List<?> queryColumn(SafeSql sql, String column) throws Exception {
     return sql.query(connection(), resultSet -> resultSet.getObject(column));
   }
 
+  private List<?> queryColumnStream(SafeSql sql, String column) throws Exception {
+    try (Stream<?> stream = sql.queryLazily(connection(), 1, resultSet -> resultSet.getObject(column))) {
+      return stream.collect(Collectors.toList());
+    }
+  }
+
   private Connection connection() throws Exception {
     return getConnection().getConnection();
   }