71 lines
2.2 KiB
Python
71 lines
2.2 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from dataclasses import dataclass, field
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class RequestDraft:
|
|
category: str | None = None
|
|
description: str | None = None
|
|
attachments: list[dict] = field(default_factory=list)
|
|
state: str = "idle"
|
|
|
|
|
|
class RequestStateStore:
|
|
async def get(self, user_id: int) -> RequestDraft:
|
|
raise NotImplementedError
|
|
|
|
async def set(self, user_id: int, draft: RequestDraft) -> None:
|
|
raise NotImplementedError
|
|
|
|
async def clear(self, user_id: int) -> None:
|
|
raise NotImplementedError
|
|
|
|
|
|
class MemoryRequestStateStore(RequestStateStore):
|
|
def __init__(self) -> None:
|
|
self._storage: dict[int, RequestDraft] = {}
|
|
|
|
async def get(self, user_id: int) -> RequestDraft:
|
|
return self._storage.get(user_id, RequestDraft())
|
|
|
|
async def set(self, user_id: int, draft: RequestDraft) -> None:
|
|
self._storage[user_id] = draft
|
|
|
|
async def clear(self, user_id: int) -> None:
|
|
self._storage.pop(user_id, None)
|
|
|
|
|
|
class RedisRequestStateStore(RequestStateStore):
|
|
def __init__(self, redis_client, prefix: str = "maxbot:req:") -> None:
|
|
self._redis = redis_client
|
|
self._prefix = prefix
|
|
|
|
def _key(self, user_id: int) -> str:
|
|
return f"{self._prefix}{user_id}"
|
|
|
|
async def get(self, user_id: int) -> RequestDraft:
|
|
data = await self._redis.get(self._key(user_id))
|
|
if not data:
|
|
return RequestDraft()
|
|
payload = json.loads(data)
|
|
return RequestDraft(
|
|
category=payload.get("category"),
|
|
description=payload.get("description"),
|
|
attachments=payload.get("attachments") or [],
|
|
state=payload.get("state", "idle"),
|
|
)
|
|
|
|
async def set(self, user_id: int, draft: RequestDraft) -> None:
|
|
payload = {
|
|
"category": draft.category,
|
|
"description": draft.description,
|
|
"attachments": draft.attachments,
|
|
"state": draft.state,
|
|
}
|
|
await self._redis.set(self._key(user_id), json.dumps(payload, ensure_ascii=False), ex=60 * 60 * 6)
|
|
|
|
async def clear(self, user_id: int) -> None:
|
|
await self._redis.delete(self._key(user_id))
|
|
|