feat: Restrict decorating async methods with sync enignes of Redis, Mongo and SQL cores

Author: gencurrent

src/cachier/core.py

@@ -156,6 +156,13 @@ def _pop_kwds_with_deprecation(kwds, name: str, default_value: bool):
     return kwds.pop(name, default_value)
 
 
+def _is_async_redis_client(client: Any) -> bool:
+    if client is None:
+        return False
+    method_names = ("hgetall", "hset", "keys", "delete", "hget")
+    return all(inspect.iscoroutinefunction(getattr(client, name, None)) for name in method_names)
+
+
 def cachier(
     hash_func: Optional[HashFunc] = None,
     hash_params: Optional[HashFunc] = None,
@@ -300,6 +307,42 @@ def cachier(
 
     def _cachier_decorator(func):
         core.set_func(func)
+        is_coroutine = inspect.iscoroutinefunction(func)
+
+        if backend == "mongo":
+            if is_coroutine and not inspect.iscoroutinefunction(mongetter):
+                msg = "Async cached functions with Mongo backend require an async mongetter."
+                raise TypeError(msg)
+            if (not is_coroutine) and inspect.iscoroutinefunction(mongetter):
+                msg = "Async mongetter requires an async cached function."
+                raise TypeError(msg)
+
+        if backend == "redis":
+            if is_coroutine:
+                if callable(redis_client):
+                    if not inspect.iscoroutinefunction(redis_client):
+                        msg = "Async cached functions with Redis backend require an async redis_client callable."
+                        raise TypeError(msg)
+                elif not _is_async_redis_client(redis_client):
+                    msg = "Async cached functions with Redis backend require an async Redis client."
+                    raise TypeError(msg)
+            else:
+                if callable(redis_client) and inspect.iscoroutinefunction(redis_client):
+                    msg = "Async redis_client callable requires an async cached function."
+                    raise TypeError(msg)
+                if _is_async_redis_client(redis_client):
+                    msg = "Async Redis client requires an async cached function."
+                    raise TypeError(msg)
+
+        if backend == "sql":
+            sql_core = core
+            assert isinstance(sql_core, _SQLCore)  # noqa: S101
+            if is_coroutine and not sql_core.has_async_engine():
+                msg = "Async cached functions with SQL backend require an AsyncEngine sql_engine."
+                raise TypeError(msg)
+            if (not is_coroutine) and sql_core.has_async_engine():
+                msg = "Async SQL engines require an async cached function."
+                raise TypeError(msg)
 
         last_cleanup = datetime.min
         cleanup_lock = threading.Lock()
@@ -501,8 +544,6 @@ async def _call_async(*args, max_age: Optional[timedelta] = None, **kwds):
         # argument.
         # For async functions, we create an async wrapper that calls
         # _call_async.
-        is_coroutine = inspect.iscoroutinefunction(func)
-
         if is_coroutine:
 
             @wraps(func)
@@ -522,6 +563,14 @@ def _clear_being_calculated():
             """Mark all entries in this cache as not being calculated."""
             core.clear_being_calculated()
 
+        async def _aclear_cache():
+            """Clear the cache asynchronously."""
+            await core.aclear_cache()
+
+        async def _aclear_being_calculated():
+            """Mark all entries in this cache as not being calculated asynchronously."""
+            await core.aclear_being_calculated()
+
         def _cache_dpath():
             """Return the path to the cache dir, if exists; None if not."""
             return getattr(core, "cache_dir", None)
@@ -541,6 +590,8 @@ def _precache_value(*args, value_to_cache, **kwds):  # noqa: D417
 
         func_wrapper.clear_cache = _clear_cache
         func_wrapper.clear_being_calculated = _clear_being_calculated
+        func_wrapper.aclear_cache = _aclear_cache
+        func_wrapper.aclear_being_calculated = _aclear_being_calculated
         func_wrapper.cache_dpath = _cache_dpath
         func_wrapper.precache_value = _precache_value
         return func_wrapper

src/cachier/cores/mongo.py

@@ -13,7 +13,6 @@
 import warnings  # to warn if pymongo is missing
 from contextlib import suppress
 from datetime import datetime, timedelta
-from inspect import isawaitable
 from typing import Any, Optional, Tuple
 
 from .._types import HashFunc, Mongetter
@@ -68,16 +67,7 @@ def _ensure_collection(self) -> Any:
 
         with self.lock:
             if self.mongo_collection is None:
-                coll = self.mongetter()
-                if isawaitable(coll):
-                    # Avoid "coroutine was never awaited" warnings.
-                    close = getattr(coll, "close", None)
-                    if callable(close):
-                        with suppress(Exception):
-                            close()
-                    msg = "async mongetter is only supported for async cached functions"
-                    raise TypeError(msg)
-                self.mongo_collection = coll
+                self.mongo_collection = self.mongetter()
 
             if not self._index_verified:
                 index_inf = self.mongo_collection.index_information()
@@ -96,23 +86,17 @@ async def _ensure_collection_async(self) -> Any:
         if self.mongo_collection is not None and self._index_verified:
             return self.mongo_collection
 
-        coll = self.mongetter()
-        if isawaitable(coll):
-            coll = await coll
+        coll = await self.mongetter()
         self.mongo_collection = coll
 
         if not self._index_verified:
-            index_inf = self.mongo_collection.index_information()
-            if isawaitable(index_inf):
-                index_inf = await index_inf
+            index_inf = await self.mongo_collection.index_information()
             if _MongoCore._INDEX_NAME not in index_inf:
                 func1key1 = IndexModel(
                     keys=[("func", ASCENDING), ("key", ASCENDING)],
                     name=_MongoCore._INDEX_NAME,
                 )
-                res = self.mongo_collection.create_indexes([func1key1])
-                if isawaitable(res):
-                    await res
+                await self.mongo_collection.create_indexes([func1key1])
             self._index_verified = True
 
         return self.mongo_collection
@@ -144,9 +128,7 @@ async def aget_entry(self, args, kwds) -> Tuple[str, Optional[CacheEntry]]:
 
     async def aget_entry_by_key(self, key: str) -> Tuple[str, Optional[CacheEntry]]:
         mongo_collection = await self._ensure_collection_async()
-        res = mongo_collection.find_one({"func": self._func_str, "key": key})
-        if isawaitable(res):
-            res = await res
+        res = await mongo_collection.find_one({"func": self._func_str, "key": key})
         if not res:
             return key, None
         val = None
@@ -188,7 +170,7 @@ async def aset_entry(self, key: str, func_res: Any) -> bool:
             return False
         mongo_collection = await self._ensure_collection_async()
         thebytes = pickle.dumps(func_res)
-        res = mongo_collection.update_one(
+        await mongo_collection.update_one(
             filter={"func": self._func_str, "key": key},
             update={
                 "$set": {
@@ -203,8 +185,6 @@ async def aset_entry(self, key: str, func_res: Any) -> bool:
             },
             upsert=True,
         )
-        if isawaitable(res):
-            await res
         return True
 
     def mark_entry_being_calculated(self, key: str) -> None:
@@ -217,13 +197,11 @@ def mark_entry_being_calculated(self, key: str) -> None:
 
     async def amark_entry_being_calculated(self, key: str) -> None:
         mongo_collection = await self._ensure_collection_async()
-        res = mongo_collection.update_one(
+        await mongo_collection.update_one(
             filter={"func": self._func_str, "key": key},
             update={"$set": {"processing": True}},
             upsert=True,
         )
-        if isawaitable(res):
-            await res
 
     def mark_entry_not_calculated(self, key: str) -> None:
         mongo_collection = self._ensure_collection()
@@ -240,13 +218,11 @@ def mark_entry_not_calculated(self, key: str) -> None:
     async def amark_entry_not_calculated(self, key: str) -> None:
         mongo_collection = await self._ensure_collection_async()
         with suppress(OperationFailure):
-            res = mongo_collection.update_one(
+            await mongo_collection.update_one(
                 filter={"func": self._func_str, "key": key},
                 update={"$set": {"processing": False}},
                 upsert=False,
             )
-            if isawaitable(res):
-                await res
 
     def wait_on_entry_calc(self, key: str) -> Any:
         time_spent = 0
@@ -266,9 +242,7 @@ def clear_cache(self) -> None:
 
     async def aclear_cache(self) -> None:
         mongo_collection = await self._ensure_collection_async()
-        res = mongo_collection.delete_many(filter={"func": self._func_str})
-        if isawaitable(res):
-            await res
+        await mongo_collection.delete_many(filter={"func": self._func_str})
 
     def clear_being_calculated(self) -> None:
         mongo_collection = self._ensure_collection()
@@ -279,12 +253,10 @@ def clear_being_calculated(self) -> None:
 
     async def aclear_being_calculated(self) -> None:
         mongo_collection = await self._ensure_collection_async()
-        res = mongo_collection.update_many(
+        await mongo_collection.update_many(
             filter={"func": self._func_str, "processing": True},
             update={"$set": {"processing": False}},
         )
-        if isawaitable(res):
-            await res
 
     def delete_stale_entries(self, stale_after: timedelta) -> None:
         """Delete stale entries from the MongoDB cache."""
@@ -296,6 +268,4 @@ async def adelete_stale_entries(self, stale_after: timedelta) -> None:
         """Delete stale entries from the MongoDB cache."""
         mongo_collection = await self._ensure_collection_async()
         threshold = datetime.now() - stale_after
-        res = mongo_collection.delete_many(filter={"func": self._func_str, "time": {"$lt": threshold}})
-        if isawaitable(res):
-            await res
+        await mongo_collection.delete_many(filter={"func": self._func_str, "time": {"$lt": threshold}})