Add pagination to SafeSql

Author: fluentfuture

mug-safesql/src/main/java/com/google/mu/safesql/OrderBy.java

@@ -0,0 +1,86 @@
+package com.google.mu.safesql;
+
+import static com.google.mu.safesql.SafeSqlUtils.checkArgument;
+import static com.google.mu.util.Substring.last;
+import static java.util.stream.Collectors.toList;
+
+import java.util.HashSet;
+import java.util.List;
+import java.util.Locale;
+import java.util.Map;
+import java.util.Set;
+import java.util.stream.Collector;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import com.google.mu.util.StringFormat;
+import com.google.mu.util.Substring;
+import com.google.mu.util.stream.BiStream;
+
+final class OrderBy {
+  private static final StringFormat QUOTED_IDENTIFIER = new StringFormat("\"{identifier}\"");
+
+  final String columnName;
+  private final boolean descending;
+ 
+  private OrderBy(String columnName, boolean descending) {
+    this.columnName =
+        QUOTED_IDENTIFIER.parse(columnName, identifier -> identifier).orElse(columnName);
+    this.descending = descending;
+  }
+  
+  SafeSql sql() {
+    return descending
+        ? SafeSql.of("\"{column}\" DESC", columnName)
+        : SafeSql.of("\"{column}\"", columnName);
+  }
+  
+  static Collector<String, ?, List<OrderBy>> parsingToDistinctOrderByList() {
+    return Collectors.collectingAndThen(
+        toList(),
+        columns -> {
+          Set<String> distinctColumns = new HashSet<>();
+          return columns.stream()
+              .map(column -> {
+                OrderBy orderBy = parse(column);
+                checkArgument(
+                    distinctColumns.add(orderBy.columnName),
+                    "Column %s specified twice in order by", orderBy.columnName);
+                return orderBy;
+              })
+              .collect(toList());
+        });
+  }
+  
+  static SafeSql after(Map<String, ?> columnValues, List<OrderBy> orderByList) {
+    checkArgument(orderByList.size() > 0, "orderByList cannot be empty");
+    return IntStream.rangeClosed(1, orderByList.size())
+        .mapToObj(prefixSize -> after(columnValues, orderByList, prefixSize))
+        .collect(SafeSql.or());
+  }
+  
+  private static SafeSql after(Map<String, ?> columnValues, List<OrderBy> orderByList, int prefixSize) {
+    checkArgument(prefixSize > 0, "prefixSize must be positive");
+    SafeSql equalTo = SafeSql.of("=");
+    return BiStream.zip(IntStream.rangeClosed(1, prefixSize).boxed(), orderByList.stream())
+        .mapKeys((i, item) -> i < prefixSize ? equalTo : item.after())
+        .mapToObj((relativeTo, item) -> item.is(relativeTo, columnValues.get(item.columnName)))
+        .collect(SafeSql.and());
+  }
+  
+  private SafeSql is(SafeSql relativeTo, Object pageEnd) {
+    return SafeSql.of("\"{column}\" {relative_to} {page_end}", columnName, relativeTo, pageEnd);
+  }
+  
+  private SafeSql after() {
+    return descending ? SafeSql.of("<") : SafeSql.of(">");
+  }
+  
+  private static OrderBy parse(String orderBy) {
+    orderBy = orderBy.trim();
+    return last('"').or(Substring.BEGINNING).then(last(Character::isWhitespace))
+        .splitThenTrim(orderBy)
+        .map((name, dir) -> new OrderBy(name, dir.toUpperCase(Locale.US).equals("DESC")))
+        .orElse(new OrderBy(orderBy, false));
+  }
+}

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

@@ -30,6 +30,7 @@
 import static java.util.stream.Collectors.collectingAndThen;
 import static java.util.stream.Collectors.mapping;
 import static java.util.stream.Collectors.toCollection;
+import static java.util.stream.Collectors.toList;
 
 import java.sql.Connection;
 import java.sql.PreparedStatement;
@@ -38,9 +39,12 @@
 import java.sql.Statement;
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.HashMap;
 import java.util.Iterator;
+import java.util.LinkedHashMap;
 import java.util.LinkedHashSet;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.Spliterators;
@@ -69,6 +73,7 @@
 import com.google.mu.util.StringFormat.Template;
 import com.google.mu.util.Substring;
 import com.google.mu.util.stream.BiStream;
+import com.google.mu.util.stream.Iteration;
 
 /**
  * An injection-safe <em>dynamic SQL</em>, constructed using compile-time enforced templates.
@@ -959,6 +964,223 @@ public <T> Optional<T> queryForOne(
       throw e.asChecked();
     }
   }
+  
+  /**
+   * Paginate the encapsulated query by fetching up to {@code pageSize} rows for each page.
+   *
+   * <p>If {@code orderByKeys} is non-empty, the column values are used to control pagination.
+   * For example, if {@code [groupId, timestamp]} can uniquely identify a row, the generated
+   * SQL will use {@code "groupId" > ? OR ("groupId" = ? AND "timestamp" > ?)} for the next page,
+   * comparing against the previous page's last row.
+   * 
+   * <p>You can also use descending order, for example by passing {@code "NAME DESC"}.
+   * The generated paginated query will then use {@code "NAME" < ?} instead of {@code "NAME > ?}
+   * for the next page. 
+   * 
+   * <p>The standard SQL {@code FETCH NEXT ROWS} syntax is used for pagination.
+   * If {@code orderByKeys} is empty, the paginated query will use {@code OFFSET} keyword
+   * to specify the next page offset.
+   * 
+   * <p>All column names are automatically double-quoted to avoid clashing with keywords or spaces
+   * in identifiers. If your DB is case sensitive, make sure you pass in the accurate case.
+   * 
+   * <p>The method is lazy. Paginated queries are only sent when the returned stream is consumed.
+   * The caller should close the {@code connection} after done, which will close the cached
+   * JDBC statements used for pagination.
+   *
+   * @param orderByKeys for example {@code "firstName"} or {@code "group_id DESC"}.
+   * @param pageSize must be positive.
+   * @param resultType the record or class type to be constructed for each row.
+   * @return a lazy stream of pages. If the query returns no rows, the stream will be empty.
+   * @throws IllegalArgumentException if {@code pageSize} isn't positive.
+   * @throws NullPointerException if any parameter is null.
+   * @since 9.9.3
+   */
+  public <T> Stream<List<T>> paginateLazily(
+      Connection connection,
+      Collection<String> orderByKeys,
+      int pageSize,
+      Class<? extends T> resultType) {
+    return paginateLazily(connection, stmt -> {}, orderByKeys, pageSize, resultType);
+  }
+  
+  /**
+   * Paginate the encapsulated query by fetching up to {@code pageSize} rows for each page.
+   *
+   * <p>If {@code orderByKeys} is non-empty, the column values are used to control pagination.
+   * For example, if {@code [groupId, timestamp]} can uniquely identify a row, the generated
+   * SQL will use {@code "groupId" > ? OR ("groupId" = ? AND "timestamp" > ?)} for the next page,
+   * comparing against the previous page's last row.
+   * 
+   * <p>You can also use descending order, for example by passing {@code "NAME DESC"}.
+   * The generated paginated query will then use {@code "NAME" < ?} instead of {@code "NAME > ?}
+   * for the next page. 
+   * 
+   * <p>The standard SQL {@code FETCH NEXT ROWS} syntax is used for pagination.
+   * If {@code orderByKeys} is empty, the paginated query will use {@code OFFSET} keyword
+   * to specify the next page offset.
+   * 
+   * <p>All column names are automatically double-quoted to avoid clashing with keywords or spaces
+   * in identifiers. If your DB is case sensitive, make sure you pass in the accurate case.
+   * 
+   * <p>The method is lazy. Paginated queries are only sent when the returned stream is consumed.
+   * The caller should close the {@code connection} after done, which will close the cached
+   * JDBC statements used for pagination.
+   *
+   * @param settings to configure the {@code PreparedStatement} used for the paginated queries.
+   * @param orderByKeys for example {@code "firstName"} or {@code "group_id DESC"}.
+   * @param pageSize must be positive.
+   * @param resultType the record or class type to be constructed for each row.
+   * @return a lazy stream of pages. If the query returns no rows, the stream will be empty.
+   * @throws IllegalArgumentException if {@code pageSize} isn't positive.
+   * @throws NullPointerException if any parameter is null.
+   * @since 9.9.3
+   */
+  public <T> Stream<List<T>> paginateLazily(
+      Connection connection,
+      StatementSettings settings,
+      Collection<String> orderByKeys,
+      int pageSize,
+      Class<? extends T> resultType) {
+    return paginateLazily(
+        connection, settings, orderByKeys, pageSize, ResultMapper.toResultOf(resultType)::from);
+  }
+  
+  /**
+   * Paginate the encapsulated query by fetching up to {@code pageSize} rows for each page.
+   *
+   * <p>If {@code orderByKeys} is non-empty, the column values are used to control pagination.
+   * For example, if {@code [groupId, timestamp]} can uniquely identify a row, the generated
+   * SQL will use {@code "groupId" > ? OR ("groupId" = ? AND "timestamp" > ?)} for the next page,
+   * comparing against the previous page's last row.
+   * 
+   * <p>You can also use descending order, for example by passing {@code "NAME DESC"}.
+   * The generated paginated query will then use {@code "NAME" < ?} instead of {@code "NAME > ?}
+   * for the next page. 
+   * 
+   * <p>The standard SQL {@code FETCH NEXT ROWS} syntax is used for pagination.
+   * If {@code orderByKeys} is empty, the paginated query will use {@code OFFSET} keyword
+   * to specify the next page offset.
+   * 
+   * <p>All column names are automatically double-quoted to avoid clashing with keywords or spaces
+   * in identifiers. If your DB is case sensitive, make sure you pass in the accurate case.
+   * 
+   * <p>The method is lazy. Paginated queries are only sent when the returned stream is consumed.
+   * The caller should close the {@code connection} after done, which will close the cached
+   * JDBC statements used for pagination.
+   *
+   * @param orderByKeys for example {@code "firstName"} or {@code "group_id DESC"}.
+   * @param pageSize must be positive.
+   * @param rowMapper maps each row (of type {@code ResultSet}) to a new result object.
+   * @return a lazy stream of pages. If the query returns no rows, the stream will be empty.
+   * @throws IllegalArgumentException if {@code pageSize} isn't positive.
+   * @throws NullPointerException if any parameter is null.
+   * @since 9.9.3
+   */
+  public <T> Stream<List<T>> paginateLazily(
+      Connection connection,
+      Collection<String> orderByKeys,
+      int pageSize,
+      SqlFunction<? super ResultSet, ? extends T> rowMapper) {
+    return paginateLazily(connection, stmt -> {}, orderByKeys, pageSize, rowMapper);
+  }
+
+  /**
+   * Paginate the encapsulated query by fetching up to {@code pageSize} rows for each page.
+   *
+   * <p>If {@code orderByKeys} is non-empty, the column values are used to control pagination.
+   * For example, if {@code [groupId, timestamp]} can uniquely identify a row, the generated
+   * SQL will use {@code "groupId" > ? OR ("groupId" = ? AND "timestamp" > ?)} for the next page,
+   * comparing against the previous page's last row.
+   * 
+   * <p>You can also use descending order, for example by passing {@code "NAME DESC"}.
+   * The generated paginated query will then use {@code "NAME" < ?} instead of {@code "NAME > ?}
+   * for the next page. 
+   * 
+   * <p>The standard SQL {@code FETCH NEXT ROWS} syntax is used for pagination.
+   * If {@code orderByKeys} is empty, the paginated query will use {@code OFFSET} keyword
+   * to specify the next page offset.
+   * 
+   * <p>All column names are automatically double-quoted to avoid clashing with keywords or spaces
+   * in identifiers. If your DB is case sensitive, make sure you pass in the accurate case.
+   * 
+   * <p>The method is lazy. Paginated queries are only sent when the returned stream is consumed.
+   * The caller should close the {@code connection} after done, which will close the cached
+   * JDBC statements used for pagination.
+   *
+   * @param settings to configure the {@code PreparedStatement} used for the paginated queries.
+   * @param orderByKeys for example {@code "firstName"} or {@code "group_id DESC"}.
+   * @param pageSize must be positive.
+   * @param rowMapper maps each row (of type {@code ResultSet}) to a new result object.
+   * @return a lazy stream of pages. If the query returns no rows, the stream will be empty.
+   * @throws IllegalArgumentException if {@code pageSize} isn't positive.
+   * @throws NullPointerException if any parameter is null.
+   * @since 9.9.3
+   */
+  public <T> Stream<List<T>> paginateLazily(
+      Connection connection,
+      StatementSettings settings,
+      Collection<String> orderByKeys,
+      int pageSize,
+      SqlFunction<? super ResultSet, ? extends T> rowMapper) {
+    requireNonNull(connection);
+    requireNonNull(settings);
+    requireNonNull(rowMapper);
+    List<OrderBy> orderByList = orderByKeys.stream().collect(OrderBy.parsingToDistinctOrderByList());
+    checkArgument(pageSize > 0, "pageSize (%s) isn't positive", pageSize);
+    class Pagination extends Iteration<List<T>> {
+      Map<String, Object> lastPageKeys = new HashMap<>();
+      int offset = 0;
+      final Template<List<T>> fetch = prepare(
+          connection,
+          "SELECT * FROM ({this})"
+              + " WHERE {seek} {order_by? -> ORDER BY order_by?}"
+              + " OFFSET {offset} ROWS FETCH NEXT {page_size} ROWS ONLY",
+          stmt -> {
+            try (ResultSet resultSet = stmt.executeQuery()) {
+              List<T> results = mapResults(resultSet, row -> {
+                lastPageKeys = new LinkedHashMap<>();
+                for (OrderBy orderBy : orderByList) {
+                  Object columnValue = row.getObject(orderBy.columnName);
+                  if (columnValue == null) {
+                    throw new NullPointerException(
+                        "null value not supported for order by column " + orderBy.columnName);
+                  }
+                  lastPageKeys.put(orderBy.columnName, columnValue);
+                }
+                return rowMapper.apply(row);
+              });
+              if (orderByList.isEmpty()) {
+                // If order by is present, we always seek and start from offset 0.
+                offset += results.size();
+              }
+              return results;
+            }
+          });
+      
+      Pagination() {
+        lazily(this::fetchPage);
+      }
+      
+      void fetchPage() {
+        SafeSql seek = lastPageKeys.isEmpty() ? TRUE : OrderBy.after(lastPageKeys, orderByList);
+        List<T> page = fetch.with(
+            SafeSql.this,
+            seek,
+            orderByList.stream().map(OrderBy::sql).collect(toList()),
+            offset,
+            pageSize);
+        if (page.isEmpty()) {  // no more data
+          return;
+        }
+        emit(page);
+        if (page.size() >= pageSize) {
+          lazily(this::fetchPage);
+        }
+      }
+    }
+    return new Pagination().iterate();
+  }
 
   /**
    * Executes the encapsulated SQL as a query against {@code connection},