98 lines
3.0 KiB
Python
98 lines
3.0 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import time
|
|
from dataclasses import dataclass
|
|
from typing import Awaitable, Callable, Hashable
|
|
|
|
from app.pixiv.oauth import PixivOauthToken
|
|
|
|
|
|
DEFAULT_REFRESH_MARGIN_S = 60.0
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class CachedAccessToken:
|
|
access_token: str
|
|
expires_at: float
|
|
|
|
def is_valid(self, *, now: float, refresh_margin_s: float) -> bool:
|
|
if not self.access_token:
|
|
return False
|
|
return now < (self.expires_at - refresh_margin_s)
|
|
|
|
|
|
class AccessTokenCache:
|
|
def __init__(
|
|
self,
|
|
*,
|
|
now: Callable[[], float] | None = None,
|
|
refresh_margin_s: float = DEFAULT_REFRESH_MARGIN_S,
|
|
) -> None:
|
|
self._now = now or time.time
|
|
self._refresh_margin_s = float(refresh_margin_s)
|
|
if self._refresh_margin_s < 0:
|
|
self._refresh_margin_s = 0.0
|
|
|
|
self._items: dict[Hashable, CachedAccessToken] = {}
|
|
self._locks: dict[Hashable, asyncio.Lock] = {}
|
|
|
|
def get(self, key: Hashable) -> str | None:
|
|
item = self._items.get(key)
|
|
if item is None:
|
|
return None
|
|
if not item.is_valid(now=float(self._now()), refresh_margin_s=self._refresh_margin_s):
|
|
return None
|
|
return item.access_token
|
|
|
|
def set(self, key: Hashable, *, access_token: str, expires_in_s: float) -> None:
|
|
access_token = (access_token or "").strip()
|
|
expires_in_s = float(expires_in_s)
|
|
if not access_token or expires_in_s <= 0:
|
|
self._items.pop(key, None)
|
|
return
|
|
expires_at = float(self._now()) + expires_in_s
|
|
self._items[key] = CachedAccessToken(access_token=access_token, expires_at=expires_at)
|
|
|
|
def invalidate(self, key: Hashable) -> None:
|
|
self._items.pop(key, None)
|
|
|
|
async def get_or_refresh(
|
|
self,
|
|
key: Hashable,
|
|
*,
|
|
refresher: Callable[[], Awaitable[PixivOauthToken]],
|
|
) -> PixivOauthToken:
|
|
cached = self.get(key)
|
|
if cached is not None:
|
|
return PixivOauthToken(
|
|
access_token=cached,
|
|
token_type="bearer",
|
|
expires_in=int(self._items[key].expires_at - float(self._now())),
|
|
refresh_token=None,
|
|
scope=None,
|
|
user_id=None,
|
|
)
|
|
|
|
lock = self._locks.get(key)
|
|
if lock is None:
|
|
lock = asyncio.Lock()
|
|
self._locks[key] = lock
|
|
|
|
async with lock:
|
|
cached2 = self.get(key)
|
|
if cached2 is not None:
|
|
return PixivOauthToken(
|
|
access_token=cached2,
|
|
token_type="bearer",
|
|
expires_in=int(self._items[key].expires_at - float(self._now())),
|
|
refresh_token=None,
|
|
scope=None,
|
|
user_id=None,
|
|
)
|
|
|
|
token = await refresher()
|
|
self.set(key, access_token=token.access_token, expires_in_s=float(token.expires_in))
|
|
return token
|
|
|