302 lines
12 KiB
Python
302 lines
12 KiB
Python
"""Клиент ClickHouse: батч-запись событий без блокировки бота.
|
||
|
||
Свойства:
|
||
- запись неблокирующая: :meth:`AnalyticsClient.record` только кладёт строку в буфер;
|
||
- батчинг: отправка по размеру буфера и по таймеру;
|
||
- недоступность ClickHouse **не роняет бота**: события попадают в stdout (fallback);
|
||
- graceful-flush при остановке (:func:`close_analytics`).
|
||
|
||
Импорт ``clickhouse_connect`` ленивый (внутри :meth:`AnalyticsClient._ensure_client`),
|
||
поэтому пакет не является обязательным для импорта модуля (нужно для тестов/админки).
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import contextlib
|
||
import json
|
||
import logging
|
||
from collections import deque
|
||
from datetime import UTC, datetime
|
||
from typing import Any, Protocol
|
||
|
||
from common.analytics.events import EVENT_COLUMNS, EventType
|
||
from common.config import Settings, get_settings
|
||
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
|
||
|
||
|
||
class ClickHouseClient(Protocol):
|
||
"""Минимальный контракт клиента ClickHouse (используется и для подмены в тестах)."""
|
||
|
||
async def insert(self, table: str, data: list[list[Any]], column_names: list[str]) -> Any:
|
||
"""Вставить строки в таблицу."""
|
||
|
||
async def close(self) -> None:
|
||
"""Закрыть соединения."""
|
||
|
||
|
||
class AnalyticsClient:
|
||
"""Батч-запись событий в ClickHouse с деградацией в лог."""
|
||
|
||
def __init__(self, settings: Settings | None = None, *, client: ClickHouseClient | None = None) -> None:
|
||
self._settings = settings or get_settings()
|
||
self._client = client
|
||
self._rows: deque[list[Any]] = deque()
|
||
self._flush_now = asyncio.Event()
|
||
self._task: asyncio.Task[None] | None = None
|
||
self._usernames: dict[int, str] = {}
|
||
self._unavailable_logged = False
|
||
self._overflow = 0
|
||
self._table = f"{self._settings.clickhouse_db}.events"
|
||
|
||
@property
|
||
def enabled(self) -> bool:
|
||
"""Включена ли аналитика (``CLICKHOUSE_ENABLED``)."""
|
||
return self._settings.clickhouse_enabled
|
||
|
||
@property
|
||
def pending(self) -> int:
|
||
"""Сколько событий ждёт отправки (для тестов и диагностики)."""
|
||
return len(self._rows)
|
||
|
||
async def start(self) -> None:
|
||
"""Подключиться к ClickHouse (best-effort) и запустить фоновый сброс буфера."""
|
||
if not self.enabled:
|
||
logger.info("Аналитика ClickHouse выключена (CLICKHOUSE_ENABLED=false)")
|
||
return
|
||
|
||
await self._ensure_client()
|
||
self._task = asyncio.create_task(self._run(), name="analytics-flush")
|
||
logger.info(
|
||
"Аналитика ClickHouse: %s:%s/%s (батч %s, интервал %sс)",
|
||
self._settings.clickhouse_host,
|
||
self._settings.clickhouse_port,
|
||
self._settings.clickhouse_db,
|
||
self._settings.analytics_batch_size,
|
||
self._settings.analytics_flush_interval,
|
||
)
|
||
|
||
async def _ensure_client(self) -> bool:
|
||
"""Вернуть True, если клиент ClickHouse доступен."""
|
||
if self._client is not None:
|
||
return True
|
||
|
||
settings = self._settings
|
||
try:
|
||
import clickhouse_connect
|
||
|
||
self._client = await clickhouse_connect.get_async_client(
|
||
host=settings.clickhouse_host,
|
||
port=settings.clickhouse_port,
|
||
username=settings.clickhouse_user,
|
||
password=settings.clickhouse_password,
|
||
database=settings.clickhouse_db,
|
||
)
|
||
except Exception:
|
||
if not self._unavailable_logged:
|
||
self._unavailable_logged = True
|
||
logger.warning(
|
||
"ClickHouse %s:%s недоступен: события пишутся в stdout (fallback)",
|
||
settings.clickhouse_host,
|
||
settings.clickhouse_port,
|
||
exc_info=True,
|
||
)
|
||
return False
|
||
|
||
self._unavailable_logged = False
|
||
return True
|
||
|
||
def remember_user(self, user_id: int, username: str | None) -> None:
|
||
"""Запомнить связку ``id_user -> username`` для обогащения будущих событий."""
|
||
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).replace(tzinfo=None),
|
||
str(event_type),
|
||
int(user_id),
|
||
username or "",
|
||
int(chat_id),
|
||
self._dump_payload(payload),
|
||
__version__,
|
||
1 if self._settings.is_prod else 0,
|
||
]
|
||
)
|
||
|
||
while len(self._rows) > self._settings.analytics_buffer_max:
|
||
self._rows.popleft()
|
||
self._overflow += 1
|
||
|
||
if len(self._rows) >= self._settings.analytics_batch_size:
|
||
self._flush_now.set()
|
||
|
||
async def _run(self) -> None:
|
||
"""Периодически (или по достижении батча) сбрасывать буфер в ClickHouse."""
|
||
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 в лог (без исключений)."""
|
||
overflow, self._overflow = self._overflow, 0
|
||
if overflow:
|
||
logger.warning("Буфер событий переполнен, потеряно %s событий (не записаны в ClickHouse)", overflow)
|
||
|
||
if not self._rows:
|
||
return
|
||
|
||
if not self.enabled:
|
||
self._rows.clear()
|
||
return
|
||
|
||
rows = list(self._rows)
|
||
self._rows.clear()
|
||
|
||
if not await self._ensure_client():
|
||
self._log_fallback(rows, "ClickHouse недоступен")
|
||
return
|
||
|
||
client = self._client
|
||
if client is None:
|
||
self._log_fallback(rows, "клиент ClickHouse не создан")
|
||
return
|
||
|
||
try:
|
||
await client.insert(self._table, rows, list(EVENT_COLUMNS))
|
||
except Exception:
|
||
logger.exception("Не удалось записать %s событий в ClickHouse", len(rows))
|
||
self._log_fallback(rows, "ошибка записи в ClickHouse")
|
||
|
||
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
|
||
|
||
if self._client is not None:
|
||
await self.flush()
|
||
client, self._client = self._client, None
|
||
try:
|
||
await client.close()
|
||
except Exception:
|
||
logger.warning("Не удалось закрыть клиент ClickHouse", exc_info=True)
|
||
elif self._rows:
|
||
# Клиента нет (ClickHouse недоступен) — не пробуем подключиться на выходе.
|
||
rows = list(self._rows)
|
||
self._rows.clear()
|
||
self._log_fallback(rows, "остановка бота")
|
||
|
||
def _log_fallback(self, rows: list[list[Any]], reason: str) -> None:
|
||
"""Вывести потерянные события в лог, чтобы данные не пропадали молча."""
|
||
if not rows:
|
||
return
|
||
lines = " | ".join(
|
||
f"{row[0]:%Y-%m-%d %H:%M:%S} {row[1]} user={row[2]} chat={row[4]} {row[5]}"
|
||
for row in rows[:FALLBACK_SAMPLE]
|
||
)
|
||
logger.warning(
|
||
"Fallback (%s): %s событий не записаны в ClickHouse: %s%s",
|
||
reason,
|
||
len(rows),
|
||
lines,
|
||
" ..." if len(rows) > FALLBACK_SAMPLE else "",
|
||
)
|
||
|
||
@staticmethod
|
||
def _dump_payload(payload: dict[str, Any] | None) -> str:
|
||
"""Сериализовать payload в JSON-строку с обрезкой по длине."""
|
||
if not payload:
|
||
return "{}"
|
||
try:
|
||
text = json.dumps(payload, ensure_ascii=False, default=str)
|
||
except (TypeError, ValueError):
|
||
text = json.dumps({"_error": "payload не сериализуется"})
|
||
return text[:PAYLOAD_MAX_LENGTH]
|
||
|
||
|
||
#: Активный клиент (если аналитика не инициализирована — None).
|
||
_client: AnalyticsClient | None = None
|
||
|
||
|
||
async def init_analytics(
|
||
settings: Settings | None = None,
|
||
*,
|
||
client: ClickHouseClient | None = None,
|
||
) -> AnalyticsClient:
|
||
"""Создать и запустить единственный клиент аналитики. Идемпотентно."""
|
||
global _client
|
||
if _client is not None:
|
||
return _client
|
||
_client = AnalyticsClient(settings, client=client)
|
||
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:
|
||
"""Вернуть активный клиент аналитики (или 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:
|
||
"""Записать событие. Если аналитика не инициализирована — no-op."""
|
||
if _client is None:
|
||
return
|
||
_client.record(event_type, user_id=user_id, username=username, chat_id=chat_id, payload=payload)
|
||
|
||
|
||
def remember_user(user_id: int, username: str | None) -> None:
|
||
"""Запомнить username пользователя (для обогащения событий из сервисов)."""
|
||
if _client is not None:
|
||
_client.remember_user(user_id, username)
|