Impelement get_homogeneous_pages.

sqlakeyset/paging.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 from functools import partial
+from dataclasses import dataclass
 from typing import (
     Any,
     List,
@@ -438,3 +439,87 @@ def get_page(
     place, backwards = process_args(after, before, page)
 
     return orm_get_page(query, per_page, place, backwards)
+
+
+# Do we need to support python 3.6?
+@dataclass
+class PageRequest(Generic[_TP]):
+    """See ``get_page()`` documentation for parameter explanations."""
+    query: Query[_TP]
+    per_page: int = PER_PAGE_DEFAULT
+    after: OptionalKeyset = None
+    before: OptionalKeyset = None
+    page: Optional[Union[MarkerLike, str]] = None
+
+
+def get_homogeneous_page(queries: list[PageRequest[_TP]]) -> list[Page[Row[_TP]]]:
+    """Get multiple pages of results for homogeneous legacy ORM queries.
+
+    This only involves a single round trip to the database. To do that, under the
+    hood it generates a UNION ALL. That means each query must select exactly the
+    same columns. They may have different filters or ordering, but must result in
+    selecting the same columns with the same names.
+
+    Resulting pages are returned in the same order as the original page requests.
+    """
+    if not queries:
+        return []
+
+    paging_queries = [_prepare_homogeneous_page(query, i) for i, query in enumerate(queries)]
+
+    query = paging_queries[0].query.query
+    query = query.union_all(*[q.query.query for q in paging_queries[1:]])
+
+    results = query.all()
+
+    # We need to make sure there's an entry for every page in case some return
+    # empty.
+    page_to_rows = {i: list() for i in range(len(queries))}
+    for row in results:
+        page_to_rows[row._page_identifier].append(row)
+
+    pages = []
+    for i in range(len(queries)):
+        rows = page_to_rows[i]
+        pages.append(paging_queries[i].page_from_rows(rows))
+
+
+@dataclass
+class _HomogeneousPagingQuery:
+    query: _PagingQuery
+    page_from_rows: Callable[[list[Row[_TP]]], Page[Row[_TP]]]
+
+
+def _prepare_homogeneous_page(
+        query: PageRequest[_TP], page_identifier: int
+) -> _HomogeneousPagingQuery:
+    """page_identifier MUST be a trusted int, as it is not escaped."""
+    # Should have no effect if called correctly, but let's be extra safe.
+    page_identifier = int(page_identifier)
+    place, backwards = process_args(query.after, query.before, query.page)
+
+    # Grab result_type and keys before adding the _page_identifier so that
+    # it isn't included in the results.
+    result_type = orm_result_type(query.query)
+    keys = orm_query_keys(query.query)
+    query.query = query.query.add_columns(
+        sa.sql.expression.literal_column(str(page_identifier)).label("_page_identifier")
+    )
+
+    # Could we order by page identifier to do the page collation in the DB?
+
+    paging_query = prepare_paging(
+        q=query.query,
+        per_page=query.per_page,
+        place=place,
+        backwards=backwards,
+        orm=True,
+        dialect=query.query.session.get_bind().dialect,
+    )
+
+    def page_from_rows(rows):
+        return orm_page_from_rows(
+            paging_query, rows, keys, result_type, per_page, backwards, current_place=place
+        )
+
+    return _HomogeneousPagingQuery(query=paging_query, page_from_rows=page_from_rows)