HacKanCuBa · March 17, 2024 07:04 · Oct 11, 2023 · Oct 11, 2023 · Jun 22, 2023 · Jun 22, 2023
diff --git a/cache_helpers.py b/cache_helpers.py
@@ -1,6 +1,15 @@
 """Handy cache helpers.
 
 These are not yet production ready, as I haven't toroughly tested them, but close.
+
+---
+
+Mandatory license blah blah:
+
+Copyright (C) 2023 HacKan (https://hackan.net)
+This Source Code Form is subject to the terms of the Mozilla Public
+License, v. 2.0. If a copy of the MPL was not distributed with this
+file, You can obtain one at https://mozilla.org/MPL/2.0/.
 """
 
 from abc import ABC, abstractmethod

diff --git a/cache_helpers.py b/cache_helpers.py
@@ -1,3 +1,8 @@
+"""Handy cache helpers.
+
+These are not yet production ready, as I haven't toroughly tested them, but close.
+"""
+
 from abc import ABC, abstractmethod
 from contextlib import asynccontextmanager
 from typing import AsyncGenerator, Sequence

diff --git a/cache_helpers.py b/cache_helpers.py
@@ -21,6 +21,7 @@ async def delete(self, key: str | bytes) -> None:
     async def keys(self) -> tuple[bytes, ...]:
         ...
 
+
 class RedisAsyncCache(AsyncCacheBackend):
     def __init__(self, conn: redis.Redis):
         self._redis = conn

diff --git a/cache.py → cache_helpers.py b/cache.py → cache_helpers.py
diff --git a/cache.py b/cache.py
@@ -17,6 +17,9 @@ async def set(self, key: str | bytes, value: bytes, ttl: int | None = None) -> N
     async def delete(self, key: str | bytes) -> None:
         ...
 
+    @abstractmethod
+    async def keys(self) -> tuple[bytes, ...]:
+        ...
 
 class RedisAsyncCache(AsyncCacheBackend):
     def __init__(self, conn: redis.Redis):
@@ -31,7 +34,7 @@ async def set(self, key: str | bytes, value: bytes, ttl: int | None = None) -> N
     async def delete(self, key: str | bytes) -> None:
         await self._redis.delete(key)
 
-    async def keys(self) -> tuple[bytes]:
+    async def keys(self) -> tuple[bytes, ...]:
         keys: list[bytes] = await self._redis.keys()
 
         return tuple(keys)
@@ -60,7 +63,8 @@ async def delete(self, key: str | bytes) -> None:
         except KeyError:
             pass
 
-    async def keys(self) -> tuple[bytes]:
+    async def keys(self) -> tuple[bytes, ...]:
+        # noinspection PyTypeChecker
         return tuple(self._cache.keys())
 
 
@@ -81,17 +85,22 @@ async def set(self, key: str | bytes, value: bytes, ttl: int | None = None) -> N
     async def delete(self, key: str | bytes) -> None:
         await self._cache.delete(key)
 
+    async def keys(self) -> tuple[bytes, ...]:
+        return await self._cache.keys()
+
 
 class AsyncCachePassthrough:
     def __init__(self, backends: Sequence[AsyncCacheBackend], /) -> None:
         self._caches = tuple(backends)
 
     async def get(self, key: str | bytes) -> bytes:
         value = None
-        to_set = set()
+        to_set: set[AsyncCacheBackend] = set()
         for cache in self._caches:
             value = await cache.get(key)
             if value is not None:
+                # This assumes that every other cache after this one already has the value
+                # Not doing so would be quite slow
                 break
 
         if value is None:
@@ -110,7 +119,10 @@ async def delete(self, key: str | bytes) -> None:
         for cache in self._caches:
             await cache.delete(key)
 
-    async def warmup(self) -> None:
+    async def synchronize(self) -> None:
+        if len(self._caches) == 1:
+            return
+
         # This will assume that if two caches have the same key, they also have the same value
         keys: defaultdict[bytes, set[AsyncCacheBackend]] = defaultdict(set)
         for cache in self._caches:
@@ -124,6 +136,7 @@ async def warmup(self) -> None:
 
                 other = next(iter(keys[key]))
                 value = await other.get(key)
+                assert value is not None
                 await cache.set(key, value)
 
 
@@ -134,7 +147,7 @@ async def async_redis_connection(host: str, *, port: int = 6379, db: int = 0, **
     }
     params.update(kwargs)
 
-    conn = redis.Redis(host=host, port=port, db=db, **params)
+    conn: redis.Redis = redis.Redis(host=host, port=port, db=db, **params)  # type: ignore[arg-type]
 
     try:
         yield conn
@@ -155,3 +168,28 @@ async def async_cache(host: str | None = None, **kwargs: Any) -> AsyncGenerator[
         cache = AsyncCache(RedisAsyncCache(conn))
 
         yield cache
+
+
+@asynccontextmanager
+async def async_cache_passthrough(
+    host: str | None = None,
+    *,
+    use_in_memory: bool = True,
+    **kwargs: Any,
+) -> AsyncGenerator[AsyncCachePassthrough | None, None]:
+    if host:
+        async with async_redis_connection(host, **kwargs) as conn:
+            if use_in_memory:
+                backends = [InMemoryAsyncCache(), RedisAsyncCache(conn)]
+            else:
+                backends = [RedisAsyncCache(conn)]
+
+            cache = AsyncCachePassthrough(backends)
+
+    elif use_in_memory:
+        cache = AsyncCachePassthrough([InMemoryAsyncCache()])
+
+    else:
+        cache = None
+
+    yield cache
diff --git a/cache.py b/cache.py
@@ -22,19 +22,20 @@ class RedisAsyncCache(AsyncCacheBackend):
     def __init__(self, conn: redis.Redis):
         self._redis = conn
 
-    async def get(self, key: str | bytes) -> bytes:
-        data = await self._redis.get(key)
-        if data is None:
-            raise KeyError(key)
-
-        return data
+    async def get(self, key: str | bytes) -> bytes | None:
+        return await self._redis.get(key)
 
     async def set(self, key: str | bytes, value: bytes, ttl: int | None = None) -> None:
         await self._redis.set(key, value, ex=ttl)
 
     async def delete(self, key: str | bytes) -> None:
         await self._redis.delete(key)
 
+    async def keys(self) -> tuple[bytes]:
+        keys: list[bytes] = await self._redis.keys()
+
+        return tuple(keys)
+
 
 class InMemoryAsyncCache(AsyncCacheBackend):
     def __init__(self):
@@ -47,8 +48,8 @@ def _build_key(key: str | bytes) -> bytes:
 
         return key
 
-    async def get(self, key: str | bytes) -> bytes:
-        return self._cache[self._build_key(key)]
+    async def get(self, key: str | bytes) -> bytes | None:
+        return self._cache.get(self._build_key(key))
 
     async def set(self, key: str | bytes, value: bytes, _: int | None = None) -> None:
         self._cache[self._build_key(key)] = value
@@ -59,13 +60,20 @@ async def delete(self, key: str | bytes) -> None:
         except KeyError:
             pass
 
+    async def keys(self) -> tuple[bytes]:
+        return tuple(self._cache.keys())
+
 
 class AsyncCache:
     def __init__(self, backend: AsyncCacheBackend, /) -> None:
         self._cache = backend
 
     async def get(self, key: str | bytes) -> bytes:
-        return await self._cache.get(key)
+        value = await self._cache.get(key)
+        if value is None:
+            raise KeyError(key)
+
+        return value
 
     async def set(self, key: str | bytes, value: bytes, ttl: int | None = None) -> None:
         await self._cache.set(key, value, ttl)

diff --git a/cache.py b/cache.py
@@ -1,6 +1,6 @@
 from abc import ABC, abstractmethod
 from contextlib import asynccontextmanager
-from typing import AsyncGenerator
+from typing import AsyncGenerator, Sequence
 import redis.asyncio as redis
 
 
@@ -74,6 +74,51 @@ async def delete(self, key: str | bytes) -> None:
         await self._cache.delete(key)
 
 
+class AsyncCachePassthrough:
+    def __init__(self, backends: Sequence[AsyncCacheBackend], /) -> None:
+        self._caches = tuple(backends)
+
+    async def get(self, key: str | bytes) -> bytes:
+        value = None
+        to_set = set()
+        for cache in self._caches:
+            value = await cache.get(key)
+            if value is not None:
+                break
+
+        if value is None:
+            raise KeyError(key)
+
+        for cache in to_set:
+            await cache.set(key, value)
+
+        return value
+
+    async def set(self, key: str | bytes, value: bytes) -> None:
+        for cache in self._caches:
+            await cache.set(key, value)
+
+    async def delete(self, key: str | bytes) -> None:
+        for cache in self._caches:
+            await cache.delete(key)
+
+    async def warmup(self) -> None:
+        # This will assume that if two caches have the same key, they also have the same value
+        keys: defaultdict[bytes, set[AsyncCacheBackend]] = defaultdict(set)
+        for cache in self._caches:
+            for key in await cache.keys():
+                keys[key].add(cache)
+
+        for cache in self._caches:
+            for key in keys:
+                if cache in keys[key]:
+                    continue
+
+                other = next(iter(keys[key]))
+                value = await other.get(key)
+                await cache.set(key, value)
+
+
 @asynccontextmanager
 async def async_redis_connection(host: str, *, port: int = 6379, db: int = 0, **kwargs: Any) -> AsyncGenerator[redis.Redis, None]:
     params = {

diff --git a/cache.py b/cache.py
@@ -1,5 +1,6 @@
 from abc import ABC, abstractmethod
 from contextlib import asynccontextmanager
+from typing import AsyncGenerator
 import redis.asyncio as redis
 
 

diff --git a/cache.py b/cache.py
@@ -89,7 +89,7 @@ async def async_redis_connection(host: str, *, port: int = 6379, db: int = 0, **
 
 
 @asynccontextmanager
-async def async_cache(host: str | None = None, **kwargs: Any) -> AsyncCache:
+async def async_cache(host: str | None = None, **kwargs: Any) -> AsyncGenerator[AsyncCache, None]:
     if not host:
         cache = AsyncCache(InMemoryAsyncCache())
 

diff --git a/cache.py b/cache.py
@@ -85,4 +85,19 @@ async def async_redis_connection(host: str, *, port: int = 6379, db: int = 0, **
     try:
         yield conn
     finally:
-        await conn.close()
+        await conn.close()
+
+
+@asynccontextmanager
+async def async_cache(host: str | None = None, **kwargs: Any) -> AsyncCache:
+    if not host:
+        cache = AsyncCache(InMemoryAsyncCache())
+
+        yield cache
+
+        return
+
+    async with async_redis_connection(host, **kwargs) as conn:
+        cache = AsyncCache(RedisAsyncCache(conn))
+
+        yield cache
diff --git a/cache.py b/cache.py
@@ -0,0 +1,88 @@
+from abc import ABC, abstractmethod
+from contextlib import asynccontextmanager
+import redis.asyncio as redis
+
+
+class AsyncCacheBackend(ABC):
+    @abstractmethod
+    async def get(self, key: str | bytes) -> bytes:
+        ...
+
+    @abstractmethod
+    async def set(self, key: str | bytes, value: bytes, ttl: int | None = None) -> None:
+        ...
+
+    @abstractmethod
+    async def delete(self, key: str | bytes) -> None:
+        ...
+
+
+class RedisAsyncCache(AsyncCacheBackend):
+    def __init__(self, conn: redis.Redis):
+        self._redis = conn
+
+    async def get(self, key: str | bytes) -> bytes:
+        data = await self._redis.get(key)
+        if data is None:
+            raise KeyError(key)
+
+        return data
+
+    async def set(self, key: str | bytes, value: bytes, ttl: int | None = None) -> None:
+        await self._redis.set(key, value, ex=ttl)
+
+    async def delete(self, key: str | bytes) -> None:
+        await self._redis.delete(key)
+
+
+class InMemoryAsyncCache(AsyncCacheBackend):
+    def __init__(self):
+        self._cache = {}
+
+    @staticmethod
+    def _build_key(key: str | bytes) -> bytes:
+        if isinstance(key, str):
+            return key.encode()
+
+        return key
+
+    async def get(self, key: str | bytes) -> bytes:
+        return self._cache[self._build_key(key)]
+
+    async def set(self, key: str | bytes, value: bytes, _: int | None = None) -> None:
+        self._cache[self._build_key(key)] = value
+
+    async def delete(self, key: str | bytes) -> None:
+        try:
+            del self._cache[self._build_key(key)]
+        except KeyError:
+            pass
+
+
+class AsyncCache:
+    def __init__(self, backend: AsyncCacheBackend, /) -> None:
+        self._cache = backend
+
+    async def get(self, key: str | bytes) -> bytes:
+        return await self._cache.get(key)
+
+    async def set(self, key: str | bytes, value: bytes, ttl: int | None = None) -> None:
+        await self._cache.set(key, value, ttl)
+
+    async def delete(self, key: str | bytes) -> None:
+        await self._cache.delete(key)
+
+
+@asynccontextmanager
+async def async_redis_connection(host: str, *, port: int = 6379, db: int = 0, **kwargs: Any) -> AsyncGenerator[redis.Redis, None]:
+    params = {
+        "auto_close_connection_pool": True,
+    }
+    params.update(kwargs)
+
+    conn = redis.Redis(host=host, port=port, db=db, **params)
+
+    try:
+        yield conn
+    finally:
+        await conn.close()
No results found