litellm/tests/e2e/logging/datadog_reader.py

"""Read-back for the DataDog logging tests against the real DataDog Logs
Search API.

Delivery is judged on what DataDog itself ingested: the proxy ships logs with
DD_API_KEY exactly as in production (no base-URL override, no local sink), and
the tests search the ingested events back with POST /api/v2/logs/events/search,
authenticated with the same DD_API_KEY plus a DD_APP_KEY application key. On
the cluster the secret manager injects both keys; locally tests/e2e/.env
provides them. Missing keys or a failed search call are hard failures, never an
empty result. External reads go through ``e2e_http``.
"""

from __future__ import annotations

import time
from dataclasses import dataclass

import pytest
from pydantic import BaseModel, ConfigDict, Field

from e2e_config import (
    DD_API_KEY,
    DD_APP_KEY,
    DD_SEARCH_FROM,
    DD_SEARCH_INTERVAL,
    DD_SETTLE_SECONDS,
    DD_SITE,
    POLL_TIMEOUT,
)
from e2e_http import URL, Headers, RateLimitedError, Success, post

#: How many rate-limited responses in a row one search tolerates before the
#: hard fail; each retry sleeps a full search interval, so this rides out a
#: burst from a concurrent consumer of the org-wide search budget.
_RATE_LIMIT_RETRIES = 5


class _DdAuthHeaders(Headers):
    api_key: str = Field(serialization_alias="DD-API-KEY")
    app_key: str = Field(serialization_alias="DD-APPLICATION-KEY")


class _SearchFilter(BaseModel):
    query: str
    #: Wide enough to cover a full suite run plus DataDog's ingestion lag;
    #: markers are unique per test, so a wide window cannot match foreign events.
    #: Override via E2E_DD_SEARCH_FROM when CI lookback needs more than the default.
    from_: str = Field(default_factory=lambda: DD_SEARCH_FROM, serialization_alias="from")
    to: str = "now"


class _SearchPage(BaseModel):
    limit: int = 100


class _SearchRequest(BaseModel):
    filter: _SearchFilter
    page: _SearchPage = _SearchPage()
    sort: str = "timestamp"


class DdLogEvent(BaseModel):
    """One ingested log event as the search API returns it: the indexed
    envelope (service/status/tags) plus ``attributes`` - DataDog's parse of the
    JSON message the integration shipped, i.e. the StandardLoggingPayload
    fields."""

    model_config = ConfigDict(extra="ignore")

    service: str | None = None
    status: str | None = None
    tags: list[str] = []
    attributes: dict[str, object] = {}


class _SearchEvent(BaseModel):
    model_config = ConfigDict(extra="ignore")

    attributes: DdLogEvent


class _SearchResponse(BaseModel):
    model_config = ConfigDict(extra="ignore")

    data: list[_SearchEvent] = []


@dataclass(frozen=True, slots=True)
class DdLogsReader:
    site: str
    api_key: str
    app_key: str

    def events_for_marker(self, marker: str) -> list[DdLogEvent]:
        """Every ingested event whose attributes carry the marker. DataDog
        consumes the shipped JSON message into ``attributes`` and leaves the
        indexed ``message`` empty, so a plain full-text query matches nothing;
        ``*:`` extends the scan to every attribute (the marker sits in the
        prompt, e.g. ``messages.content``, wherever the route's payload puts
        it)."""
        return self.events_for_query(f"*:*{marker}*")

    def events_for_query(self, query: str) -> list[DdLogEvent]:
        """Every ingested event the search query matches (failure payloads
        carry no prompt to mark, so failure scenarios query indexed attributes
        like ``@model_group:...`` instead of a body marker). More than one hit
        for one call IS the duplicate-delivery bug, so this never collapses to
        a single event. A 429 backs off and retries - the search budget is
        org-wide, so another consumer can empty it under us - while any other
        failure stays a hard fail."""
        for _ in range(_RATE_LIMIT_RETRIES):
            result = post(
                URL(f"https://api.{self.site}/api/v2/logs/events/search"),
                headers=_DdAuthHeaders(api_key=self.api_key, app_key=self.app_key),
                json=_SearchRequest(filter=_SearchFilter(query=query)),
                response_type=_SearchResponse,
                timeout=30.0,
            )
            match result:
                case Success(data=page):
                    return [event.attributes for event in page.data]
                case RateLimitedError(retry_after_seconds=retry_after):
                    time.sleep(retry_after if retry_after else DD_SEARCH_INTERVAL)
                case failure:
                    pytest.fail(f"DataDog Logs Search API at api.{self.site} failed: {failure}")
        pytest.fail(
            f"DataDog Logs Search API at api.{self.site} still rate-limited after "
            f"{_RATE_LIMIT_RETRIES} retries {DD_SEARCH_INTERVAL}s apart - the org-wide "
            "logs_public_search_api budget (2 requests per 10s) is exhausted by another consumer"
        )

    def poll_events_for_marker(self, marker: str) -> list[DdLogEvent]:
        """``poll_events_for_query`` over the every-attribute marker scan."""
        return self.poll_events_for_query(f"*:*{marker}*")

    def poll_events_for_query(self, query: str) -> list[DdLogEvent]:
        """Poll until at least one matching event is searchable (the callback
        flushes in periodic batches and DataDog ingestion adds seconds of lag),
        then keep re-reading for DD_SETTLE_SECONDS so a late duplicate cannot
        hide from the exactly-one assertion - real-DataDog jitter can surface
        one call's two events tens of seconds apart. Searches pace at
        DD_SEARCH_INTERVAL, not POLL_INTERVAL, to respect the search API's
        request budget. At the deadline the last result is returned as-is."""
        deadline = time.monotonic() + POLL_TIMEOUT
        while time.monotonic() < deadline:
            events = self.events_for_query(query)
            if events:
                return self._settled_events_for_query(query, events)
            time.sleep(DD_SEARCH_INTERVAL)
        return self.events_for_query(query)

    def _settled_events_for_query(self, query: str, events: list[DdLogEvent]) -> list[DdLogEvent]:
        """Re-read at every search interval until the settle window closes; a
        duplicate ends the watch early because more waiting cannot clear it.

        Keep the last non-empty result: a transient empty search (index lag)
        must not erase events already confirmed earlier in the settle window.
        """
        settle_deadline = time.monotonic() + DD_SETTLE_SECONDS
        last_nonempty = events
        while time.monotonic() < settle_deadline:
            time.sleep(DD_SEARCH_INTERVAL)
            latest = self.events_for_query(query)
            if not latest:
                continue
            if len(latest) > 1:
                return latest
            last_nonempty = latest
        return last_nonempty


def build_dd_logs_reader() -> DdLogsReader:
    if not DD_API_KEY or not DD_APP_KEY:
        pytest.fail(
            "DD_API_KEY and DD_APP_KEY must be set: the DataDog tests deliver to and "
            "read back from the real DataDog API (on the cluster the secret manager "
            "injects them; locally set them in tests/e2e/.env)"
        )
    return DdLogsReader(site=DD_SITE, api_key=DD_API_KEY, app_key=DD_APP_KEY)