162 lines
6.1 KiB
Python
162 lines
6.1 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from dataclasses import dataclass, field
|
|
from datetime import UTC, datetime
|
|
from typing import Any
|
|
|
|
from fastapi import WebSocket
|
|
|
|
from app.db.session import async_session_factory
|
|
from app.services.system_logs import (
|
|
DEFAULT_LOG_LINE_LIMIT,
|
|
LOG_SOURCES,
|
|
MAX_LOG_LINE_LIMIT,
|
|
read_database_log_events,
|
|
read_log_events,
|
|
)
|
|
|
|
DATABASE_LOG_SOURCE_IDS = {"system-db", "audit-db"}
|
|
LOG_TAIL_CHANNEL = "logs_tail"
|
|
LOG_TAIL_INTERVAL_SECONDS = 1.5
|
|
LOG_TAIL_SCAN_MULTIPLIER = 5
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class LogTailConfig:
|
|
source_id: str
|
|
limit: int = DEFAULT_LOG_LINE_LIMIT
|
|
level: str = "all"
|
|
levels: str | None = None
|
|
start_date: str | None = None
|
|
end_date: str | None = None
|
|
search: str | None = None
|
|
|
|
|
|
@dataclass
|
|
class LogTailSubscription:
|
|
config: LogTailConfig
|
|
emitted_cursors: set[str] = field(default_factory=set)
|
|
task: asyncio.Task | None = None
|
|
|
|
|
|
class LogTailManager:
|
|
def __init__(self) -> None:
|
|
self._subscriptions: dict[WebSocket, LogTailSubscription] = {}
|
|
|
|
def normalize_config(self, payload: dict[str, Any]) -> LogTailConfig:
|
|
source_id = str(payload.get("source_id") or payload.get("source") or "").strip()
|
|
if not source_id:
|
|
raise ValueError("source_id is required")
|
|
if source_id not in LOG_SOURCES and source_id not in DATABASE_LOG_SOURCE_IDS:
|
|
raise ValueError("Log source not found")
|
|
try:
|
|
limit = int(payload.get("limit") or DEFAULT_LOG_LINE_LIMIT)
|
|
except (TypeError, ValueError) as exc:
|
|
raise ValueError("limit must be a number") from exc
|
|
if limit < 1 or limit > MAX_LOG_LINE_LIMIT:
|
|
raise ValueError(f"limit must be between 1 and {MAX_LOG_LINE_LIMIT}")
|
|
return LogTailConfig(
|
|
source_id=source_id,
|
|
limit=limit,
|
|
level=str(payload.get("level") or "all"),
|
|
levels=str(payload.get("levels")).strip() if payload.get("levels") else None,
|
|
start_date=str(payload.get("start_date")).strip() if payload.get("start_date") else None,
|
|
end_date=str(payload.get("end_date")).strip() if payload.get("end_date") else None,
|
|
search=str(payload.get("search")).strip() if payload.get("search") else None,
|
|
)
|
|
|
|
async def subscribe(self, websocket: WebSocket, payload: dict[str, Any]) -> LogTailConfig:
|
|
config = self.normalize_config(payload)
|
|
await self.unsubscribe(websocket)
|
|
subscription = LogTailSubscription(config=config)
|
|
subscription.task = asyncio.create_task(self._run_tail(websocket, subscription))
|
|
self._subscriptions[websocket] = subscription
|
|
return config
|
|
|
|
async def unsubscribe(self, websocket: WebSocket) -> None:
|
|
subscription = self._subscriptions.pop(websocket, None)
|
|
if subscription and subscription.task:
|
|
subscription.task.cancel()
|
|
try:
|
|
await subscription.task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
|
|
async def disconnect(self, websocket: WebSocket) -> None:
|
|
await self.unsubscribe(websocket)
|
|
|
|
async def _run_tail(self, websocket: WebSocket, subscription: LogTailSubscription) -> None:
|
|
first_frame = True
|
|
while True:
|
|
events = await self._read_events(subscription.config)
|
|
if first_frame:
|
|
visible_events = events[-subscription.config.limit :]
|
|
subscription.emitted_cursors.update(event.cursor for event in visible_events)
|
|
await self._send_frame(websocket, subscription.config, "snapshot", visible_events)
|
|
first_frame = False
|
|
else:
|
|
new_events = [
|
|
event
|
|
for event in events
|
|
if event.cursor not in subscription.emitted_cursors
|
|
]
|
|
if new_events:
|
|
visible_events = new_events[-subscription.config.limit :]
|
|
subscription.emitted_cursors.update(event.cursor for event in visible_events)
|
|
await self._send_frame(websocket, subscription.config, "append", visible_events)
|
|
await asyncio.sleep(LOG_TAIL_INTERVAL_SECONDS)
|
|
|
|
async def _read_events(self, config: LogTailConfig):
|
|
scan_limit = max(config.limit * LOG_TAIL_SCAN_MULTIPLIER, config.limit)
|
|
if config.source_id in DATABASE_LOG_SOURCE_IDS:
|
|
async with async_session_factory() as db:
|
|
events = await read_database_log_events(
|
|
config.source_id,
|
|
scan_limit=scan_limit,
|
|
level=config.level,
|
|
levels=config.levels,
|
|
start_date=config.start_date,
|
|
end_date=config.end_date,
|
|
search=config.search,
|
|
db=db,
|
|
)
|
|
return events or []
|
|
events = read_log_events(
|
|
config.source_id,
|
|
scan_limit=scan_limit,
|
|
level=config.level,
|
|
levels=config.levels,
|
|
start_date=config.start_date,
|
|
end_date=config.end_date,
|
|
search=config.search,
|
|
)
|
|
return events or []
|
|
|
|
async def _send_frame(self, websocket: WebSocket, config: LogTailConfig, mode: str, events) -> None:
|
|
await websocket.send_json(
|
|
{
|
|
"type": "data_frame",
|
|
"channel": LOG_TAIL_CHANNEL,
|
|
"timestamp": datetime.now(UTC).isoformat(),
|
|
"payload": {
|
|
"mode": mode,
|
|
"source_id": config.source_id,
|
|
"line_count": len(events),
|
|
"lines": [event.line for event in events],
|
|
"filters": {
|
|
"limit": config.limit,
|
|
"level": config.level,
|
|
"levels": config.levels,
|
|
"start_date": config.start_date,
|
|
"end_date": config.end_date,
|
|
"search": config.search,
|
|
},
|
|
"status": "ok",
|
|
},
|
|
}
|
|
)
|
|
|
|
|
|
log_tail_manager = LogTailManager()
|