Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 35 additions & 1 deletion cloud_utils/cache/stores/redis.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
16 changes: 14 additions & 2 deletions cloud_utils/cache/stores/redis_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -23,11 +23,23 @@ 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)
try:
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"