290 lines
9.9 KiB
Python
290 lines
9.9 KiB
Python
"""Аналитика в Postgres: батч-запись событий и дневных счётчиков сообщений.
|
||
|
||
Свойства:
|
||
- запись неблокирующая: ``record``/``record_message`` только кладут данные в буфер;
|
||
- батчинг по размеру/таймеру;
|
||
- сбой Postgres не роняет бота: батч уходит в лог (fallback) и отбрасывается;
|
||
- сообщения не пишутся поштучно: инкрементится ``analytics.messages_daily``.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import contextlib
|
||
import json
|
||
import logging
|
||
from collections import deque
|
||
from collections.abc import Awaitable, Callable
|
||
from datetime import UTC, datetime
|
||
from typing import Any
|
||
|
||
from psycopg.types.json import Jsonb
|
||
|
||
from common.analytics.events import EventType
|
||
from common.config import Settings, get_settings
|
||
from common.db.pool import get_pool
|
||
from common.version import __version__
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
#: Ограничение длины JSON-payload (защита от «толстых» событий).
|
||
PAYLOAD_MAX_LENGTH = 4000
|
||
|
||
#: Ограничение кэша ``id_user -> username``.
|
||
USERNAME_CACHE_MAX = 5000
|
||
|
||
#: Сколько строк события выводить в лог при fallback.
|
||
FALLBACK_SAMPLE = 5
|
||
|
||
#: Строка события: (time, type, user_id, username, chat_id, payload, version, prod).
|
||
EventRow = list[Any]
|
||
#: (day, user_id, chat_id) -> (username, count)
|
||
MessageCounts = dict[tuple[Any, int, int], list[Any]]
|
||
|
||
#: Куда сбрасывать батч (подменяется в тестах).
|
||
Writer = Callable[[list[EventRow], MessageCounts], Awaitable[None]]
|
||
|
||
|
||
async def _write_db(events: list[EventRow], messages: MessageCounts) -> None:
|
||
"""Записать батч в Postgres."""
|
||
async with get_pool().connection() as conn, conn.cursor() as cur:
|
||
if events:
|
||
await cur.executemany(
|
||
"""
|
||
INSERT INTO analytics.events
|
||
(event_time, event_type, user_id, username, chat_id, payload, bot_version, prod)
|
||
VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
|
||
""",
|
||
[
|
||
[
|
||
row[0],
|
||
row[1],
|
||
row[2],
|
||
row[3],
|
||
row[4],
|
||
Jsonb(row[5]),
|
||
row[6],
|
||
row[7],
|
||
]
|
||
for row in events
|
||
],
|
||
)
|
||
|
||
for (day, user_id, chat_id), (username, count) in messages.items():
|
||
await cur.execute(
|
||
"""
|
||
INSERT INTO analytics.messages_daily (day, user_id, username, chat_id, count)
|
||
VALUES (%s, %s, %s, %s, %s)
|
||
ON CONFLICT (day, user_id, chat_id) DO UPDATE
|
||
SET count = analytics.messages_daily.count + EXCLUDED.count,
|
||
username = EXCLUDED.username
|
||
""",
|
||
(day, user_id, username, chat_id, count),
|
||
)
|
||
|
||
|
||
class AnalyticsClient:
|
||
"""Батч-запись событий и счётчиков в Postgres с деградацией в лог."""
|
||
|
||
def __init__(
|
||
self,
|
||
settings: Settings | None = None,
|
||
*,
|
||
writer: Writer | None = None,
|
||
) -> None:
|
||
self._settings = settings or get_settings()
|
||
self._writer = writer or _write_db
|
||
self._rows: deque[EventRow] = deque()
|
||
self._messages: MessageCounts = {}
|
||
self._flush_now = asyncio.Event()
|
||
self._task: asyncio.Task[None] | None = None
|
||
self._usernames: dict[int, str] = {}
|
||
|
||
@property
|
||
def enabled(self) -> bool:
|
||
return self._settings.analytics_enabled
|
||
|
||
@property
|
||
def pending(self) -> int:
|
||
return len(self._rows) + sum(int(value[1]) for value in self._messages.values())
|
||
|
||
async def start(self) -> None:
|
||
if not self.enabled:
|
||
logger.info("Аналитика выключена (ANALYTICS_ENABLED=false)")
|
||
return
|
||
self._task = asyncio.create_task(self._run(), name="analytics-flush")
|
||
logger.info(
|
||
"Аналитика Postgres: батч %s, интервал %sс",
|
||
self._settings.analytics_batch_size,
|
||
self._settings.analytics_flush_interval,
|
||
)
|
||
|
||
def remember_user(self, user_id: int, username: str | None) -> None:
|
||
if not user_id or not username:
|
||
return
|
||
if user_id not in self._usernames and len(self._usernames) >= USERNAME_CACHE_MAX:
|
||
self._usernames.pop(next(iter(self._usernames)))
|
||
self._usernames[user_id] = username
|
||
|
||
def record(
|
||
self,
|
||
event_type: EventType | str,
|
||
*,
|
||
user_id: int = 0,
|
||
username: str | None = None,
|
||
chat_id: int = 0,
|
||
payload: dict[str, Any] | None = None,
|
||
) -> None:
|
||
"""Положить событие в буфер (никогда не бросает исключений)."""
|
||
if not self.enabled:
|
||
return
|
||
|
||
if username is None:
|
||
username = self._usernames.get(user_id, "")
|
||
|
||
self._rows.append(
|
||
[
|
||
datetime.now(UTC),
|
||
str(event_type),
|
||
int(user_id),
|
||
username or "",
|
||
int(chat_id),
|
||
self._payload(payload),
|
||
__version__,
|
||
bool(self._settings.is_prod),
|
||
]
|
||
)
|
||
|
||
while len(self._rows) > self._settings.analytics_buffer_max:
|
||
self._rows.popleft()
|
||
|
||
if self.pending >= self._settings.analytics_batch_size:
|
||
self._flush_now.set()
|
||
|
||
def record_message(self, *, user_id: int, username: str | None, chat_id: int) -> None:
|
||
"""Увеличить дневной счётчик сообщений пользователя."""
|
||
if not self.enabled or not user_id:
|
||
return
|
||
|
||
day = datetime.now(self._settings.tzinfo).date()
|
||
key = (day, int(user_id), int(chat_id))
|
||
entry = self._messages.get(key)
|
||
if entry is None:
|
||
self._messages[key] = [username or self._usernames.get(user_id, ""), 1]
|
||
else:
|
||
entry[0] = username or entry[0]
|
||
entry[1] = int(entry[1]) + 1
|
||
|
||
if self.pending >= self._settings.analytics_batch_size:
|
||
self._flush_now.set()
|
||
|
||
async def _run(self) -> None:
|
||
interval = self._settings.analytics_flush_interval
|
||
while True:
|
||
with contextlib.suppress(TimeoutError):
|
||
await asyncio.wait_for(self._flush_now.wait(), timeout=interval)
|
||
self._flush_now.clear()
|
||
await self.flush()
|
||
|
||
async def flush(self) -> None:
|
||
"""Отправить накопленное. При сбое — fallback в лог (без исключений)."""
|
||
if not self._rows and not self._messages:
|
||
return
|
||
|
||
events = list(self._rows)
|
||
messages = self._messages
|
||
self._rows.clear()
|
||
self._messages = {}
|
||
|
||
if not self.enabled:
|
||
return
|
||
|
||
try:
|
||
await self._writer(events, messages)
|
||
except Exception:
|
||
self._log_fallback(events, messages)
|
||
|
||
async def close(self) -> None:
|
||
if self._task is not None:
|
||
self._task.cancel()
|
||
with contextlib.suppress(asyncio.CancelledError):
|
||
await self._task
|
||
self._task = None
|
||
await self.flush()
|
||
|
||
def _log_fallback(self, events: list[EventRow], messages: MessageCounts) -> None:
|
||
sample = " | ".join(
|
||
f"{row[0]:%Y-%m-%d %H:%M:%S} {row[1]} user={row[2]} chat={row[4]}" for row in events[:FALLBACK_SAMPLE]
|
||
)
|
||
logger.warning(
|
||
"Fallback: не записано в Postgres событий=%s, счётчиков=%s%s",
|
||
len(events),
|
||
len(messages),
|
||
f": {sample}" if sample else "",
|
||
)
|
||
|
||
@staticmethod
|
||
def _payload(payload: dict[str, Any] | None) -> dict[str, Any]:
|
||
if not payload:
|
||
return {}
|
||
try:
|
||
text = json.dumps(payload, ensure_ascii=False, default=str)
|
||
except (TypeError, ValueError):
|
||
return {"_error": "payload не сериализуется"}
|
||
if len(text) > PAYLOAD_MAX_LENGTH:
|
||
return {"_truncated": text[:PAYLOAD_MAX_LENGTH]}
|
||
# Пересобираем из JSON: так в jsonb не попадёт объект, который не сериализуется.
|
||
return json.loads(text)
|
||
|
||
|
||
_client: AnalyticsClient | None = None
|
||
|
||
|
||
async def init_analytics(
|
||
settings: Settings | None = None,
|
||
*,
|
||
writer: Writer | None = None,
|
||
) -> AnalyticsClient:
|
||
"""Создать и запустить единственный клиент аналитики. Идемпотентно."""
|
||
global _client
|
||
if _client is not None:
|
||
return _client
|
||
_client = AnalyticsClient(settings, writer=writer)
|
||
await _client.start()
|
||
return _client
|
||
|
||
|
||
async def close_analytics() -> None:
|
||
global _client
|
||
if _client is not None:
|
||
await _client.close()
|
||
_client = None
|
||
|
||
|
||
def get_analytics() -> AnalyticsClient | None:
|
||
return _client
|
||
|
||
|
||
def record(
|
||
event_type: EventType | str,
|
||
*,
|
||
user_id: int = 0,
|
||
username: str | None = None,
|
||
chat_id: int = 0,
|
||
payload: dict[str, Any] | None = None,
|
||
) -> None:
|
||
if _client is None:
|
||
return
|
||
_client.record(event_type, user_id=user_id, username=username, chat_id=chat_id, payload=payload)
|
||
|
||
|
||
def record_message(*, user_id: int, username: str | None, chat_id: int) -> None:
|
||
if _client is None:
|
||
return
|
||
_client.record_message(user_id=user_id, username=username, chat_id=chat_id)
|
||
|
||
|
||
def remember_user(user_id: int, username: str | None) -> None:
|
||
if _client is not None:
|
||
_client.remember_user(user_id, username)
|