diff --git a/swh/scheduler/backend.py b/swh/scheduler/backend.py --- a/swh/scheduler/backend.py +++ b/swh/scheduler/backend.py @@ -934,8 +934,10 @@ """, OriginVisitStats.select_columns(), ) - psycopg2.extras.execute_values(cur=cur, sql=query, argslist=primary_keys) - return [OriginVisitStats(**row) for row in cur.fetchall()] + rows = psycopg2.extras.execute_values( + cur=cur, sql=query, argslist=primary_keys, fetch=True + ) + return [OriginVisitStats(**row) for row in rows] @db_transaction() def update_metrics( diff --git a/swh/scheduler/tests/test_scheduler.py b/swh/scheduler/tests/test_scheduler.py --- a/swh/scheduler/tests/test_scheduler.py +++ b/swh/scheduler/tests/test_scheduler.py @@ -12,6 +12,7 @@ import uuid import attr +from psycopg2.extras import execute_values import pytest from swh.model.hashutil import hash_to_bytes @@ -943,6 +944,31 @@ def test_origin_visit_stats_get_empty(self, swh_scheduler) -> None: assert swh_scheduler.origin_visit_stats_get([]) == [] + def test_origin_visit_stats_get_pagination(self, swh_scheduler) -> None: + page_size = inspect.signature(execute_values).parameters["page_size"].default + + visit_stats = [ + OriginVisitStats( + url=f"https://example.com/origin-{i:03d}", + visit_type="git", + last_eventful=utcnow(), + last_uneventful=None, + last_failed=None, + last_notfound=None, + ) + for i in range( + page_size + 1 + ) # Ensure overflow of the psycopg2.extras.execute_values page_size + ] + + swh_scheduler.origin_visit_stats_upsert(visit_stats) + + assert set( + swh_scheduler.origin_visit_stats_get( + [(ovs.url, ovs.visit_type) for ovs in visit_stats] + ) + ) == set(visit_stats) + def test_origin_visit_stats_upsert(self, swh_scheduler) -> None: eventful_date = utcnow() url = "https://github.com/test"