From bfa6a94a6793e46eea8edc2c0e90386467e8d13d Mon Sep 17 00:00:00 2001
From: Mattia Almansi <m.almansi@bopen.eu>
Date: Tue, 23 Jul 2024 12:34:19 +0200
Subject: [PATCH] add expire cache entries (#120)

---
 cacholote/__init__.py  |  8 +++++++-
 cacholote/clean.py     | 24 ++++++++++++++++++++++++
 tests/test_60_clean.py | 34 ++++++++++++++++++++++++++++++++++
 3 files changed, 65 insertions(+), 1 deletion(-)

diff --git a/cacholote/__init__.py b/cacholote/__init__.py
index 570a2f1..8cd06d2 100644
--- a/cacholote/__init__.py
+++ b/cacholote/__init__.py
@@ -17,7 +17,12 @@
 
 from . import config, database, extra_encoders, utils
 from .cache import cacheable
-from .clean import clean_cache_files, clean_invalid_cache_entries, delete
+from .clean import (
+    clean_cache_files,
+    clean_invalid_cache_entries,
+    delete,
+    expire_cache_entries,
+)
 from .decode import loads
 from .encode import dumps
 
@@ -40,6 +45,7 @@
     "database",
     "delete",
     "dumps",
+    "expire_cache_entries",
     "extra_encoders",
     "loads",
     "utils",
diff --git a/cacholote/clean.py b/cacholote/clean.py
index a04a6d0..c24cc9a 100644
--- a/cacholote/clean.py
+++ b/cacholote/clean.py
@@ -23,6 +23,7 @@
 import pydantic
 import sqlalchemy as sa
 import sqlalchemy.orm
+from sqlalchemy import BinaryExpression, ColumnElement
 
 from . import config, database, decode, encode, extra_encoders, utils
 
@@ -354,3 +355,26 @@ def clean_invalid_cache_entries(
                     decode.loads(cache_entry._result_as_string)
                 except decode.DecodeError:
                     _delete_cache_entry(session, cache_entry)
+
+
+def expire_cache_entries(
+    tags: list[str] | None = None,
+    before: datetime.datetime | None = None,
+    after: datetime.date | None = None,
+) -> None:
+    now = utils.utcnow()
+
+    filters: list[BinaryExpression[bool] | ColumnElement[bool]] = []
+    if tags is not None:
+        filters.append(database.CacheEntry.tag.in_(tags))
+    if before is not None:
+        filters.append(database.CacheEntry.timestamp < before)
+    if after is not None:
+        filters.append(database.CacheEntry.timestamp > after)
+
+    with config.get().instantiated_sessionmaker() as session:
+        for cache_entry in session.scalars(
+            sa.select(database.CacheEntry).filter(*filters)
+        ):
+            cache_entry.expiration = now
+        database._commit_or_rollback(session)
diff --git a/tests/test_60_clean.py b/tests/test_60_clean.py
index d8f2bf7..d624f86 100644
--- a/tests/test_60_clean.py
+++ b/tests/test_60_clean.py
@@ -16,6 +16,9 @@
 from cacholote import cache, clean, config, utils
 
 ONE_BYTE = os.urandom(1)
+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
 
 
@@ -30,6 +33,11 @@ def open_urls(*urls: pathlib.Path) -> list[fsspec.spec.AbstractBufferedFile]:
     return [fsspec.open(url).open() for url in urls]
 
 
+@cache.cacheable
+def cached_now() -> datetime.datetime:
+    return datetime.datetime.now()
+
+
 @pytest.mark.parametrize("method", ["LRU", "LFU"])
 @pytest.mark.parametrize("set_cache", ["file", "cads"], indirect=True)
 def test_clean_cache_files(
@@ -301,3 +309,29 @@ def test_clean_multiple_files(tmp_path: pathlib.Path) -> None:
 
     clean.clean_cache_files(0)
     assert len(fs.ls(dirname)) == 0
+
+
+@pytest.mark.parametrize(
+    "tags,before,after",
+    [
+        (["foo"], None, None),
+        (None, TOMORROW, None),
+        (None, None, YESTERDAY),
+        (["foo"], TOMORROW, YESTERDAY),
+    ],
+)
+def test_expire_cache_entries(
+    tags: None | list[str],
+    before: None | datetime.datetime,
+    after: None | datetime.datetime,
+) -> None:
+    with config.set(tag="foo"):
+        now = cached_now()
+
+    # Do not expire
+    clean.expire_cache_entries(tags=["bar"], before=YESTERDAY, after=TOMORROW)
+    assert now == cached_now()
+
+    # Expire
+    clean.expire_cache_entries(tags=tags, before=before, after=after)
+    assert now != cached_now()