From d167a711c249e3baf7e0005b4d0a089d5721c87c Mon Sep 17 00:00:00 2001
From: Yoni Isaacs <yoni@hyro.ai>
Date: Wed, 10 May 2023 15:16:12 +0300
Subject: [PATCH] add store for bytes for redis

---
 cloud_utils/cache/stores/redis.py      | 36 +++++++++++++++++++++++++-
 cloud_utils/cache/stores/redis_test.py | 16 ++++++++++--
 2 files changed, 49 insertions(+), 3 deletions(-)

diff --git a/cloud_utils/cache/stores/redis.py b/cloud_utils/cache/stores/redis.py
index 0cdaf26..eaf24b9 100644
--- a/cloud_utils/cache/stores/redis.py
+++ b/cloud_utils/cache/stores/redis.py
@@ -21,7 +21,7 @@ async def wrapper(*args, **kwargs):
     return wrapper
 
 
-def make_store(
+def make_json_store(
     redis_client: redis.Redis,
     ttl: int,
     name: str,
@@ -53,3 +53,37 @@ async def set_item(key: str, value):
             )
 
     return get_item, set_item
+
+
+def make_bytes_store(
+    redis_client: redis.Redis,
+    ttl: int,
+    name: str,
+) -> Tuple[Callable, Callable]:
+    utils.log_initialized_cache("redis", name)
+
+    async def get_item(key: str):
+        cache_key = utils.cache_key_name(name, key)
+        result = await redis_error_handler(redis_client.get)(cache_key)
+        if result is None:
+            logging.error(f"{key} is not in {name}")
+            raise KeyError
+        try:
+            return result
+        except ValueError:  # Key contents are malformed (will force key to update).
+            logging.error(f"Malformed key detected: {key} in {name}.")
+            raise KeyError
+
+    async def set_item(key: str, value):
+        if ttl == 0:
+            await redis_error_handler(
+                redis_client.set,
+            )(utils.cache_key_name(name, key), value)
+        else:
+            await redis_error_handler(redis_client.setex)(
+                utils.cache_key_name(name, key),
+                ttl,
+                json.dumps(value),
+            )
+
+    return get_item, set_item
diff --git a/cloud_utils/cache/stores/redis_test.py b/cloud_utils/cache/stores/redis_test.py
index 41b6ef8..4155411 100644
--- a/cloud_utils/cache/stores/redis_test.py
+++ b/cloud_utils/cache/stores/redis_test.py
@@ -10,7 +10,7 @@
 
 async def test_redis_store_unbounded():
     client = aioredis.FakeRedis(server=_SERVER)
-    get_item, set_item = redis.make_store(client, 0, "unbound_store")
+    get_item, set_item = redis.make_json_store(client, 0, "unbound_store")
 
     await set_item("1", 1)
     await set_item("2", 2)
@@ -23,7 +23,7 @@ async def test_redis_store_unbounded():
 
 async def test_redis_store_ttl():
     client = aioredis.FakeRedis(server=_SERVER)
-    get_item, set_item = redis.make_store(client, 1, "ttl_1")
+    get_item, set_item = redis.make_json_store(client, 1, "ttl_1")
 
     await set_item("1", 1)
     await asyncio.sleep(2)
@@ -31,3 +31,15 @@ async def test_redis_store_ttl():
         await get_item("1")
     except KeyError:
         assert utils.cache_key_name("ttl_1", "1") not in await client.keys()
+
+
+async def test_make_bytes_store():
+    client = aioredis.FakeRedis(server=_SERVER)
+    get_item, set_item = redis.make_bytes_store(client, 0, "unbound_bytes_store")
+    await set_item("1", 1)
+    await set_item("2", 2)
+    await set_item("3", "3")
+
+    assert await get_item("1") == b"1"
+    assert await get_item("2") == b"2"
+    assert await get_item("3") == b"3"