more tests

cacholote/clean.py

@@ -366,7 +366,10 @@ def clean_cache_files(
 
 
 def clean_invalid_cache_entries(
-    check_expiration: bool = True, try_decode: bool = False
+    check_expiration: bool = True,
+    try_decode: bool = False,
+    batch_size: int | None = None,
+    batch_delay: float = 0,
 ) -> None:
     """Clean invalid cache entries.
 
@@ -376,16 +379,27 @@ def clean_invalid_cache_entries(
         Whether or not to delete expired entries
     try_decode: bool
         Whether or not to delete entries that raise DecodeError (this can be slow!)
+    batch_size: int, optional, default: None
+        Number of entries to process in each batch.
+        If None, all entries are processed in a single batch.
+    batch_delay: float, default: 0
+        Time in seconds to wait between processing consecutive batches.
     """
-    filters = []
     if check_expiration:
-        filters.append(database.CacheEntry.expiration <= utils.utcnow())
-    if filters:
+        id_stmt = (
+            sa.select(database.CacheEntry.id)
+            .filter(database.CacheEntry.expiration <= utils.utcnow())
+            .execution_options(yield_per=batch_size)
+        )
         with config.get().instantiated_sessionmaker() as session:
-            for cache_entry in session.scalars(
-                sa.select(database.CacheEntry).filter(*filters)
-            ):
-                _delete_cache_entries(session, cache_entry)
+            partitions = list(session.scalars(id_stmt).partitions())
+        for i, partition in enumerate(partitions):
+            entry_stmt = sa.select(database.CacheEntry).filter(
+                database.CacheEntry.id.in_(partition)
+            )
+            time.sleep(batch_delay if i else 0)
+            with config.get().instantiated_sessionmaker() as session:
+                _delete_cache_entries(session, *list(session.scalars(entry_stmt)))
 
     if try_decode:
         with config.get().instantiated_sessionmaker() as session:

tests/test_60_clean.py

@@ -518,3 +518,38 @@ def test_expire_cache_entries_dry_run(dry_run: bool, cache_entries: int) -> None
 
     cur.execute("SELECT COUNT(*) FROM cache_entries", ())
     assert cur.fetchone() == (cache_entries,)
+
+
+@pytest.mark.parametrize(
+    "batch_size,batch_delay,expected_time",
+    [
+        (1, 0, 1),
+        (2, 0, 1),
+        (1, 1, 2),
+        (2, 1, 1),
+    ],
+)
+def test_clean_invalid_cache_entries_batch(
+    batch_size: int,
+    batch_delay: float,
+    expected_time: float,
+) -> None:
+    con = config.get().engine.raw_connection()
+    cur = con.cursor()
+
+    for i in range(2):
+        cached_now(i)
+    clean.expire_cache_entries(before=TOMORROW, delete=False)
+
+    tic = time.perf_counter()
+    clean.clean_invalid_cache_entries(
+        check_expiration=True,
+        try_decode=False,
+        batch_size=batch_size,
+        batch_delay=batch_delay,
+    )
+    toc = time.perf_counter()
+    assert expected_time - 1 < toc - tic < expected_time
+
+    cur.execute("SELECT COUNT(*) FROM cache_entries", ())
+    assert cur.fetchone() == (0,)