Hotfix: allow getting disk usage from database

ecmwf-projects · Sep 4, 2024 · 2019f1a · 2019f1a
1 parent 5981fd7
commit 2019f1a
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 22 deletions.
diff --git a/cads_worker/entry_points.py b/cads_worker/entry_points.py
@@ -1,6 +1,6 @@
 import datetime
 import os
-from typing import Annotated
+from typing import Annotated, TypedDict
 
 import cacholote
 import structlog
@@ -14,19 +14,35 @@
 cacholote.config.set(logger=LOGGER)
 
 
+def strtobool(value: str) -> bool:
+    if value.lower() in ("y", "yes", "t", "true", "on", "1"):
+        return True
+    if value.lower() in ("n", "no", "f", "false", "off", "0"):
+        return False
+    raise ValueError(f"invalid truth value {value!r}")
+
+
+class CleanerKwargs(TypedDict):
+    maxsize: int
+    method: str
+    delete_unknown_files: bool
+    lock_validity_period: float
+    use_database: bool
+
+
 def _cache_cleaner() -> None:
-    max_size = int(os.environ.get("MAX_SIZE", 1_000_000_000))
     cache_bucket = os.environ.get("CACHE_BUCKET", None)
-    LOGGER.info("Running cache cleaner", max_size=max_size, cache_bucket=cache_bucket)
+    use_database = strtobool(os.environ.get("USE_DATABASE", "1"))
+    cleaner_kwargs = CleanerKwargs(
+        maxsize=int(os.environ.get("MAX_SIZE", 1_000_000_000)),
+        method=os.environ.get("METHOD", "LRU"),
+        delete_unknown_files=not use_database,
+        lock_validity_period=float(os.environ.get("LOCK_VALIDITY_PERIOD", 86400)),
+        use_database=use_database,
+    )
+    LOGGER.info("Running cache cleaner", cache_bucket=cache_bucket, **cleaner_kwargs)
     try:
-        cacholote.clean_cache_files(
-            maxsize=max_size,
-            method=os.environ.get("METHOD", "LRU"),  # type: ignore[arg-type] # let cacholote handle it
-            delete_unknown_files=bool(os.environ.get("DELETE_UNKNOWN_FILES", 1)),
-            lock_validity_period=float(
-                os.environ.get("LOCK_VALIDITY_PERIOD", 60 * 60 * 24)
-            ),
-        )
+        cacholote.clean_cache_files(**cleaner_kwargs)
     except Exception:
         LOGGER.exception("cache_cleaner crashed")
         raise

diff --git a/tests/test_10_cache_cleaner.py b/tests/test_10_cache_cleaner.py
@@ -1,11 +1,16 @@
-import os
 import pathlib
 import subprocess
 
 import cacholote
+import pytest
 
 
-def test_cache_cleaner(tmp_path: pathlib.Path) -> None:
+@pytest.mark.parametrize("use_database", ["true", "false"])
+def test_cache_cleaner(
+    tmp_path: pathlib.Path,
+    monkeypatch: pytest.MonkeyPatch,
+    use_database: str,
+) -> None:
     # create dummy file
     dummy_path = tmp_path / "dummy.txt"
     dummy_path.write_text("dummy")
@@ -21,13 +26,10 @@ def test_cache_cleaner(tmp_path: pathlib.Path) -> None:
     assert cached_path.exists()
 
     # clean cache
-    cache_env = os.environ.copy()
-    cache_env.update(
-        {
-            "MAX_SIZE": "0",
-            "CACHOLOTE_CACHE_DB_URLPATH": cache_db_urlpath,
-            "CACHOLOTE_CACHE_FILES_URLPATH": cache_files_urlpath,
-        }
-    )
-    subprocess.run("cache-cleaner", check=True, env=cache_env)
+    monkeypatch.setenv("MAX_SIZE", "0")
+    monkeypatch.setenv("USE_DATABASE", use_database)
+    monkeypatch.setenv("CACHOLOTE_CACHE_DB_URLPATH", cache_db_urlpath)
+    monkeypatch.setenv("CACHOLOTE_CACHE_FILES_URLPATH", cache_files_urlpath)
+    monkeypatch.setenv("USE_DATABASE", use_database)
+    subprocess.run("cache-cleaner", check=True)
     assert not cached_path.exists()