rediskit 0.0.1__py3-none-any.whl

rediskit/__init__.py

@@ -0,0 +1,19 @@
+"""
+rediskit - Redis-backed performance and concurrency primitives for Python applications.
+
+Provides caching, distributed coordination, and data protection using Redis.
+"""
+
+from rediskit.encrypter import Encrypter
+from src.rediskit.memoize import RedisMemoize
+
+__all__ = [
+    "RedisMemoize",
+    "InitRedisConnectionPool",
+    "InitAsyncRedisConnectionPool",
+    "GetRedisConnection",
+    "GetAsyncRedisConnection",
+    "GetRedisMutexLock",
+    "GetAsyncRedisMutexLock",
+    "Encrypter",
+]

rediskit/config.py

@@ -0,0 +1,23 @@
+import os
+
+from DockerBuildSystem import TerminalTools
+
+from src.rediskit.utils import base64JsonToDict
+
+TerminalTools.LoadDefaultEnvironmentVariablesFile("private.env")
+TerminalTools.LoadDefaultEnvironmentVariablesFile(".env")
+
+# Redis Settings
+REDISKIT_REDIS_HOST = os.environ.get("REDISKIT_REDIS_HOST", "localhost")
+REDISKIT_REDIS_PORT = int(os.environ.get("REDISKIT_REDIS_PORT", "6379"))
+REDISKIT_REDIS_PASSWORD = os.environ.get("REDISKIT_REDIS_PASSWORD", "")
+REDISKIT_REDIS_TOP_NODE = os.environ.get("REDISKIT_REDIS_TOP_NODE", "redis_kit_node")
+REDISKIT_REDIS_SCAN_COUNT = int(os.environ.get("REDISKIT_REDIS_SCAN_COUNT", "10000"))
+REDISKIT_REDIS_SKIP_CACHING = os.environ.get("REDISKIT_REDIS_SKIP_CACHING", "false").upper() == "TRUE"
+
+# Lock Settings
+REDISKIT_LOCK_SETTINGS_REDIS_NAMESPACE = os.environ.get("REDISKIT_LOCK_SETTINGS_REDIS_NAMESPACE", f"{REDISKIT_REDIS_TOP_NODE}:LOCK")
+REDISKIT_LOCK_ASYNC_SETTINGS_REDIS_NAMESPACE = os.environ.get("REDISKIT_LOCK_ASYNC_SETTINGS_REDIS_NAMESPACE", f"{REDISKIT_REDIS_TOP_NODE}:LOCK_ASYNC")
+REDISKIT_LOCK_CACHE_REDIS_MUTEX = os.environ.get("REDISKIT_LOCK_CACHE_REDIS_MUTEX", "REDISKIT_LOCK_CACHE_REDIS_MUTEX")
+
+REDISKIT_ENCRYPTION_SECRET = base64JsonToDict(os.environ.get("REDISKIT_ENCRYPTION_SECRET", ""))

rediskit/encrypter.py

@@ -0,0 +1,168 @@
+import base64
+import gzip
+import json
+import re
+
+import zstd
+from nacl import encoding, secret
+from nacl.utils import random
+
+from rediskit import config
+
+
+class Encrypter:
+    VERSION_PREFIX = "__enc_v"
+
+    # EncryptedCompressedBase64Box
+    def __init__(self, keyHexDict: dict[str, str] = config.REDISKIT_ENCRYPTION_SECRET) -> None:
+        """
+        keysBase64 shall have the following format {"__enc_v1": "32-byte key"...,"__enc_vn": ...}
+        """
+        self.encryptionKeys = keyHexDict
+        self.latestVersion = list(self.encryptionKeys.keys())[-1]
+
+    def _getSecretBox(self, version: str) -> secret.SecretBox:
+        hexKey = self.encryptionKeys.get(version)
+        if not hexKey:
+            raise ValueError(f"Encryption key for version {version} not found")
+        # Ensure the key is in the proper format (32 bytes after decoding from hex)
+        return secret.SecretBox(hexKey.encode(), encoder=encoding.HexEncoder)
+
+    def encrypt[T: str | bytes | None](self, data: T, raiseIfEncrypted: bool = True, useZstd: bool = True) -> T:
+        if data is None:
+            return None  # type: ignore # not able to check this properly
+        elif isinstance(data, str):
+            dataToEncrypt: bytes = data.encode()
+            isText = True
+        elif isinstance(data, bytes):
+            dataToEncrypt = data
+            isText = False
+        else:
+            raise ValueError("data expected to be bytes or str")
+
+        if self.isEncrypted(dataToEncrypt, raiseIfEncrypted):
+            return data
+
+        compressedData: bytes = zstd.compress(dataToEncrypt) if useZstd else gzip.compress(dataToEncrypt)
+        tagBytes = b"zstd" if useZstd else b"gzip"
+
+        cipherText = self._getSecretBox(self.latestVersion).encrypt(compressedData, encoder=encoding.Base64Encoder)
+        token = self.latestVersion.encode() + b"|" + tagBytes + b":" + cipherText
+
+        return token.decode() if isText else token  # type: ignore # not able to check this properly
+
+    def decrypt[T: str | bytes | None](self, data: T) -> T:
+        if data is None:
+            return None  # type: ignore # not able to check this properly
+        elif isinstance(data, str):
+            dataToDecrypt: bytes = data.encode()
+            isText = True
+        elif isinstance(data, bytes):
+            dataToDecrypt = data
+            isText = False
+        else:
+            raise ValueError("data expected to be bytes or str")
+
+        try:
+            if dataToDecrypt.startswith(self.VERSION_PREFIX.encode()):
+                if b"|" in dataToDecrypt:
+                    # Format: __enc_vX|compression:ciphertext
+                    pre, rest = dataToDecrypt.split(b"|", 1)
+                    version = pre
+                    compressionTag, ciphertext = rest.split(b":", 1)
+                elif b":" in dataToDecrypt:
+                    # Format: __enc_vX:ciphertext
+                    version, ciphertext = dataToDecrypt.split(b":", 1)
+                    compressionTag = b"gzip"
+                else:
+                    raise ValueError("Invalid encrypted data format")
+            else:
+                # No version: legacy, fallback to v1 + gzip
+                version = f"{self.VERSION_PREFIX}1".encode()
+                compressionTag = b"gzip"
+                ciphertext = dataToDecrypt
+        except ValueError:
+            raise ValueError("Invalid encrypted data format; expected version prefix.")
+
+        secret_box = self._getSecretBox(version.decode())
+        compressed_data = secret_box.decrypt(ciphertext, encoder=encoding.Base64Encoder)
+
+        if compressionTag == b"zstd":
+            deCompressed = zstd.decompress(compressed_data)
+        elif compressionTag == b"gzip":
+            deCompressed = gzip.decompress(compressed_data)
+        else:
+            raise ValueError(f"Unknown compression '{compressionTag.decode()}' in encrypted data.")
+
+        return deCompressed.decode() if isText else deCompressed  # type: ignore # not able to check this properly
+
+    @staticmethod
+    def getEncryptionKeyVersionNumber(versionKey: str) -> int:
+        match = re.match(r"^__enc_v(\d+)$", versionKey)
+        if not match:
+            raise ValueError(f"Invalid version string: {versionKey}")
+        return int(match.group(1))
+
+    @staticmethod
+    def isEncrypted(data: str | bytes | None, raiseIfEncrypted: bool = False) -> bool:
+        # NB! can be false positive if a secret starts with "__enc_v"..., this risk is just accepted for now
+        if data is None:
+            return False
+        dataBytes = data.encode() if isinstance(data, str) else data
+        if not dataBytes.startswith(Encrypter.VERSION_PREFIX.encode()) or b":" not in dataBytes:
+            return False
+        if raiseIfEncrypted:
+            raise Exception("The data is already encrypted")
+        return True
+
+    @staticmethod
+    def generateNewHexKey() -> str:
+        """
+        Generate a new 32-byte key and return it as a hex-encoded string.
+
+        >>> key = Encrypter.generateNewHexKey()
+        >>> isinstance(key, str)
+        True
+        >>> len(key) == 64
+        True
+        """
+        newKey = random(secret.SecretBox.KEY_SIZE)  # Generates 32 random bytes.
+        hexKey = encoding.HexEncoder.encode(newKey).decode("utf-8")
+        return hexKey
+
+    @staticmethod
+    def encodeKeysDictToBase64(keys: dict) -> str:
+        """
+        Convert the keys dictionary to a JSON string and encode it in base64.
+
+        >>> Encrypter.encodeKeysDictToBase64({"v1": "abcdef"})
+        'eyJ2MSI6ICJhYmNkZWYifQ=='
+        """
+        json_str = json.dumps(keys)
+        base64_str = base64.b64encode(json_str.encode("utf-8")).decode("utf-8")
+        return base64_str
+
+    @staticmethod
+    def decodeKeysFromBase64(encodedKeys: str) -> dict:
+        """
+        Decode a base64-encoded JSON string into a dictionary of keys.
+
+        >>> Encrypter.decodeKeysFromBase64("eyJ2MSI6ICJhYmNkZWYifQ==")
+        {'v1': 'abcdef'}
+        """
+        try:
+            # Decode the base64 string into a JSON string.
+            json_str = base64.b64decode(encodedKeys).decode("utf-8")
+            # Convert the JSON string into a dictionary.
+            keys_dict = json.loads(json_str)
+            return keys_dict
+        except Exception as e:
+            raise ValueError("Failed to decode keys from base64 string.") from e
+
+    @staticmethod
+    def appendNewEncryptKey(currentKeys: str):
+        decodedkeys = Encrypter.decodeKeysFromBase64(currentKeys)
+        latestKey = list(decodedkeys.keys())[-1]
+        newVersion = f"{Encrypter.VERSION_PREFIX}{Encrypter.getEncryptionKeyVersionNumber(latestKey) + 1}"
+        decodedkeys.update({newVersion: Encrypter.generateNewHexKey()})
+        return {"newKeysDecoded": decodedkeys, "EncryptedKeys": Encrypter.encodeKeysDictToBase64(decodedkeys)}

rediskit/memoize.py

@@ -0,0 +1,258 @@
+import base64
+import functools
+import inspect
+import json
+import logging
+import pickle
+from typing import Any, Awaitable, Callable, Literal, overload
+
+import zstd
+from redis import Redis
+
+from rediskit import config, redisClient
+from rediskit.encrypter import Encrypter
+from rediskit.redisClient import HGetCacheFromRedis, HSetCacheToRedis
+from rediskit.redisLock import GetAsyncRedisMutexLock, GetRedisMutexLock
+
+log = logging.getLogger(__name__)
+CacheTypeOptions = Literal["zipPickled", "zipJson"]
+RedisStorageOptions = Literal["string", "hash"]
+
+
+def splitHashKey(key: str) -> tuple[str, str]:
+    *parts, field = key.split(":")
+    if not parts:
+        raise ValueError("Cannot use a single-part key with hash storage.")
+    return ":".join(parts), field
+
+
+def compressAndSign(data: Any, serializeFn: Callable[[Any], bytes], enableEncryption: bool = False) -> str:
+    serializedData = serializeFn(data)
+    if enableEncryption:
+        compressedData = Encrypter().encrypt(serializedData)
+    else:
+        compressedData = zstd.compress(serializedData)
+
+    return base64.b64encode(compressedData).decode("utf-8")
+
+
+def verifyAndDecompress(payload: bytes, deserializeFn: Callable[[bytes], Any], enableEncryption: bool = False) -> Any:
+    if enableEncryption:
+        serializedData = Encrypter().decrypt(payload)
+    else:
+        serializedData = zstd.decompress(payload)
+    return deserializeFn(serializedData)
+
+
+def deserializeData(
+    data: Any,
+    cacheType: CacheTypeOptions,
+    enableEncryption: bool = False,
+) -> bytes:
+    if cacheType == "zipPickled":
+        cachedData = verifyAndDecompress(base64.b64decode(data), lambda b: pickle.loads(b), enableEncryption)
+    elif cacheType == "zipJson":
+        cachedData = verifyAndDecompress(base64.b64decode(data), lambda b: json.loads(b.decode("utf-8")), enableEncryption)
+    else:
+        raise ValueError("Unknown cacheType specified.")
+
+    return cachedData
+
+
+def serializeData(
+    data: Any,
+    cacheType: CacheTypeOptions,
+    enableEncryption: bool = False,
+) -> str:
+    if cacheType == "zipPickled":
+        payload = compressAndSign(data, lambda d: pickle.dumps(d), enableEncryption)
+    elif cacheType == "zipJson":
+        payload = compressAndSign(data, lambda d: json.dumps(d).encode("utf-8"), enableEncryption)
+    else:
+        raise ValueError("Unknown cacheType specified.")
+    return payload
+
+
+def computeValue[T](param: T | Callable[..., T], *args, **kwargs) -> T:
+    if callable(param):
+        sig = inspect.signature(param)
+        accepts_kwargs = any(p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values())
+
+        if accepts_kwargs:
+            # Pass all kwargs directly
+            value = param(*args, **kwargs)
+        else:
+            # Filter only matching kwargs
+            filtered_kwargs = {k: v for k, v in kwargs.items() if k in sig.parameters}
+            bound = sig.bind(*args, **filtered_kwargs)
+            bound.apply_defaults()
+            value = param(*bound.args, **bound.kwargs)
+        return value
+    else:
+        return param
+
+
+def maybeDataInCache(
+    tenantId: str | None,
+    computedMemoizeKey: str,
+    computedTtl: int | None,
+    cacheType: CacheTypeOptions,
+    resetTtlUponRead: bool,
+    byPassCachedData: bool,
+    enableEncryption: bool,
+    storageType: RedisStorageOptions = "string",
+    connection: Redis | None = None,
+) -> Any:
+    if byPassCachedData:
+        log.info(f"Cache bypassed for tenantId: {tenantId}, key {computedMemoizeKey}")
+        return None
+
+    cachedData = None
+    if storageType == "string":
+        cached = redisClient.LoadBlobFromRedis(
+            tenantId, match=computedMemoizeKey, setTtlOnRead=computedTtl if resetTtlUponRead and computedTtl is not None else None, connection=connection
+        )
+        if cached:
+            log.info(f"Cache hit tenantId: {tenantId}, key: {computedMemoizeKey}")
+            cachedData = cached
+    elif storageType == "hash":
+        hashKey, field = splitHashKey(computedMemoizeKey)
+        cachedDict = HGetCacheFromRedis(
+            tenantId, hashKey, field, setTtlOnRead=computedTtl if resetTtlUponRead and computedTtl is not None else None, connection=connection
+        )
+        if cachedDict and field in cachedDict and cachedDict[field] is not None:
+            log.info(f"HASH cache hit tenantId: {tenantId}, key: {hashKey}, field: {field}")
+            cachedData = cachedDict[field]
+    else:
+        raise ValueError(f"Unknown storageType: {storageType}")
+
+    if cachedData:
+        return deserializeData(cachedData, cacheType, enableEncryption)
+    else:
+        log.info(f"No cache found tenantId: {tenantId}, key: {computedMemoizeKey}")
+        return None
+
+
+def dumpData(
+    data: Any,
+    tenantId: str,
+    computedMemoizeKey: str,
+    cacheType: CacheTypeOptions,
+    computedTtl: int | None,
+    enableEncryption: bool,
+    storageType: RedisStorageOptions = "string",
+    connection: Redis | None = None,
+) -> None:
+    payload = serializeData(data, cacheType, enableEncryption)
+    if storageType == "string":
+        redisClient.DumpBlobToRedis(tenantId, computedMemoizeKey, payload=payload, ttl=computedTtl, connection=connection)
+    elif storageType == "hash":
+        hashKey, field = splitHashKey(computedMemoizeKey)
+        HSetCacheToRedis(tenantId, hashKey, fields={field: payload}, ttl=computedTtl, connection=connection)
+    else:
+        raise ValueError(f"Unknown storageType: {storageType}")
+
+
+def RedisMemoize[T](
+    memoizeKey: Callable[..., str] | str,
+    ttl: Callable[..., int] | int | None = None,
+    bypassCache: Callable[..., bool] | bool = False,
+    cacheType: CacheTypeOptions = "zipJson",
+    resetTtlUponRead: bool = True,
+    enableEncryption: bool = False,
+    storageType: RedisStorageOptions = "string",
+    connection: Redis | None = None,
+) -> Callable[[Callable[..., T]], Callable[..., T]]:
+    """Caches the result of any function in Redis using either pickle or JSON.
+
+    The decorated function must have 'tenantId' as an arg or kwarg.
+
+    Params:
+    -------
+    - memoizeKey: Callable computing a memoize key based on wrapped funcs args and kwargs, callable shall define the logic to compute the correct memoize key.
+    - ttl: Time To Live, either fixed value, or callable consuming args+kwargs to return a ttl. Default None, if None no ttl is set.
+    - bypassCache: Don't get data from cache, run wrapped func and update cache. run new values.
+    - cacheType: "zipPickled" Uses pickle for arbitrary Python objects, "zipJson" Uses JSON for data that is JSON serializable.
+    - resetTtlUponRead: Set the ttl to the initial value upon reading the value from redis cache
+    - connection: Custom Redis connection to use instead of the default connection pool
+    """
+
+    def computeMemoizeKey(*args, **kwargs) -> str:
+        if not (isinstance(memoizeKey, str) or callable(memoizeKey)):
+            raise ValueError(f"Expected memoizeKey to be Callable or a str. got {type(memoizeKey)}")
+        return computeValue(memoizeKey, *args, **kwargs)
+
+    def computeTtl(*args, **kwargs) -> int | None:
+        if ttl is None:
+            return None
+        if not (isinstance(ttl, int) or callable(ttl)):
+            raise ValueError(f"Expected ttl to be Callable or an int. got {type(ttl)}")
+        return computeValue(ttl, *args, **kwargs)
+
+    def computeByPassCache(*args, **kwargs) -> bool:
+        if not (isinstance(bypassCache, bool) or callable(bypassCache)):
+            raise ValueError(f"Expected bypassCache to be Callable or an int. got {type(bypassCache)}")
+        return computeValue(bypassCache, *args, **kwargs)
+
+    def computeTenantId(wrappedFunc: Callable[..., Any], *args, **kwargs) -> str | None:
+        boundArgs = inspect.signature(wrappedFunc).bind(*args, **kwargs)
+        boundArgs.apply_defaults()
+        tenantId = boundArgs.arguments.get("tenantId") or boundArgs.kwargs.get("tenantId")
+        # if tenantId is None:
+        #    raise ValueError("tenantId not provided in either args or kwargs")
+        return tenantId
+
+    def getLockName(tenantId: str, computedMemoizeKey: str) -> str:
+        lockName = f"{config.REDISKIT_LOCK_CACHE_REDIS_MUTEX}:{tenantId}:{computedMemoizeKey}"
+        return lockName
+
+    def getParams(func, *args, **kwargs) -> tuple[str, int | None, str, str, bool]:
+        computedMemoizeKey = computeMemoizeKey(*args, **kwargs)
+        computedTtl = computeTtl(*args, **kwargs)
+        tenantId = computeTenantId(func, *args, **kwargs)
+        if tenantId is None:
+            raise ValueError("tenantId cannot be None")
+        lockName = getLockName(tenantId, computedMemoizeKey)
+        byPassCachedData = computeByPassCache(*args, **kwargs)
+
+        return computedMemoizeKey, computedTtl, tenantId, lockName, byPassCachedData
+
+    def decorator(func: Callable[..., T]) -> Callable[..., T]:
+        isAsyncFunc = inspect.iscoroutinefunction(func)
+        if isAsyncFunc:
+
+            @functools.wraps(func)
+            async def async_wrapper(*args, **kwargs) -> T:
+                computedMemoizeKey, computedTtl, tenantId, lockName, byPassCachedData = getParams(func, *args, **kwargs)
+                async with await GetAsyncRedisMutexLock(lockName, expire=60):
+                    inCache = maybeDataInCache(
+                        tenantId, computedMemoizeKey, computedTtl, cacheType, resetTtlUponRead, byPassCachedData, enableEncryption, storageType, connection
+                    )
+                    if inCache is not None:
+                        return inCache
+                    result = await func(*args, **kwargs)
+                    if result is not None:
+                        dumpData(result, tenantId, computedMemoizeKey, cacheType, computedTtl, enableEncryption, storageType, connection)
+                    return result
+
+            return async_wrapper
+
+        else:
+
+            @functools.wraps(func)
+            def wrapper(*args, **kwargs) -> T:
+                computedMemoizeKey, computedTtl, tenantId, lockName, byPassCachedData = getParams(func, *args, **kwargs)
+                with GetRedisMutexLock(lockName, auto_renewal=True, expire=60):
+                    inCache = maybeDataInCache(
+                        tenantId, computedMemoizeKey, computedTtl, cacheType, resetTtlUponRead, byPassCachedData, enableEncryption, storageType, connection
+                    )
+                    if inCache is not None:
+                        return inCache
+                    result = func(*args, **kwargs)
+                    if result is not None:
+                        dumpData(result, tenantId, computedMemoizeKey, cacheType, computedTtl, enableEncryption, storageType, connection)
+                    return result
+
+            return wrapper
+
+    return decorator