Return cache entry

cacholote/cache.py

@@ -21,7 +21,7 @@
 import functools
 import json
 import warnings
-from typing import Any, Callable, Dict, Iterator, Optional, TypeVar, Union, cast
+from typing import Any, Callable, Iterator, Optional, TypeVar, Union, cast
 
 import sqlalchemy
 import sqlalchemy.orm
@@ -30,10 +30,6 @@
 
 F = TypeVar("F", bound=Callable[..., Any])
 
-LAST_PRIMARY_KEYS: contextvars.ContextVar[Dict[str, Any]] = contextvars.ContextVar(
-    "cacholote_last_primary_keys"
-)
-
 _LOCKER = "__locked__"
 
 
@@ -47,7 +43,9 @@ def _decode_and_update(
     if settings.tag is not None:
         cache_entry.tag = settings.tag
     database._commit_or_rollback(session)
-    LAST_PRIMARY_KEYS.set(cache_entry._primary_keys)
+    if settings.return_cache_entry:
+        session.refresh(cache_entry)
+        return cache_entry
     return result
 
 
@@ -116,15 +114,15 @@ def wrapper(
             for key, value in __context__.items():
                 key.set(value)
 
-        LAST_PRIMARY_KEYS.set({})
-
         settings = config.get()
         if not settings.use_cache:
             return func(*args, **kwargs)
 
         try:
             hexdigest = _hexdigestify_python_call(func, *args, **kwargs)
         except encode.EncodeError as ex:
+            if settings.return_cache_entry:
+                raise ex
             warnings.warn(f"can NOT encode python call: {ex!r}", UserWarning)
             return func(*args, **kwargs)
 
@@ -163,6 +161,8 @@ def wrapper(
                     cache_entry.result = json.loads(encode.dumps(result))
                     return _decode_and_update(session, cache_entry, settings)
                 except encode.EncodeError as ex:
+                    if settings.return_cache_entry:
+                        raise ex
                     warnings.warn(f"can NOT encode output: {ex!r}", UserWarning)
                     return result
 

cacholote/config.py

@@ -49,6 +49,7 @@ class Settings(pydantic.BaseSettings):
     raise_all_encoding_errors: bool = False
     expiration: Optional[str] = None
     tag: Optional[str] = None
+    return_cache_entry: bool = False
 
     @pydantic.validator("expiration")
     def expiration_must_be_isoformat(
@@ -64,6 +65,16 @@ def expiration_must_be_isoformat(
                 ) from ex
         return expiration
 
+    @pydantic.validator("return_cache_entry")
+    def validate_return_cache_entry(
+        cls: pydantic.BaseSettings, return_cache_entry: bool, values: Dict[str, Any]
+    ) -> bool:
+        if return_cache_entry is True and values["use_cache"] is False:
+            raise ValueError(
+                "`use_cache` must be True when `return_cache_entry` is True"
+            )
+        return return_cache_entry
+
     def make_cache_dir(self) -> None:
         fs, _, (urlpath, *_) = fsspec.get_fs_token_paths(
             self.cache_files_urlpath,
@@ -127,6 +138,8 @@ class set:
     tag: str, optional, default: None
         Tag for the cache entry. If None, do NOT tag.
         Note that existing tags are overwritten.
+    return_cache_entry: bool, default: False
+        Whether to return the cache database entry rather than decoded resuts.
     """
 
     def __init__(self, **kwargs: Any):

tests/test_30_cache.py

@@ -76,7 +76,6 @@ class Dummy:
         with pytest.warns(UserWarning, match="can NOT encode python call"):
             res = cfunc(inst)
         assert res == {"a": inst, "args": (), "b": None, "kwargs": {}}
-        assert cache.LAST_PRIMARY_KEYS.get() == {}
 
     if raise_all_encoding_errors:
         with pytest.raises(AttributeError):
@@ -85,7 +84,6 @@ class Dummy:
         with pytest.warns(UserWarning, match="can NOT encode output"):
             res = cfunc("test", b=1)
         assert res.__class__.__name__ == "LocalClass"
-        assert cache.LAST_PRIMARY_KEYS.get() == {}
 
     # cache-db must be empty
     con = database.ENGINE.get().raw_connection()
@@ -115,34 +113,26 @@ def test_use_cache(use_cache: bool) -> None:
 
     if use_cache:
         assert cached_now() == cached_now()
-        assert cache.LAST_PRIMARY_KEYS.get() == {
-            "key": "c3d9e414d0d32337c3672cb29b1b3cc9",
-            "expiration": datetime.datetime(9999, 12, 31, 23, 59, 59, 999999),
-        }
     else:
         assert cached_now() < cached_now()
-        assert cache.LAST_PRIMARY_KEYS.get() == {}
 
 
-def test_expiration() -> None:
-    first = cached_now()
-    assert cache.LAST_PRIMARY_KEYS.get() == {
-        "key": "c3d9e414d0d32337c3672cb29b1b3cc9",
-        "expiration": datetime.datetime(9999, 12, 31, 23, 59, 59, 999999),
-    }
+def test_expiration_and_return_cache_entry() -> None:
+    config.set(return_cache_entry=True)
+    first: database.CacheEntry = cached_now()  # type: ignore[assignment]
+    first.key = "c3d9e414d0d32337c3672cb29b1b3cc9"
+    first.expiration = datetime.datetime(9999, 12, 31, 23, 59, 59, 999999)
 
     with config.set(expiration="1908-03-09T00:00:00"):
-        assert cached_now() != first
-        assert cache.LAST_PRIMARY_KEYS.get() == {
-            "key": "c3d9e414d0d32337c3672cb29b1b3cc9",
-            "expiration": datetime.datetime(1908, 3, 9),
-        }
+        second: database.CacheEntry = cached_now()  # type: ignore[assignment]
+        assert second.result != first.result
+        assert second.key == "c3d9e414d0d32337c3672cb29b1b3cc9"
+        assert second.expiration == datetime.datetime(1908, 3, 9)
 
-    assert first == cached_now()
-    assert cache.LAST_PRIMARY_KEYS.get() == {
-        "key": "c3d9e414d0d32337c3672cb29b1b3cc9",
-        "expiration": datetime.datetime(9999, 12, 31, 23, 59, 59, 999999),
-    }
+    third: database.CacheEntry = cached_now()  # type: ignore[assignment]
+    assert first.result == third.result
+    assert third.key == "c3d9e414d0d32337c3672cb29b1b3cc9"
+    assert third.expiration == datetime.datetime(9999, 12, 31, 23, 59, 59, 999999)
 
 
 def test_tag(tmpdir: pathlib.Path) -> None:
@@ -171,19 +161,6 @@ def test_tag(tmpdir: pathlib.Path) -> None:
     assert cur.fetchall() == [("2", 4)]
 
 
-def test_contextvar() -> None:
-    cache.LAST_PRIMARY_KEYS.set({})
-
-    ctx = contextvars.copy_context()
-    ctx.run(cached_now)
-    assert ctx[cache.LAST_PRIMARY_KEYS] == {
-        "key": "c3d9e414d0d32337c3672cb29b1b3cc9",
-        "expiration": datetime.datetime(9999, 12, 31, 23, 59, 59, 999999),
-    }
-
-    assert cache.LAST_PRIMARY_KEYS.get() == {}
-
-
 def test_cached_error() -> None:
     con = database.ENGINE.get().raw_connection()
     cur = con.cursor()