artemis/core/data/cache.py

68 lines
2.3 KiB
Python
Raw Normal View History

from typing import Any, Callable
from functools import wraps
import hashlib
import pickle
import logging
from core.config import CoreConfig
2023-03-09 16:38:58 +00:00
cfg: CoreConfig = None # type: ignore
# Make memcache optional
try:
import pylibmc # type: ignore
2023-03-09 16:38:58 +00:00
has_mc = True
except ModuleNotFoundError:
has_mc = False
2023-03-09 16:38:58 +00:00
def cached(lifetime: int = 10, extra_key: Any = None) -> Callable:
def _cached(func: Callable) -> Callable:
2023-10-06 02:16:50 +00:00
if has_mc and (cfg and cfg.database.enable_memcached):
hostname = "127.0.0.1"
if cfg:
hostname = cfg.database.memcached_host
memcache = pylibmc.Client([hostname], binary=True)
memcache.behaviors = {"tcp_nodelay": True, "ketama": True}
2023-03-09 16:38:58 +00:00
@wraps(func)
def wrapper(*args: Any, **kwargs: Any) -> Any:
if lifetime is not None:
# Hash function args
items = kwargs.items()
hashable_args = (args[1:], sorted(list(items)))
args_key = hashlib.md5(pickle.dumps(hashable_args)).hexdigest()
# Generate unique cache key
cache_key = f'{func.__module__}-{func.__name__}-{args_key}-{extra_key() if hasattr(extra_key, "__call__") else extra_key}'
# Return cached version if allowed and available
try:
result = memcache.get(cache_key)
except pylibmc.Error as e:
logging.getLogger("database").error(f"Memcache failed: {e}")
result = None
2023-03-09 16:38:58 +00:00
if result is not None:
logging.getLogger("database").debug(f"Cache hit: {result}")
return result
# Generate output
result = func(*args, **kwargs)
# Cache output if allowed
if lifetime is not None and result is not None:
logging.getLogger("database").debug(f"Setting cache: {result}")
memcache.set(cache_key, result, lifetime)
return result
2023-03-09 16:38:58 +00:00
else:
2023-03-09 16:38:58 +00:00
@wraps(func)
def wrapper(*args: Any, **kwargs: Any) -> Any:
return func(*args, **kwargs)
return wrapper
return _cached