refactor: replace current windowed query

tests/unit/utils/db/test_windowed_query.py

@@ -14,6 +14,9 @@
 
 import pytest
 
+from sqlalchemy import select
+from sqlalchemy.orm import aliased
+
 from warehouse.packaging.models import Project
 from warehouse.utils.db.windowed_query import windowed_query
 
@@ -22,11 +25,20 @@
 
 @pytest.mark.parametrize("window_size", [1, 2])
 def test_windowed_query(db_session, query_recorder, window_size):
-    projects = set(ProjectFactory.create_batch(10))
+    projects = ProjectFactory.create_batch(10)
+    project_set = {(project.name, project.id) for project in projects}
+
     expected = math.ceil(len(projects) / window_size) + 1
 
-    query = db_session.query(Project)
+    subquery = select(Project.normalized_name).order_by(Project.id).subquery()
+    pa = aliased(Project, subquery)
+
+    query = select(Project.name).select_from(pa).distinct(Project.id)
+
     with query_recorder:
-        assert set(windowed_query(query, Project.name, window_size)) == projects
+        assert (
+            set(windowed_query(db_session, query, Project.id, window_size))
+            == project_set
+        )
 
     assert len(query_recorder.queries) == expected

warehouse/search/tasks.py

@@ -21,7 +21,8 @@
 
 from elasticsearch.helpers import parallel_bulk
 from elasticsearch_dsl import serializer
-from sqlalchemy import func, text
+from sqlalchemy import func, select, text
+from sqlalchemy.orm import aliased
 
 from warehouse import tasks
 from warehouse.packaging.models import (
@@ -38,7 +39,7 @@
 
 def _project_docs(db, project_name=None):
     releases_list = (
-        db.query(Release.id)
+        select(Release.id)
         .filter(Release.yanked.is_(False), Release.files)
         .order_by(
             Release.project_id,
@@ -52,9 +53,10 @@ def _project_docs(db, project_name=None):
         releases_list = releases_list.join(Project).filter(Project.name == project_name)
 
     releases_list = releases_list.subquery()
+    rlist = aliased(Release, releases_list)
 
     classifiers = (
-        db.query(func.array_agg(Classifier.classifier))
+        select(func.array_agg(Classifier.classifier))
         .select_from(release_classifiers)
         .join(Classifier, Classifier.id == release_classifiers.c.trove_id)
         .filter(Release.id == release_classifiers.c.release_id)
@@ -64,7 +66,7 @@ def _project_docs(db, project_name=None):
     )
 
     release_data = (
-        db.query(
+        select(
             Description.raw.label("description"),
             Release.version.label("latest_version"),
             Release.author,
@@ -81,13 +83,13 @@ def _project_docs(db, project_name=None):
             Project.normalized_name,
             Project.name,
         )
-        .select_from(releases_list)
-        .join(Release, Release.id == releases_list.c.id)
+        .select_from(rlist)
+        .join(Release, Release.id == rlist.id)
         .join(Description)
         .outerjoin(Release.project)
     )
 
-    for release in windowed_query(release_data, Project.id, 25000):
+    for release in windowed_query(db, release_data, Project.id, 25000):
         p = ProjectDocument.from_db(release)
         p._index = None
         p.full_clean()

warehouse/utils/db/windowed_query.py

@@ -12,54 +12,22 @@
 
 # Taken from "Theatrum Chemicum" at
 # https://github.com/sqlalchemy/sqlalchemy/wiki/RangeQuery-and-WindowedRangeQuery
+# updated from, with minor tweaks:
+# https://github.com/sqlalchemy/sqlalchemy/discussions/7948#discussioncomment-2597083
 
-from sqlalchemy import and_, func, text
 
+def windowed_query(s, q, column, windowsize):
+    """Break a Query into chunks on a given column."""
 
-def column_windows(session, column, windowsize):
-    """
-    Return a series of WHERE clauses against a given column that break it into
-    windows.
+    q = q.add_columns(column).order_by(column)
+    last_id = None
 
-    Result is an iterable of tuples, consisting of ((start, end), whereclause),
-    where (start, end) are the ids.
-
-    Requires a database that supports window functions, i.e. Postgresql,
-    SQL Server, Oracle.
-
-    Enhance this yourself !  Add a "where" argument so that windows of just a
-    subset of rows can be computed.
-    """
-
-    def int_for_range(start_id, end_id):
-        if end_id:
-            return and_(column >= start_id, column < end_id)
-        else:
-            return column >= start_id
-
-    q = session.query(
-        column, func.row_number().over(order_by=column).label("rownum")
-    ).from_self(column)
-
-    if windowsize > 1:
-        q = q.filter(text("rownum %% %d=1" % windowsize))
-
-    intervals = [row[0] for row in q]
-
-    while intervals:
-        start = intervals.pop(0)
-        if intervals:
-            end = intervals[0]
-        else:
-            end = None
-
-        yield int_for_range(start, end)
-
-
-def windowed_query(q, column, windowsize):
-    """
-    Break a Query into windows on a given column.
-    """
-
-    for whereclause in column_windows(q.session, column, windowsize):
-        yield from q.filter(whereclause).order_by(column)
+    while True:
+        subq = q
+        if last_id is not None:
+            subq = subq.filter(column > last_id)
+        chunk = s.execute(subq.limit(windowsize)).all()
+        if not chunk:
+            break
+        last_id = chunk[-1][-1]
+        yield from chunk