GX-gp-notify/vendor/loghive_client/async_client.py

"""Async LogHive client — for use in asyncio-based projects (e.g., FastAPI, aiohttp)."""

import asyncio
import logging
import traceback
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional
from urllib.parse import urljoin

import httpx

logger = logging.getLogger(__name__)


class AsyncLogHiveLogger:
    """Async logger for asyncio applications.

    Uses an async background task to batch and send log entries.
    Ideal for FastAPI / Starlette / aiohttp projects.

    Usage:
        logger = AsyncLogHiveLogger("my-project", "api-key", "http://localhost:8000")
        await logger.start()

        await logger.info("Request processed", extra={"path": "/api/users"})

        await logger.stop()
    """

    def __init__(
        self,
        project: str,
        api_key: str,
        endpoint: str = "http://localhost:8000",
        batch_size: int = 50,
        flush_interval: float = 2.0,
        max_retries: int = 3,
        timeout: float = 5.0,
    ):
        self.project = project
        self.api_key = api_key
        self.endpoint = endpoint.rstrip("/")
        self.batch_size = batch_size
        self.flush_interval = flush_interval
        self.max_retries = max_retries
        self.timeout = timeout

        self._queue: asyncio.Queue = asyncio.Queue()
        self._task: Optional[asyncio.Task] = None
        self._client: Optional[httpx.AsyncClient] = None
        self._stop_event = asyncio.Event()

    async def start(self):
        """Start the background flush task."""
        if self._task and not self._task.done():
            return
        self._client = httpx.AsyncClient(timeout=self.timeout)
        self._stop_event.clear()
        self._task = asyncio.create_task(self._flush_loop())
        logger.debug("AsyncLogHiveLogger started for project '%s'", self.project)

    async def stop(self, flush: bool = True):
        """Stop the background task."""
        self._stop_event.set()
        if flush:
            await self._flush_now()
        if self._task:
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                pass
        if self._client:
            await self._client.aclose()
        logger.debug("AsyncLogHiveLogger stopped for project '%s'", self.project)

    async def __aenter__(self):
        await self.start()
        return self

    async def __aexit__(self, *args):
        await self.stop()

    async def _flush_loop(self):
        """Background loop that periodically flushes the queue."""
        while not self._stop_event.is_set():
            await self._flush_now()
            await asyncio.sleep(self.flush_interval)

    async def _flush_now(self):
        """Flush all currently queued entries."""
        entries = []
        while len(entries) < self.batch_size:
            try:
                entry = self._queue.get_nowait()
                entries.append(entry)
            except asyncio.QueueEmpty:
                break

        if not entries:
            return

        await self._send_batch(entries)

    async def _send_batch(self, entries: List[Dict[str, Any]]):
        """Send a batch with retries."""
        url = urljoin(self.endpoint, "/api/logs/ingest")
        payload = {"project": self.project, "entries": entries}

        for attempt in range(self.max_retries):
            try:
                resp = await self._client.post(
                    url,
                    json=payload,
                    headers={"Authorization": f"Bearer {self.api_key}"},
                )
                if resp.status_code == 201:
                    return
                elif resp.status_code == 401:
                    logger.error("LogHive: Invalid API key — dropping batch")
                    return
                else:
                    logger.warning(
                        "LogHive: HTTP %d (attempt %d/%d)",
                        resp.status_code,
                        attempt + 1,
                        self.max_retries,
                    )
            except httpx.RequestError as e:
                logger.warning(
                    "LogHive: Connection error (attempt %d/%d): %s",
                    attempt + 1,
                    self.max_retries,
                    e,
                )

            if attempt < self.max_retries - 1:
                await asyncio.sleep(2 ** attempt)

        logger.error("LogHive: Failed to send %d entries after %d retries", len(entries), self.max_retries)

    def _enqueue(self, level: str, message: str, **kwargs):
        """Enqueue a log entry."""
        entry = {
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "level": level,
            "message": message,
            "logger": kwargs.pop("logger", None) or __name__,
            "extra": kwargs,
        }

        exc_info = kwargs.pop("exc_info", None)
        if exc_info:
            if isinstance(exc_info, BaseException):
                entry["exception"] = "".join(
                    traceback.format_exception(type(exc_info), exc_info, exc_info.__traceback__)
                )
            elif exc_info is True:
                entry["exception"] = traceback.format_exc()

        self._queue.put_nowait(entry)

    # ── Public API ─────────────────────────────────────────────

    async def debug(self, message: str, **kwargs):
        self._enqueue("debug", message, **kwargs)

    async def info(self, message: str, **kwargs):
        self._enqueue("info", message, **kwargs)

    async def warning(self, message: str, **kwargs):
        self._enqueue("warning", message, **kwargs)

    async def error(self, message: str, **kwargs):
        self._enqueue("error", message, **kwargs)

    async def critical(self, message: str, **kwargs):
        self._enqueue("critical", message, **kwargs)