From 554b5224001d29c53eee5289380cad6101f70c9e Mon Sep 17 00:00:00 2001 From: profitroll Date: Sun, 16 Feb 2025 16:45:22 +0100 Subject: [PATCH] Added experimental cache support --- pyproject.toml | 1 + requirements/cache.txt | 2 + src/libbot/cache/__init__.py | 2 + src/libbot/cache/classes/__init__.py | 3 + src/libbot/cache/classes/cache.py | 44 ++++++++++ src/libbot/cache/classes/cache_memcached.py | 89 +++++++++++++++++++++ src/libbot/cache/classes/cache_redis.py | 89 +++++++++++++++++++++ src/libbot/cache/utils/__init__.py | 1 + src/libbot/cache/utils/_objects.py | 42 ++++++++++ src/libbot/cache/utils/manager.py | 24 ++++++ 10 files changed, 297 insertions(+) create mode 100644 requirements/cache.txt create mode 100644 src/libbot/cache/__init__.py create mode 100644 src/libbot/cache/classes/__init__.py create mode 100644 src/libbot/cache/classes/cache.py create mode 100644 src/libbot/cache/classes/cache_memcached.py create mode 100644 src/libbot/cache/classes/cache_redis.py create mode 100644 src/libbot/cache/utils/__init__.py create mode 100644 src/libbot/cache/utils/_objects.py create mode 100644 src/libbot/cache/utils/manager.py diff --git a/pyproject.toml b/pyproject.toml index 5748f8d..9fc3893 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,6 +31,7 @@ dev = { file = "requirements/dev.txt" } pycord = { file = "requirements/pycord.txt" } pyrogram = { file = "requirements/pyrogram.txt" } speed = { file = "requirements/speed.txt" } +cache = { file = "requirements/cache.txt" } [project.urls] Source = "https://git.end-play.xyz/profitroll/LibBotUniversal" diff --git a/requirements/cache.txt b/requirements/cache.txt new file mode 100644 index 0000000..f0f2e14 --- /dev/null +++ b/requirements/cache.txt @@ -0,0 +1,2 @@ +pymemcache~=4.0.0 +redis~=5.2.1 \ No newline at end of file diff --git a/src/libbot/cache/__init__.py b/src/libbot/cache/__init__.py new file mode 100644 index 0000000..def5795 --- /dev/null +++ b/src/libbot/cache/__init__.py @@ -0,0 +1,2 @@ +# This file is left empty on purpose +# Adding imports here will cause import errors when libbot[pycord] is not installed diff --git a/src/libbot/cache/classes/__init__.py b/src/libbot/cache/classes/__init__.py new file mode 100644 index 0000000..dce54a8 --- /dev/null +++ b/src/libbot/cache/classes/__init__.py @@ -0,0 +1,3 @@ +from .cache import Cache +from .cache_memcached import CacheMemcached +from .cache_redis import CacheRedis diff --git a/src/libbot/cache/classes/cache.py b/src/libbot/cache/classes/cache.py new file mode 100644 index 0000000..8b0f617 --- /dev/null +++ b/src/libbot/cache/classes/cache.py @@ -0,0 +1,44 @@ +from abc import ABC, abstractmethod +from typing import Any, Dict + +import pymemcache +import redis + + +class Cache(ABC): + client: pymemcache.Client | redis.Redis + + @classmethod + @abstractmethod + def from_config(cls, engine_config: Dict[str, Any]) -> Any: + pass + + @abstractmethod + def get_json(self, key: str) -> Any | None: + # TODO This method must also carry out ObjectId conversion! + pass + + @abstractmethod + def get_string(self, key: str) -> str | None: + pass + + @abstractmethod + def get_object(self, key: str) -> Any | None: + pass + + @abstractmethod + def set_json(self, key: str, value: Any) -> None: + # TODO This method must also carry out ObjectId conversion! + pass + + @abstractmethod + def set_string(self, key: str, value: str) -> None: + pass + + @abstractmethod + def set_object(self, key: str, value: Any) -> None: + pass + + @abstractmethod + def delete(self, key: str) -> None: + pass diff --git a/src/libbot/cache/classes/cache_memcached.py b/src/libbot/cache/classes/cache_memcached.py new file mode 100644 index 0000000..41b654c --- /dev/null +++ b/src/libbot/cache/classes/cache_memcached.py @@ -0,0 +1,89 @@ +import logging +from logging import Logger +from typing import Dict, Any + +from pymemcache import Client + +from .cache import Cache +from ..utils._objects import _json_to_string, _string_to_json + +logger: Logger = logging.getLogger(__name__) + + +class CacheMemcached(Cache): + client: Client + + def __init__(self, client: Client): + self.client = client + + logger.info("Initialized Memcached for caching") + + @classmethod + def from_config(cls, engine_config: Dict[str, Any]) -> "CacheMemcached": + if "uri" not in engine_config: + raise KeyError( + "Cache configuration is invalid. Please check if all keys are set (engine: memcached)" + ) + + return cls(Client(engine_config["uri"], default_noreply=True)) + + def get_json(self, key: str) -> Any | None: + try: + result: Any | None = self.client.get(key, None) + + logger.debug( + "Got json cache key '%s'%s", + key, + "" if result is not None else " (not found)", + ) + except Exception as exc: + logger.error("Could not get json cache key '%s' due to: %s", key, exc) + return None + + return None if result is None else _string_to_json(result) + + def get_string(self, key: str) -> str | None: + try: + result: str | None = self.client.get(key, None) + + logger.debug( + "Got string cache key '%s'%s", + key, + "" if result is not None else " (not found)", + ) + + return result + except Exception as exc: + logger.error("Could not get string cache key '%s' due to: %s", key, exc) + return None + + # TODO Implement binary deserialization + def get_object(self, key: str) -> Any | None: + raise NotImplementedError() + + def set_json(self, key: str, value: Any) -> None: + try: + self.client.set(key, _json_to_string(value)) + logger.debug("Set json cache key '%s'", key) + except Exception as exc: + logger.error("Could not set json cache key '%s' due to: %s", key, exc) + return None + + def set_string(self, key: str, value: str) -> None: + try: + self.client.set(key, value) + logger.debug("Set string cache key '%s'", key) + except Exception as exc: + logger.error("Could not set string cache key '%s' due to: %s", key, exc) + return None + + # TODO Implement binary serialization + def set_object(self, key: str, value: Any) -> None: + raise NotImplementedError() + + def delete(self, key: str) -> None: + try: + self.client.delete(key) + logger.debug("Deleted cache key '%s'", key) + except Exception as exc: + logger.error("Could not delete cache key '%s' due to: %s", key, exc) diff --git a/src/libbot/cache/classes/cache_redis.py b/src/libbot/cache/classes/cache_redis.py new file mode 100644 index 0000000..2edfa43 --- /dev/null +++ b/src/libbot/cache/classes/cache_redis.py @@ -0,0 +1,89 @@ +import logging +from logging import Logger +from typing import Dict, Any + +from redis import Redis + +from .cache import Cache +from ..utils._objects import _string_to_json, _json_to_string + +logger: Logger = logging.getLogger(__name__) + + +class CacheRedis(Cache): + client: Redis + + def __init__(self, client: Redis): + self.client = client + + logger.info("Initialized Redis for caching") + + @classmethod + def from_config(cls, engine_config: Dict[str, Any]) -> Any: + if "uri" not in engine_config: + raise KeyError( + "Cache configuration is invalid. Please check if all keys are set (engine: memcached)" + ) + + return cls(Redis.from_url(engine_config["uri"])) + + def get_json(self, key: str) -> Any | None: + try: + result: Any | None = self.client.get(key) + + logger.debug( + "Got json cache key '%s'%s", + key, + "" if result is not None else " (not found)", + ) + except Exception as exc: + logger.error("Could not get json cache key '%s' due to: %s", key, exc) + return None + + return None if result is None else _string_to_json(result) + + def get_string(self, key: str) -> str | None: + try: + result: str | None = self.client.get(key) + + logger.debug( + "Got string cache key '%s'%s", + key, + "" if result is not None else " (not found)", + ) + + return result + except Exception as exc: + logger.error("Could not get string cache key '%s' due to: %s", key, exc) + return None + + # TODO Implement binary deserialization + def get_object(self, key: str) -> Any | None: + raise NotImplementedError() + + def set_json(self, key: str, value: Any) -> None: + try: + self.client.set(key, _json_to_string(value)) + logger.debug("Set json cache key '%s'", key) + except Exception as exc: + logger.error("Could not set json cache key '%s' due to: %s", key, exc) + return None + + def set_string(self, key: str, value: str) -> None: + try: + self.client.set(key, value) + logger.debug("Set string cache key '%s'", key) + except Exception as exc: + logger.error("Could not set string cache key '%s' due to: %s", key, exc) + return None + + # TODO Implement binary serialization + def set_object(self, key: str, value: Any) -> None: + raise NotImplementedError() + + def delete(self, key: str) -> None: + try: + self.client.delete(key) + logger.debug("Deleted cache key '%s'", key) + except Exception as exc: + logger.error("Could not delete cache key '%s' due to: %s", key, exc) diff --git a/src/libbot/cache/utils/__init__.py b/src/libbot/cache/utils/__init__.py new file mode 100644 index 0000000..881a7f6 --- /dev/null +++ b/src/libbot/cache/utils/__init__.py @@ -0,0 +1 @@ +from .manager import create_cache_client diff --git a/src/libbot/cache/utils/_objects.py b/src/libbot/cache/utils/_objects.py new file mode 100644 index 0000000..9840c6d --- /dev/null +++ b/src/libbot/cache/utils/_objects.py @@ -0,0 +1,42 @@ +import logging +from copy import deepcopy +from logging import Logger +from typing import Any + +try: + from ujson import dumps, loads +except ImportError: + from json import dumps, loads + +logger: Logger = logging.getLogger(__name__) + +try: + from bson import ObjectId +except ImportError: + logger.warning( + "Could not import bson.ObjectId. PyMongo conversions will not be supported by the cache. It's safe to ignore this message if you do not use MongoDB." + ) + + +def _json_to_string(json_object: Any) -> str: + json_object_copy: Any = deepcopy(json_object) + + if isinstance(json_object_copy, dict) and "_id" in json_object_copy: + json_object_copy["_id"] = str(json_object_copy["_id"]) + + return dumps(json_object_copy, ensure_ascii=False, indent=0, escape_forward_slashes=False) + + +def _string_to_json(json_string: str) -> Any: + json_object: Any = loads(json_string) + + if "_id" in json_object: + try: + json_object["_id"] = ObjectId(json_object["_id"]) + except NameError: + logger.debug( + "Tried to convert attribute '_id' with value '%s' but bson.ObjectId is not present, skipping the conversion.", + json_object["_id"], + ) + + return json_object diff --git a/src/libbot/cache/utils/manager.py b/src/libbot/cache/utils/manager.py new file mode 100644 index 0000000..6b7db2c --- /dev/null +++ b/src/libbot/cache/utils/manager.py @@ -0,0 +1,24 @@ +from typing import Dict, Any, Literal + +from ..classes import CacheMemcached, CacheRedis + + +def create_cache_client( + config: Dict[str, Any], + engine: Literal["memcached", "redis"] | None = None, +) -> CacheMemcached | CacheRedis: + if engine not in ["memcached", "redis"] or engine is None: + raise KeyError(f"Incorrect cache engine provided. Expected 'memcached' or 'redis', got '{engine}'") + + if "cache" not in config or engine not in config["cache"]: + raise KeyError( + f"Cache configuration is invalid. Please check if all keys are set (engine: '{engine}')" + ) + + match engine: + case "memcached": + return CacheMemcached.from_config(config["cache"][engine]) + case "redis": + return CacheRedis.from_config(config["cache"][engine]) + case _: + raise KeyError(f"Cache implementation for the engine '{engine}' is not present.")