improve expiration args (#46)

Author: malmans2

cads_worker/entry_points.py

@@ -1,6 +1,6 @@
 import datetime
 import os
-from typing import Annotated, Optional
+from typing import Annotated
 
 import cacholote
 import structlog
@@ -32,31 +32,34 @@ def _cache_cleaner() -> None:
         raise
 
 
-def _add_tzinfo(timestamp: datetime.datetime | None) -> datetime.datetime | None:
-    if timestamp is not None and timestamp.tzinfo is None:
+def _add_tzinfo(timestamp: datetime.datetime) -> datetime.datetime:
+    if timestamp.tzinfo is None:
         timestamp = timestamp.replace(tzinfo=datetime.timezone.utc)
     return timestamp
 
 
 def _expire_cache_entries(
-    collection_id: Annotated[
-        Optional[list[str]], Option(help="Collection ID to expire", show_default="all")
-    ] = None,
     before: Annotated[
-        Optional[datetime.datetime],
-        Option(
-            help="Expire entries created before this date",
-            show_default="no bound",
-        ),
-    ] = None,
+        datetime.datetime,
+        Option(help="Expire entries created before this date"),
+    ],
     after: Annotated[
-        Optional[datetime.datetime],
-        Option(help="Expire entries created after this date", show_default="no bound"),
-    ] = None,
+        datetime.datetime,
+        Option(help="Expire entries created after this date"),
+    ],
+    collection_id: Annotated[list[str], Option(help="Collection ID to expire")] = [],
+    all_collections: Annotated[
+        bool, Option("--all-collections", help="Expire all collections")
+    ] = False,
 ) -> int:
     """Expire cache entries."""
+    if (all_collections and collection_id) or not (all_collections or collection_id):
+        raise ValueError(
+            "Either '--collection-id' or '--all-collections' must be chosen, but not both."
+        )
+
     count = cacholote.expire_cache_entries(
-        tags=collection_id,
+        tags=None if all_collections else collection_id,
         before=_add_tzinfo(before),
         after=_add_tzinfo(after),
     )

tests/test_20_expire_cache_entries.py

@@ -1,14 +1,14 @@
+import contextlib
 import datetime
-import pathlib
+from pathlib import Path
+from typing import Any
 
 import cacholote
 import pytest
 
 from cads_worker import entry_points
 
-TODAY = datetime.datetime.now(tz=datetime.timezone.utc)
-TOMORROW = TODAY + datetime.timedelta(days=1)
-YESTERDAY = TODAY - datetime.timedelta(days=1)
+does_not_raise = contextlib.nullcontext
 
 
 @cacholote.cacheable
@@ -17,31 +17,37 @@ def cached_now() -> datetime.datetime:
 
 
 @pytest.mark.parametrize(
-    "collection_id,before,after",
+    "collection_id,all_collections,raises",
     [
-        (["foo"], None, None),
-        (None, TOMORROW, None),
-        (None, None, YESTERDAY),
-        (["foo"], TOMORROW, YESTERDAY),
+        (["foo"], False, does_not_raise()),
+        ([], True, does_not_raise()),
+        ([], False, pytest.raises(ValueError)),
+        (["foo"], True, pytest.raises(ValueError)),
     ],
 )
 def test_cache_entries(
-    tmp_path: pathlib.Path,
-    collection_id: list[str] | None,
-    before: datetime.datetime | None,
-    after: datetime.datetime | None,
+    tmp_path: Path,
+    collection_id: list[str],
+    all_collections: bool,
+    raises: contextlib.nullcontext[Any],
 ) -> None:
+    today = datetime.datetime.now(tz=datetime.timezone.utc)
+    tomorrow = today + datetime.timedelta(days=1)
+    yeasterday = today - datetime.timedelta(days=1)
+
     with cacholote.config.set(
         cache_db_urlpath=f"sqlite:///{tmp_path / 'cacholote.db'}",
         tag="foo",
     ):
         now = cached_now()
         assert now == cached_now()
 
-        count = entry_points._expire_cache_entries(
-            collection_id=collection_id,
-            before=before,
-            after=after,
-        )
-        assert count == 1
-        assert now != cached_now()
+        with raises:
+            count = entry_points._expire_cache_entries(
+                before=tomorrow,
+                after=yeasterday,
+                collection_id=collection_id,
+                all_collections=all_collections,
+            )
+            assert count == 1
+            assert now != cached_now()