COPDS-1939: multiple endpoints

README.md

@@ -32,7 +32,7 @@ False
 >>> tmpdir = tempfile.TemporaryDirectory().name
 >>> cacholote.config.set(
 ...     cache_db_urlpath="sqlite://",
-...     cache_files_urlpath=tmpdir,
+...     cache_files_urlpaths=[tmpdir],
 ... )
 <cacholote.config.set ...
 
@@ -59,7 +59,7 @@ True
 >>> tmpdir = tempfile.TemporaryDirectory().name
 >>> cacholote.config.set(
 ...     cache_db_urlpath="sqlite://",
-...     cache_files_urlpath=tmpdir,
+...     cache_files_urlpaths=[tmpdir],
 ... )
 <cacholote.config.set ...
 

cacholote/clean.py

@@ -20,6 +20,7 @@
 import posixpath
 from typing import Any, Callable, Literal, Optional
 
+import fsspec
 import pydantic
 import sqlalchemy as sa
 import sqlalchemy.orm
@@ -35,6 +36,10 @@
 )
 
 
+class CleanerError(Exception):
+    pass
+
+
 def _get_files_from_cache_entry(cache_entry: database.CacheEntry) -> dict[str, str]:
     result = cache_entry.result
     if not isinstance(result, (list, tuple, set)):
@@ -55,7 +60,7 @@ def _get_files_from_cache_entry(cache_entry: database.CacheEntry) -> dict[str, s
 def _delete_cache_entry(
     session: sa.orm.Session, cache_entry: database.CacheEntry
 ) -> None:
-    fs, _ = utils.get_cache_files_fs_dirname()
+    fs, _ = utils.get_cache_files_fs_dirnames()
     files_to_delete = _get_files_from_cache_entry(cache_entry)
     logger = config.get().logger
 
@@ -92,37 +97,45 @@ def delete(func_to_del: str | Callable[..., Any], *args: Any, **kwargs: Any) ->
 
 
 class _Cleaner:
-    def __init__(self) -> None:
+    def __init__(
+        self,
+        fs: fsspec.AbstractFileSystem,
+        dirname: str,
+    ) -> None:
         self.logger = config.get().logger
-        self.fs, self.dirname = utils.get_cache_files_fs_dirname()
+        self.fs = fs
+        self.dirname = dirname
 
-        urldir = self.fs.unstrip_protocol(self.dirname)
-
-        self.logger.info("getting disk usage")
+        self.urldir = self.fs.unstrip_protocol(self.dirname)
+        self._log_info("getting disk usage")
         self.file_sizes: dict[str, int] = collections.defaultdict(int)
         for path, size in self.fs.du(self.dirname, total=False).items():
             # Group dirs
             urlpath = self.fs.unstrip_protocol(path)
-            basename, *_ = urlpath.replace(urldir, "", 1).strip("/").split("/")
+            basename, *_ = urlpath.replace(self.urldir, "", 1).strip("/").split("/")
             if basename:
-                self.file_sizes[posixpath.join(urldir, basename)] += size
+                self.file_sizes[posixpath.join(self.urldir, basename)] += size
 
         self.disk_usage = sum(self.file_sizes.values())
         self.log_disk_usage()
 
+    def _log_info(self, *args: Any, **kwargs: Any) -> None:
+        kwargs.setdefault("dirname", self.dirname)
+        self.logger.info(*args, **kwargs)
+
     def pop_file_size(self, file: str) -> int:
         size = self.file_sizes.pop(file, 0)
         self.disk_usage -= size
         return size
 
     def log_disk_usage(self) -> None:
-        self.logger.info("check disk usage", disk_usage=self.disk_usage)
+        self._log_info("check disk usage", disk_usage=self.disk_usage)
 
     def stop_cleaning(self, maxsize: int) -> bool:
         return self.disk_usage <= maxsize
 
     def get_unknown_files(self, lock_validity_period: float | None) -> set[str]:
-        self.logger.info("getting unknown files")
+        self._log_info("getting unknown files")
 
         utcnow = utils.utcnow()
         locked_files = set()
@@ -218,7 +231,7 @@ def remove_files(
         if not files:
             return
 
-        self.logger.info("deleting files", n_files_to_delete=len(files), **kwargs)
+        self._log_info("deleting files", n_files_to_delete=len(files), **kwargs)
 
         n_tries = 0
         while files:
@@ -253,17 +266,17 @@ def delete_cache_files(
             for cache_entry in session.scalars(
                 sa.select(database.CacheEntry).filter(*filters).order_by(*sorters)
             ):
-                files = _get_files_from_cache_entry(cache_entry)
-                if files:
+                urlpaths = _get_files_from_cache_entry(cache_entry)
+                if any(urlpath.startswith(self.urldir) for urlpath in urlpaths):
                     n_entries_to_delete += 1
                     session.delete(cache_entry)
 
-                for file, file_type in files.items():
-                    self.pop_file_size(file)
-                    if file_type == "application/vnd+zarr":
-                        dirs_to_delete.append(file)
-                    else:
-                        files_to_delete.append(file)
+                    for file, file_type in urlpaths.items():
+                        self.pop_file_size(file)
+                        if file_type == "application/vnd+zarr":
+                            dirs_to_delete.append(file)
+                        else:
+                            files_to_delete.append(file)
 
                 if self.stop_cleaning(maxsize):
                     break
@@ -279,7 +292,7 @@ def delete_cache_files(
         self.log_disk_usage()
 
         if not self.stop_cleaning(maxsize):
-            raise ValueError(
+            raise CleanerError(
                 (
                     f"Unable to clean {self.dirname!r}."
                     f" Final disk usage: {self.disk_usage!r}."
@@ -317,17 +330,26 @@ def clean_cache_files(
         To delete/keep untagged entries, add None in the list (e.g., [None, 'tag1', ...]).
         tags_to_clean and tags_to_keep are mutually exclusive.
     """
-    cleaner = _Cleaner()
-
-    if delete_unknown_files:
-        cleaner.delete_unknown_files(lock_validity_period, recursive)
+    fs, dirnames = utils.get_cache_files_fs_dirnames()
+    exceptions = []
+    for dirname in dirnames:
+        cleaner = _Cleaner(fs=fs, dirname=dirname)
+
+        if delete_unknown_files:
+            cleaner.delete_unknown_files(lock_validity_period, recursive)
+
+        try:
+            cleaner.delete_cache_files(
+                maxsize=maxsize,
+                method=method,
+                tags_to_clean=tags_to_clean,
+                tags_to_keep=tags_to_keep,
+            )
+        except CleanerError as exc:
+            exceptions.append(exc)
 
-    cleaner.delete_cache_files(
-        maxsize=maxsize,
-        method=method,
-        tags_to_clean=tags_to_clean,
-        tags_to_keep=tags_to_keep,
-    )
+    if exceptions:
+        raise CleanerError("\n".join(map(str, exceptions)))
 
 
 def clean_invalid_cache_entries(

cacholote/config.py

@@ -18,6 +18,7 @@
 import abc
 import datetime
 import logging
+import os
 import pathlib
 import tempfile
 from types import TracebackType
@@ -33,15 +34,27 @@
 from . import database
 
 _SETTINGS: Settings | None = None
-_DEFAULT_CACHE_DIR = pathlib.Path(tempfile.gettempdir()) / "cacholote"
-_DEFAULT_CACHE_DIR.mkdir(exist_ok=True)
-_DEFAULT_CACHE_DB_URLPATH = f"sqlite:///{_DEFAULT_CACHE_DIR / 'cacholote.db'}"
-_DEFAULT_CACHE_FILES_URLPATH = f"{_DEFAULT_CACHE_DIR / 'cache_files'}"
 _DEFAULT_LOGGER = structlog.get_logger(
     wrapper_class=structlog.make_filtering_bound_logger(logging.WARNING)
 )
 
 
+def _get_tmp_path() -> pathlib.Path:
+    tmp_path = pathlib.Path(tempfile.gettempdir()) / "cacholote"
+    tmp_path.mkdir(exist_ok=True)
+    return tmp_path
+
+
+def _default_cache_db_urlpath() -> str:
+    return f"sqlite:///{_get_tmp_path() / 'cacholote.db'}"
+
+
+def _default_cache_files_urlpaths() -> list[str]:
+    if (config := os.getenv("CACHOLOTE_CACHE_FILES_URLPATHS_CONFIG")) is not None:
+        return pathlib.Path(config).read_text().splitlines()
+    return [str(_get_tmp_path() / "cache_files")]
+
+
 class Context(abc.ABC):
     @abc.abstractmethod
     def __init__(self, *args: Any, **kwargs: Any) -> None: ...
@@ -52,10 +65,14 @@ def upload_log(self, *args: Any, **kwargs: Any) -> None: ...
 
 class Settings(pydantic_settings.BaseSettings):
     use_cache: bool = True
-    cache_db_urlpath: Optional[str] = _DEFAULT_CACHE_DB_URLPATH
+    cache_db_urlpath: Optional[str] = pydantic.Field(
+        default_factory=_default_cache_db_urlpath
+    )
     create_engine_kwargs: dict[str, Any] = {}
     sessionmaker: Optional[sa.orm.sessionmaker[sa.orm.Session]] = None
-    cache_files_urlpath: str = _DEFAULT_CACHE_FILES_URLPATH
+    cache_files_urlpaths: list[str] = pydantic.Field(
+        default_factory=_default_cache_files_urlpaths
+    )
     cache_files_urlpath_readonly: Optional[str] = None
     cache_files_storage_options: dict[str, Any] = {}
     xarray_cache_type: Literal[
@@ -90,12 +107,13 @@ def validate_expiration(
         return expiration
 
     @pydantic.model_validator(mode="after")
-    def make_cache_dir(self) -> Settings:
-        fs, _, (urlpath, *_) = fsspec.get_fs_token_paths(
-            self.cache_files_urlpath,
-            storage_options=self.cache_files_storage_options,
-        )
-        fs.mkdirs(urlpath, exist_ok=True)
+    def make_cache_dirs(self) -> Settings:
+        for cache_files_urlpath in self.cache_files_urlpaths:
+            fs, _, (urlpath, *_) = fsspec.get_fs_token_paths(
+                cache_files_urlpath,
+                storage_options=self.cache_files_storage_options,
+            )
+            fs.mkdirs(urlpath, exist_ok=True)
         return self
 
     @property

cacholote/extra_encoders.py

@@ -24,6 +24,7 @@
 import mimetypes
 import pathlib
 import posixpath
+import random
 import tempfile
 import time
 from collections.abc import Generator
@@ -144,7 +145,7 @@ class FileInfoModel(pydantic.BaseModel):
 def _dictify_file(fs: fsspec.AbstractFileSystem, local_path: str) -> dict[str, Any]:
     settings = config.get()
     href = posixpath.join(
-        settings.cache_files_urlpath_readonly or settings.cache_files_urlpath,
+        settings.cache_files_urlpath_readonly or posixpath.dirname(local_path),
         posixpath.basename(local_path),
     )
     file_dict = {
@@ -318,7 +319,8 @@ def dictify_xr_object(obj: xr.Dataset | xr.DataArray) -> dict[str, Any]:
         root = dask.base.tokenize(obj)
 
     ext = mimetypes.guess_extension(settings.xarray_cache_type, strict=False)
-    urlpath_out = posixpath.join(settings.cache_files_urlpath, f"{root}{ext}")
+    cache_files_urlpath = random.choice(settings.cache_files_urlpaths)
+    urlpath_out = posixpath.join(cache_files_urlpath, f"{root}{ext}")
 
     fs_out, *_ = fsspec.get_fs_token_paths(
         urlpath_out,
@@ -395,7 +397,7 @@ def dictify_io_object(obj: _UNION_IO_TYPES) -> dict[str, Any]:
     """Encode a file object to JSON deserialized data (``dict``)."""
     settings = config.get()
 
-    cache_files_urlpath = settings.cache_files_urlpath
+    cache_files_urlpath = random.choice(settings.cache_files_urlpaths)
     fs_out, *_ = fsspec.get_fs_token_paths(
         cache_files_urlpath,
         storage_options=settings.cache_files_storage_options,

cacholote/utils.py

@@ -37,14 +37,17 @@ def hexdigestify(text: str) -> str:
     return hash_req.hexdigest()[:32]
 
 
-def get_cache_files_fs_dirname() -> tuple[fsspec.AbstractFileSystem, str]:
-    """Return the ``fsspec`` filesystem and directory name where cache files are stored."""
-    fs, _, (path,) = fsspec.get_fs_token_paths(
-        config.get().cache_files_urlpath,
-        storage_options=config.get().cache_files_storage_options,
-    )
-    fs.invalidate_cache()
-    return (fs, path)
+def get_cache_files_fs_dirnames() -> tuple[fsspec.AbstractFileSystem, list[str]]:
+    """Return the ``fsspec`` filesystem and directory names where cache files are stored."""
+    paths = []
+    for cache_files_urlpath in config.get().cache_files_urlpaths:
+        fs, _, (path,) = fsspec.get_fs_token_paths(
+            cache_files_urlpath,
+            storage_options=config.get().cache_files_storage_options,
+        )
+        paths.append(path)
+        fs.invalidate_cache()
+    return (fs, paths)
 
 
 def copy_buffered_file(

tests/conftest.py

@@ -55,14 +55,14 @@ def set_cache(
                 f"postgresql+psycopg2://{postgresql.info.user}:@{postgresql.info.host}:"
                 f"{postgresql.info.port}/{postgresql.info.dbname}"
             ),
-            cache_files_urlpath=f"s3://{test_bucket_name}",
+            cache_files_urlpaths=[f"s3://{test_bucket_name}"],
             cache_files_storage_options={"client_kwargs": client_kwargs},
         ):
             yield "cads"
     elif param.lower() in ("file", "local"):
         with config.set(
             cache_db_urlpath="sqlite:///" + str(tmp_path / "cacholote.db"),
-            cache_files_urlpath=str(tmp_path / "cache_files"),
+            cache_files_urlpaths=[str(tmp_path / "cache_files")],
         ):
             yield "file"
     elif param.lower() == "off":