wastream-mirror/wastream/utils/http_cache.py

import hashlib
import json
from typing import Any, Optional, List

from fastapi import Request
from fastapi.responses import Response

from wastream.config.settings import settings


# ===========================
# No Cache Headers
# ===========================
NO_CACHE_HEADERS = {
    "Cache-Control": "no-store, no-cache, must-revalidate, max-age=0",
    "Pragma": "no-cache",
    "Expires": "0",
}


# ===========================
# Cache Control Builder
# ===========================
class CacheControl:

    def __init__(self):
        self._directives = []
        self._max_age = None
        self._s_maxage = None
        self._stale_while_revalidate = None
        self._stale_if_error = None

    def public(self):
        self._directives.append("public")
        return self

    def private(self):
        self._directives.append("private")
        return self

    def no_cache(self):
        self._directives.append("no-cache")
        return self

    def no_store(self):
        self._directives.append("no-store")
        return self

    def must_revalidate(self):
        self._directives.append("must-revalidate")
        return self

    def immutable(self):
        self._directives.append("immutable")
        return self

    def max_age(self, seconds: int):
        self._max_age = seconds
        return self

    def s_maxage(self, seconds: int):
        self._s_maxage = seconds
        return self

    def stale_while_revalidate(self, seconds: int):
        self._stale_while_revalidate = seconds
        return self

    def stale_if_error(self, seconds: int):
        self._stale_if_error = seconds
        return self

    def build(self) -> str:
        parts = list(self._directives)
        if self._max_age is not None:
            parts.append(f"max-age={self._max_age}")
        if self._s_maxage is not None:
            parts.append(f"s-maxage={self._s_maxage}")
        if self._stale_while_revalidate is not None:
            parts.append(f"stale-while-revalidate={self._stale_while_revalidate}")
        if self._stale_if_error is not None:
            parts.append(f"stale-if-error={self._stale_if_error}")
        return ", ".join(parts)


# ===========================
# ETag Generation
# ===========================
def generate_etag(data: Any) -> str:
    if isinstance(data, bytes):
        content = data
    elif isinstance(data, str):
        content = data.encode("utf-8")
    else:
        content = json.dumps(data, sort_keys=True, separators=(",", ":")).encode()
    return f'W/"{hashlib.md5(content, usedforsecurity=False).hexdigest()[:16]}"'


# ===========================
# ETag Matching
# ===========================
def check_etag_match(request: Request, etag: str) -> bool:
    if_none_match = request.headers.get("if-none-match")
    if not if_none_match:
        return False

    normalized_etag = etag.replace('W/"', '"')
    for client_etag in if_none_match.split(","):
        normalized_client = client_etag.strip().replace('W/"', '"')
        if normalized_client == normalized_etag or client_etag.strip() == "*":
            return True
    return False


# ===========================
# Cached JSON Response
# ===========================
class CachedJSONResponse(Response):

    def __init__(
        self,
        content: Any,
        status_code: int = 200,
        cache_control: Optional[CacheControl] = None,
        etag: Optional[str] = None,
        vary: Optional[List[str]] = None,
        **kwargs,
    ):
        body = json.dumps(content, separators=(",", ":")).encode()
        super().__init__(
            content=body,
            status_code=status_code,
            media_type="application/json",
            **kwargs,
        )

        if cache_control:
            self.headers["Cache-Control"] = cache_control.build()

        self.headers["ETag"] = etag or generate_etag(body)

        if vary:
            self.headers["Vary"] = ", ".join(vary)


# ===========================
# 304 Not Modified Response
# ===========================
def not_modified_response(etag: str, cache_control: str = "must-revalidate") -> Response:
    return Response(
        status_code=304,
        headers={
            "ETag": etag,
            "Cache-Control": cache_control,
        },
    )


# ===========================
# Cache Policies
# ===========================
class CachePolicies:

    @staticmethod
    def streams() -> CacheControl:
        ttl = settings.HTTP_CACHE_STREAMS_TTL
        swr = settings.HTTP_CACHE_STALE_WHILE_REVALIDATE
        return (
            CacheControl()
            .public()
            .max_age(ttl // 2)
            .s_maxage(ttl)
            .stale_while_revalidate(swr)
            .stale_if_error(300)
        )

    @staticmethod
    def manifest() -> CacheControl:
        ttl = settings.HTTP_CACHE_MANIFEST_TTL
        swr = settings.HTTP_CACHE_STALE_WHILE_REVALIDATE
        return CacheControl().public().max_age(ttl).stale_while_revalidate(swr)

    @staticmethod
    def configure_page() -> CacheControl:
        ttl = settings.HTTP_CACHE_CONFIGURE_TTL
        swr = settings.HTTP_CACHE_STALE_WHILE_REVALIDATE
        return CacheControl().public().max_age(ttl).stale_while_revalidate(swr)

    @staticmethod
    def empty_results() -> CacheControl:
        return (
            CacheControl()
            .public()
            .max_age(15)
            .s_maxage(30)
            .stale_if_error(60)
        )

    @staticmethod
    def no_cache() -> CacheControl:
        return CacheControl().private().no_store().no_cache().max_age(0)