openobserve/tests/api-testing/support/sse.py

"""Server-Sent Events (SSE) response parser for OO's `/_search_stream` and
`/_values_stream` endpoints.

OO returns SSE-formatted responses for streaming queries:
    event: search_response_metadata
    data: {"results": {"total": N, ...}}

    event: search_response_hits
    data: {"hits": [...]}

The endpoint may emit MULTIPLE metadata + hits frames per response (one
per partition for aggregations, etc.). This helper merges them into a
single `{"results": {"total": int, "hits": [...]}}` dict that callers
can treat uniformly.

Replaces the `read_response()` function previously inlined in
test_streaming.py.
"""
from __future__ import annotations

import json
import logging
from typing import Any

import requests

logger = logging.getLogger(__name__)


def read_sse_frames(response: requests.Response) -> list[tuple[str, dict[str, Any]]]:
    """Parse an SSE-streamed response into its raw frames, in wire order.

    Returns one `(event_name, data)` pair per `event:` / `data:` line pair,
    with NO merging and NO interpretation — e.g.::

        [("search_response_metadata", {"results": {...}, "streaming_aggs": True}),
         ("search_response_hits",     {"hits": [...]}),
         ("search_response_metadata", {...}),
         ...]

    Use this (instead of `read_sse_response()`) whenever the PER-FRAME
    semantics are the thing under test. `read_sse_response()` deliberately
    collapses the frames — it keeps `max(totals)` and concatenates every hits
    frame — so it cannot tell a disjoint page from a cumulative aggregation
    state, and concatenating hits across cumulative frames duplicates rows.

    Frames whose `data:` payload is not valid JSON are skipped (logged at
    debug), matching `read_sse_response()`'s tolerance.
    """
    content = response.content.decode("utf-8")
    lines = content.split("\n")
    frames: list[tuple[str, dict[str, Any]]] = []

    for i, line in enumerate(lines):
        text = line.strip()
        if not text.startswith("event: "):
            continue
        event = text[len("event: "):].strip()
        if i + 1 >= len(lines):
            continue
        data_line = lines[i + 1].strip()
        if not data_line.startswith("data: "):
            continue
        try:
            frames.append((event, json.loads(data_line[6:])))
        except json.JSONDecodeError as e:
            logger.debug("SSE %s frame parse error: %s", event, e)

    return frames


def read_sse_response(response: requests.Response) -> dict[str, Any]:
    """Parse an SSE-streamed response into one merged `{results: {total, hits}}` dict.

    Multiple metadata frames: we pick the MAX total (for aggregation queries,
    each partition reports a partial; max is the right global). Hits are
    concatenated across all hits frames.

    Returns `{"results": {"total": 0, "hits": []}}` if no valid frames found.
    """
    content = response.content.decode("utf-8")
    lines = content.split("\n")
    metadata_frames: list[dict[str, Any]] = []
    hits_frames: list[dict[str, Any]] = []

    for i, line in enumerate(lines):
        text = line.strip()

        if text.startswith("event: search_response_metadata"):
            if i + 1 < len(lines):
                data_line = lines[i + 1].strip()
                if data_line.startswith("data: "):
                    try:
                        metadata_frames.append(json.loads(data_line[6:]))
                    except json.JSONDecodeError as e:
                        logger.debug("SSE metadata frame parse error: %s", e)

        elif text.startswith("event: search_response_hits"):
            if i + 1 < len(lines):
                data_line = lines[i + 1].strip()
                if data_line.startswith("data: "):
                    try:
                        hits_frames.append(json.loads(data_line[6:]))
                    except json.JSONDecodeError as e:
                        logger.debug("SSE hits frame parse error: %s", e)

    if not metadata_frames:
        logger.debug("read_sse_response: no metadata frames found")
        return {"results": {"total": 0, "hits": []}}

    # Concatenate all hits across hits frames
    all_hits: list[dict[str, Any]] = []
    for frame in hits_frames:
        if "hits" in frame and isinstance(frame["hits"], list):
            all_hits.extend(frame["hits"])

    # Determine combined total. For single-frame responses use the reported
    # total directly. For multi-frame, pick the MAX (each partition reports
    # its partial; for aggregation queries the max is the global total).
    combined = metadata_frames[0].copy()
    if len(metadata_frames) == 1:
        total = metadata_frames[0].get("results", {}).get("total", len(all_hits))
    else:
        totals = [
            f.get("results", {}).get("total", 0) for f in metadata_frames
        ]
        total = max(totals) if any(t for t in totals) else len(all_hits)

    combined.setdefault("results", {})
    combined["results"]["total"] = total
    combined["results"]["hits"] = all_hits
    return combined