seitime-frappe/frappe/utils/caching.py

# Copyright (c) 2022, Frappe Technologies Pvt. Ltd. and Contributors
# License: MIT. Check LICENSE

import datetime
import json
from collections import defaultdict
from collections.abc import Callable
from functools import wraps

import pytz

import frappe

_SITE_CACHE = defaultdict(lambda: defaultdict(dict))


def __generate_request_cache_key(args: tuple, kwargs: dict):
	"""Generate a key for the cache."""
	if not kwargs:
		return hash(args)
	return hash((args, frozenset(kwargs.items())))


def request_cache(func: Callable) -> Callable:
	"""Decorator to cache function calls mid-request. Cache is stored in
	frappe.local.request_cache. The cache only persists for the current request
	and is cleared when the request is over. The function is called just once
	per request with the same set of (kw)arguments.

	Usage:
	        from frappe.utils.caching import request_cache

	        @request_cache
	        def calculate_pi(num_terms=0):
	                import math, time
	                print(f"{num_terms = }")
	                time.sleep(10)
	                return math.pi

	        calculate_pi(10) # will calculate value
	        calculate_pi(10) # will return value from cache
	"""

	@wraps(func)
	def wrapper(*args, **kwargs):
		if not getattr(frappe.local, "initialised", None):
			return func(*args, **kwargs)
		if not hasattr(frappe.local, "request_cache"):
			frappe.local.request_cache = defaultdict(dict)

		try:
			args_key = __generate_request_cache_key(args, kwargs)
		except Exception:
			return func(*args, **kwargs)

		try:
			return frappe.local.request_cache[func][args_key]
		except KeyError:
			return_val = func(*args, **kwargs)
			frappe.local.request_cache[func][args_key] = return_val
			return return_val

	return wrapper


def site_cache(ttl: int | None = None, maxsize: int | None = None) -> Callable:
	"""Decorator to cache method calls across requests. The cache is stored in
	frappe.utils.caching._SITE_CACHE. The cache persists on the parent process.
	It offers a light-weight cache for the current process without the additional
	overhead of serializing / deserializing Python objects.

	Note: This cache isn't shared among workers. If you need to share data across
	workers, use redis (frappe.cache API) instead.

	Usage:
	        from frappe.utils.caching import site_cache

	        @site_cache
	        def calculate_pi():
	                import math, time
	                precision = get_precision("Math Constant", "Pi") # depends on site data
	                return round(math.pi, precision)

	        calculate_pi(10) # will calculate value
	        calculate_pi(10) # will return value from cache
	        calculate_pi.clear_cache() # clear this function's cache for all sites
	        calculate_pi(10) # will calculate value
	"""

	def time_cache_wrapper(func: Callable | None = None) -> Callable:
		func_key = f"{func.__module__}.{func.__name__}"

		def clear_cache():
			"""Clear cache for this function for all sites if not specified."""
			_SITE_CACHE[func_key].clear()

		func.clear_cache = clear_cache

		if ttl is not None and not callable(ttl):
			func.ttl = ttl
			func.expiration = datetime.datetime.now(pytz.UTC) + datetime.timedelta(seconds=func.ttl)

		if maxsize is not None and not callable(maxsize):
			func.maxsize = maxsize

		@wraps(func)
		def site_cache_wrapper(*args, **kwargs):
			if getattr(frappe.local, "initialised", None):
				func_call_key = json.dumps((args, kwargs))

				if hasattr(func, "ttl") and datetime.datetime.now(pytz.UTC) >= func.expiration:
					func.clear_cache()
					func.expiration = datetime.datetime.now(pytz.UTC) + datetime.timedelta(seconds=func.ttl)

				if hasattr(func, "maxsize") and len(_SITE_CACHE[func_key][frappe.local.site]) >= func.maxsize:
					_SITE_CACHE[func_key][frappe.local.site].pop(
						next(iter(_SITE_CACHE[func_key][frappe.local.site])), None
					)

				if func_call_key not in _SITE_CACHE[func_key][frappe.local.site]:
					_SITE_CACHE[func_key][frappe.local.site][func_call_key] = func(*args, **kwargs)

				return _SITE_CACHE[func_key][frappe.local.site][func_call_key]

			return func(*args, **kwargs)

		return site_cache_wrapper

	if callable(ttl):
		return time_cache_wrapper(ttl)

	return time_cache_wrapper


def redis_cache(ttl: int | None = 3600, user: str | bool | None = None, shared: bool = False) -> Callable:
	"""Decorator to cache method calls and its return values in Redis

	args:
	        ttl: time to expiry in seconds, defaults to 1 hour
	        user: `true` should cache be specific to session user.
	        shared: `true` should cache be shared across sites
	"""

	def wrapper(func: Callable | None = None) -> Callable:
		func_key = f"{func.__module__}.{func.__qualname__}"

		def clear_cache():
			frappe.cache.delete_keys(func_key)

		func.clear_cache = clear_cache
		func.ttl = ttl if not callable(ttl) else 3600

		@wraps(func)
		def redis_cache_wrapper(*args, **kwargs):
			func_call_key = func_key + "::" + str(__generate_request_cache_key(args, kwargs))
			if frappe.cache.exists(func_call_key, user=user, shared=shared):
				return frappe.cache.get_value(func_call_key, user=user, shared=shared)
			val = func(*args, **kwargs)
			ttl = getattr(func, "ttl", 3600)
			frappe.cache.set_value(func_call_key, val, expires_in_sec=ttl, user=user, shared=shared)
			return val

		return redis_cache_wrapper

	if callable(ttl):
		return wrapper(ttl)
	return wrapper