mirror of
https://github.com/long2ice/fastapi-cache.git
synced 2026-03-24 20:47:54 +00:00
130 lines
4.5 KiB
Python
130 lines
4.5 KiB
Python
import inspect
|
|
from functools import wraps
|
|
from typing import Callable, Optional, Type
|
|
|
|
from fastapi.concurrency import run_in_threadpool
|
|
from starlette.requests import Request
|
|
from starlette.responses import Response
|
|
|
|
from fastapi_cache import FastAPICache
|
|
from fastapi_cache.coder import Coder
|
|
|
|
|
|
def cache(
|
|
expire: int = None,
|
|
coder: Type[Coder] = None,
|
|
key_builder: Callable = None,
|
|
namespace: Optional[str] = "",
|
|
):
|
|
"""
|
|
cache all function
|
|
:param namespace:
|
|
:param expire:
|
|
:param coder:
|
|
:param key_builder:
|
|
|
|
:return:
|
|
"""
|
|
|
|
def wrapper(func):
|
|
signature = inspect.signature(func)
|
|
request_param = next(
|
|
(param for param in signature.parameters.values() if param.annotation is Request),
|
|
None,
|
|
)
|
|
response_param = next(
|
|
(param for param in signature.parameters.values() if param.annotation is Response),
|
|
None,
|
|
)
|
|
parameters = [*signature.parameters.values()]
|
|
if not request_param:
|
|
parameters.append(
|
|
inspect.Parameter(
|
|
name="request",
|
|
annotation=Request,
|
|
kind=inspect.Parameter.KEYWORD_ONLY,
|
|
),
|
|
)
|
|
if not response_param:
|
|
parameters.append(
|
|
inspect.Parameter(
|
|
name="response",
|
|
annotation=Response,
|
|
kind=inspect.Parameter.KEYWORD_ONLY,
|
|
),
|
|
)
|
|
if parameters:
|
|
signature = signature.replace(parameters=parameters)
|
|
func.__signature__ = signature
|
|
|
|
@wraps(func)
|
|
async def inner(*args, **kwargs):
|
|
nonlocal coder
|
|
nonlocal expire
|
|
nonlocal key_builder
|
|
|
|
async def ensure_async_func(*args, **kwargs):
|
|
"""Run cached sync functions in thread pool just like FastAPI."""
|
|
if inspect.iscoroutinefunction(func):
|
|
# async, return as is.
|
|
# unintuitively, we have to await once here, so that caller
|
|
# does not have to await twice. See
|
|
# https://stackoverflow.com/a/59268198/532513
|
|
return await func(*args, **kwargs)
|
|
else:
|
|
# sync, wrap in thread and return async
|
|
# see above why we have to await even although caller also awaits.
|
|
return await run_in_threadpool(func, *args, **kwargs)
|
|
|
|
|
|
copy_kwargs = kwargs.copy()
|
|
request = copy_kwargs.pop("request", None)
|
|
response = copy_kwargs.pop("response", None)
|
|
if (
|
|
request and request.headers.get("Cache-Control") == "no-store"
|
|
) or not FastAPICache.get_enable():
|
|
return await ensure_async_func(*args, **kwargs)
|
|
|
|
coder = coder or FastAPICache.get_coder()
|
|
expire = expire or FastAPICache.get_expire()
|
|
key_builder = key_builder or FastAPICache.get_key_builder()
|
|
backend = FastAPICache.get_backend()
|
|
|
|
cache_key = key_builder(
|
|
func, namespace, request=request, response=response, args=args, kwargs=copy_kwargs
|
|
)
|
|
ttl, ret = await backend.get_with_ttl(cache_key)
|
|
if not request:
|
|
if ret is not None:
|
|
return coder.decode(ret)
|
|
ret = await ensure_async_func(*args, **kwargs)
|
|
await backend.set(cache_key, coder.encode(ret), expire or FastAPICache.get_expire())
|
|
return ret
|
|
|
|
if request.method != "GET":
|
|
return await ensure_async_func(request, *args, **kwargs)
|
|
|
|
if_none_match = request.headers.get("if-none-match")
|
|
if ret is not None:
|
|
if response:
|
|
response.headers["Cache-Control"] = f"max-age={ttl}"
|
|
etag = f"W/{hash(ret)}"
|
|
if if_none_match == etag:
|
|
response.status_code = 304
|
|
return response
|
|
response.headers["ETag"] = etag
|
|
return coder.decode(ret)
|
|
if not request_param:
|
|
kwargs.pop("request")
|
|
if not response_param:
|
|
kwargs.pop("response")
|
|
|
|
ret = await ensure_async_func(*args, **kwargs)
|
|
|
|
await backend.set(cache_key, coder.encode(ret), expire or FastAPICache.get_expire())
|
|
return ret
|
|
|
|
return inner
|
|
|
|
return wrapper
|