Files
cosec_automation/cache/async_redis_cache_v2.py
T
khushalps a1c1853ca5 Squashed 'utils_v2/' content from commit 88e444a
git-subtree-dir: utils_v2
git-subtree-split: 88e444ac57ca6c090736645616d8a8a2d2af367b
2025-11-14 16:53:30 +05:30

569 lines
21 KiB
Python

"""
AUTHOR:
Khushal P Soonderji
DATE:
Saturday, 26th Oct., 2024
OBJECTIVE:
To provide an easy way to cache data for fast access. This version has the change that it can handle custom
serializers by way of dependency injection.
REFERENCES:
N/A
DOWNLOADS:
N/A
"""
# *****************************************************************************************************************
# ***** ****
# *** IMPORT ***
# ***** ****
# *****************************************************************************************************************
# To use redis:
import redis.asyncio as redis
# Other utils:
from utils_v2.string import json
from utils_v2.serialization.pickle_serializer import PickleSerializer
# For debugging:
from icecream import IceCreamDebugger
# To make a decorator:
from functools import wraps
# For hashing and shortening the hash:
import hashlib
import base64
# To work with datatypes:
from typing import List, Any
# *****************************************************************************************************************
# ***** ****
# *** MACROS / ONE-TIME INIT ***
# ***** ****
# *****************************************************************************************************************
# --- Nothing Yet
# *****************************************************************************************************************
# ***** ****
# *** VARIABLES ***
# ***** ****
# *****************************************************************************************************************
# --- Nothing Yet
# *****************************************************************************************************************
# ***** ****
# *** FUNCTIONS ***
# ***** ****
# *****************************************************************************************************************
# --- Nothing Yet
# *****************************************************************************************************************
# ***** ****
# *** WRAPPERS ***
# ***** ****
# *****************************************************************************************************************
def cache_it(cache = None, expiry = 120):
"""
This decorator factory takes an instance of the async caching class 'AsyncRedisCache' and holds your data there.
If a subsequent call is made to the same decorated function with the same inputs, the result is fetched from the
cache instead of going through the whole function again.
:param cache: The instance of 'AsyncRedisCache'.
:param expiry: The time in seconds after which the cached data must be cleared.
:return: The decorator that automatically caches your data.
"""
def decorator(func):
@wraps(func)
async def wrapper(*args, **kwargs):
# We first use the name of the function and the inputs given to it to generate a key for Redis with the
# simple hashing and shortening by way of base64 strings:
inputs_given = func.__name__ + str([_ for _ in args]) + str(kwargs)
sha256_hash = hashlib.sha256()
sha256_hash.update(inputs_given.encode("utf-8"))
hashed_key = sha256_hash.digest()
base64_key = base64.b64encode(hashed_key).decode("utf-8")
# Now we check if we have the value in cache:
try: response = await cache.get(base64_key, raise_exception = True)
# If the key doesn't exist, we pass through the function and store the results.
except:
response = await func(*args, **kwargs)
await cache.set(key = base64_key, value = response, expiry = expiry)
# Return the response from the wrapped function.
return response
return wrapper
return decorator
# ---------------------------------------------------------------------------------------------------------------------
def cache_class_methods(attr_name, expiry = 120):
"""
This decorator factory takes an instance of the async caching class 'AsyncRedisCache' and holds your data there.
If a subsequent call is made to the same decorated function with the same inputs, the result is fetched from the
cache instead of going through the whole function again.
:param attr_name: The name of the variable that has an instance of "AsyncRedisCache".
:param expiry: The time in seconds after which the cached data must be cleared.
:return: The decorator that automatically caches your data.
"""
def decorator(func):
@wraps(func)
async def wrapper(self, *args, **kwargs):
# Get the cache object first:
cache_obj = getattr(self, attr_name)
# We first use the name of the function and the inputs given to it to generate a key for Redis with the
# simple hashing and shortening by way of base64 strings:
inputs_given = func.__name__ + str([_ for _ in args]) + str(kwargs)
sha256_hash = hashlib.sha256()
sha256_hash.update(inputs_given.encode("utf-8"))
hashed_key = sha256_hash.digest()
base64_key = base64.b64encode(hashed_key).decode("utf-8")
# Now we check if we have the value in cache:
try: response = await cache_obj.get(base64_key, raise_exception = True)
# If the key doesn't exist, we pass through the function and store the results.
except:
response = await func(self, *args, **kwargs)
await cache_obj.set(key = base64_key, value = response, expiry = expiry)
# Return the response from the wrapped function.
return response
return wrapper
return decorator
# *****************************************************************************************************************
# ***** ****
# *** CLASSES ***
# ***** ****
# *****************************************************************************************************************
class AsyncRedisCache:
def __init__(
self,
connection_string,
serializer = None,
ping_counter = 1_000,
debug = False,
debug_prefix = "R-Cache | "
):
"""
Implements a simple cache in Redis which holds and returns all native datatypes like ints, floats, bools,
strings, dicts, lists, sets, and tuples :)
:param connection_string: The connection URL for connecting to Redis.
:param serializer: The serializer to use.
:param ping_counter: The number of requests to Redis after which you want to ping to ensure connection.
:param debug: Whether, or not, you want to show debugging messages from the start.
:param debug_prefix: The prefix text to show with the debugging messages.
"""
# Note down the configuration:
self.__client = None
self.__serializer = serializer or PickleSerializer()
self.__ping_counter = ping_counter
self.__requests_since_last_ping = 0
self.__connection_string = connection_string
# For debugging:
self.__printer = IceCreamDebugger(prefix = debug_prefix, includeContext = True)
if not debug: self.__printer.disable()
def enable_debug(self):
self.__printer.enable()
def disable_debug(self):
self.__printer.disable()
@staticmethod
def make_key(*args, **kwargs) -> str:
"""
Generates a key by hashing the args and kwargs sent to it. Can be useful to generate a predictable key. If the
inputs stay the same, the output stays the same.
:param args: Any number of args that you would like to use to generate the key.
:param kwargs: Any number of kwargs that you would like to use to generate the key.
:return: A string that can be used as a key to store values on Redis.
"""
# Take everything into one plain text string:
plain_text = "".join(str(a) for a in args)
plain_text += json.to_string(kwargs, no_space = True)
# Hash the plain text value, and create a key from it:
sha256_hash = hashlib.sha256()
sha256_hash.update(plain_text.encode("utf-8"))
hashed_key = sha256_hash.digest()
base64_key = base64.b64encode(hashed_key).decode("utf-8")
# Done here:
return base64_key
async def connect(self) -> bool:
"""
To make an asynchronous connection request to the Redis server to establish a connection.
:return: True if connected, else False.
"""
try:
self.__client = redis.from_url(
self.__connection_string
)
return True
except Exception as exception:
self.__printer(exception)
return False
async def disconnect(self) -> bool:
"""
Close the connection to the Redis server.
:return: True if successful, else False.
"""
if self.__client is not None:
try: await self.__client.close()
except Exception as exception:
self.__printer(exception)
return False
return True
return True
async def ensure_connection(self) -> bool:
"""
To ensure that we are connected. We keep pinging the Redis server every once in a while even when connected.
:return: True if connected, else False.
"""
# If we are not connected, we try to establish a connection:
if self.__client is None: return await self.connect()
# Else we check if we are connected. If not, we try to connect.
# But we check only once in a while. In the meantime, we assume that we are connected.
elif self.__requests_since_last_ping > self.__ping_counter:
try:
await self.__client.ping()
self.__requests_since_last_ping = 0
return True
except Exception as exception:
self.__printer(exception)
return await self.connect()
else:
self.__requests_since_last_ping += 1
return True
async def list_keys(
self,
match: str = "*",
raise_exception: bool = False
) -> List[str] | None:
"""
To get a list of all the keys that match the pattern.
:param match: To match a glob-style pattern. THIS IS NOT FULL-FLEDGED REGEX.
:param raise_exception: If you want to raise an exception if the process fails.
:return: The list of keys if successful, else None.
"""
# Standard connectivity check:
await self.ensure_connection()
# Start with an empty list:
keys_list = []
try:
# Start at the beginning,
# and break out of the loop if the pointer returns to zero:
pointer = 0
while True:
pointer, keys = await self.__client.scan(pointer, match = match)
keys_list = keys_list + [k.decode("utf-8") for k in keys]
if pointer == 0: break
# Done here:
return list(set(keys_list)) if keys_list else None
# If an exception occurs in th eprocess:
except Exception as exception:
self.__printer(exception)
if raise_exception: raise
else: return None
async def ttl(
self,
key: str | bytes,
raise_exception: bool = False
) -> int | float | None:
"""
To get the no. of seconds till the expiry of some key.
:param key: The key to check the expiry of.
:param raise_exception:
:return: -1 if the key is persistent (i.e., no expiry time set), -2 if the key does not exist, or the time left
in seconds if the key exists and is not persistent. If something goes wrong, you will get a null value.
"""
# Standard connectivity check:
await self.ensure_connection()
# Start by assuming failure:
ttl = None
try:
# Get the TTL of the key in ms,
# convert it to seconds and return the value:
ttl_ms = await self.__client.pttl(key)
if ttl_ms >= 0: ttl = ttl_ms / 1_000
return ttl
# If an exception occurs in th eprocess:
except Exception as exception:
self.__printer(exception)
if raise_exception: raise
else: return None
async def set(
self,
key: str | bytes,
value: Any,
expiry: float = None,
raise_exception: bool = False
) -> bool:
"""
Saves some value to the cache. If an expiry is specified, the data will be deleted after that many seconds.
:param key: The key with which the data will be stored and retrieved.
:param value: The value to store.
:param expiry: The time in seconds after which the data will expire. Must be a positive number.
:param raise_exception: If you want to raise an exception if the process fails.
:return: True if cached, else False.
"""
# Standard connectivity check:
await self.ensure_connection()
try:
# Here we actually try to set the data:
data = self.__serializer.serialize(value)
if expiry: response = await self.__client.setex(key, int(expiry), data)
else: response = await self.__client.set(key, data)
return response
# If an exception occurs in th eprocess:
except Exception as exception:
self.__printer(exception)
if raise_exception: raise
else: return False
async def get(
self,
key: str | bytes,
raise_exception: bool = False,
on_fail: Any = None
) -> Any:
"""
Retrieve the cached value.
:param key: The key with which the data was saved.
:param raise_exception: If you want to raise an exception if the process fails.
:param on_fail: What to return if the process fails due to an exception.
:return: The retrieved data or null if not found.
"""
# Standard connectivity check:
await self.ensure_connection()
try:
# Here we try to fetch the data:
data = await self.__client.get(key)
data = self.__serializer.deserialize(data)
return data
# If an exception occurs in th eprocess:
except Exception as exception:
self.__printer(exception)
if raise_exception: raise
else: return on_fail
async def delete(
self,
key: str | bytes,
raise_exception: bool = False
) -> bool:
"""
Prematurely delete the value from the cache before it expires.
:param key: The key with which the data was saved.
:param raise_exception: If you want to raise an exception if the process fails.
:return: True if deleted, else False.
"""
# Standard connectivity check:
await self.ensure_connection()
try:
# Try to manually delete the key before expiry:
response = await self.__client.delete(key)
return True if response else False
# If an exception occurs in th eprocess:
except Exception as exception:
self.__printer(exception)
if raise_exception: raise
else: return False
async def count(
self,
key: str | bytes,
value: int = 1,
expiry: float = None,
raise_exception: bool = False
) -> int | None:
"""
To use simple counters. If the counter (identified by the 'key') exists, it will be incremented, else the
counter will be created and the value will be incremented from 0.
:param key: The name of the counter.
:param value: The amount to increment the value by. Send negative values to count backwards.
:param expiry: The time (in seconds) in which the counter expires. Starts from the time the counter is created.
This value has to be an integer. If a float is passed, the value will be rounded off.
:param raise_exception: If you want to raise an exception if the process fails.
:return: The latest value of the counter. Will be null if something went wrong and the exception was suppressed.
"""
# Standard connectivity check:
await self.ensure_connection()
try:
# check if the key already exists,
# regardless of that, increment the counter:
already_existed = await self.__client.exists(key)
new_value = await self.__client.incrby(key, value)
# If the key didn't already exist, specify the expiry:
if expiry and not already_existed: await self.__client.expire(key, int(expiry))
# Done here:
return new_value
# If an exception occurs in th eprocess:
except Exception as exception:
self.__printer(exception)
if raise_exception: raise
else: return None
# *****************************************************************************************************************
# ***** ****
# *** MAIN PROGRAM ***
# ***** ****
# *****************************************************************************************************************
if __name__ == "__main__":
import asyncio
from utils_v2.string import json
async def main():
my_cache = AsyncRedisCache(
connection_string = str(input("Paste your conn. str.: ")),
ping_counter = 100,
debug = True
)
print("Connecting.")
connected = await my_cache.connect()
print("Success:", connected)
# Proceed only if we connected to the database successfully:
if connected:
existing_keys = await my_cache.list_keys()
print(f"KEYS ({len(existing_keys)}):", json.to_string(existing_keys, default = str))
key = str(input("Paste a key to check its TTL: "))
key_ttl = await my_cache.ttl(key = key)
print(f"KEY '{key}' has {key_ttl} seconds of TTL.")
value = await my_cache.get(key = str(input("Paste a key to get: ")))
print("GET:", json.to_string(value, default = str) if isinstance(value, (dict, list)) else value)
print("TYP:", type(value), end = "\n\n")
success = await my_cache.set(
key = "name",
value = {"first": "John", "last": "Doe"},
expiry = 10
)
print("SET:", success, end = "\n\n")
value = await my_cache.get(key = "name")
print("GET:", value)
print("TYP:", type(value), end = "\n\n")
success = await my_cache.delete(key = "name")
print("DEL:", success, end = "\n\n")
print("Starting count test.")
await my_cache.delete(key = "cnt")
for _ in range(5):
print("Step:", _)
await asyncio.sleep(1.0)
counter = await my_cache.count(key = "cnt", value = 1, expiry = 10)
print("COUNTER:", counter)
asyncio.run(main())