openobserve/tests/api-testing/tests/query_agent/conftest.py

"""Shared fixtures for query agent tests.

Generates deterministic data in-memory (no external JSON file needed).
Each query filters to its own time range within the dataset for isolation.

A timestamp-suffixed stream name ensures each test session gets a clean
slate — accumulated data from prior runs cannot bleed into query results.
"""

import json
import logging
import re
import sys
import time
from datetime import datetime, timedelta, UTC
from pathlib import Path

import pytest

from support.client import OpenObserveClient
from support.factories import search_payload
from support.wait import wait_until

# ── Paths ──────────────────────────────────────────────────────────────────
_DATA_DIR = Path(__file__).parent.parent.parent.parent / "test-data" / "query-agent"
_QUERIES_DIR = _DATA_DIR / "queries"
sys.path.insert(0, str(_DATA_DIR))
from data_gen import BASE_TS, build_dataset  # noqa: E402

# Tie stream name to BASE_TS (stable within a session: comes from
# base_ts_override.json written by compute_counts.py).  PID suffix
# prevents data accumulation across multiple test invocations against
# the same environment.
import os as _os
STREAM = f"query_agent_test_{BASE_TS}_{_os.getpid()}"
STREAM2 = f"query_agent_test2_{BASE_TS}_{_os.getpid()}"


# ── Shared query loader ────────────────────────────────────────────────────
def _values_equal(a: str, b: str, rel_tol: float = 0.05) -> bool:
    """Compare two value strings, with relative tolerance for floats.

    Cross-engine floating-point aggregates (AVG, STDDEV, etc.) differ
    at the least-significant digits due to summation order.  A 1 %
    relative tolerance catches these while keeping exact match for
    strings, integers, and timestamps.
    """
    if a == b:
        return True
    # OO serializes NULL as "" while DuckDB uses "None" — treat as equal
    if (a == "" and b == "None") or (a == "None" and b == ""):
        return True
    try:
        fa = float(a)
        fb = float(b)
        denom = max(abs(fa), abs(fb), 1.0)
        return abs(fa - fb) / denom < rel_tol
    except (ValueError, TypeError):
        return False
def load_all_queries():
    """Return {category_name: [queries]} from queries/*.json files."""
    categories: dict[str, list[dict]] = {}
    for fp in sorted(_QUERIES_DIR.glob("*.json")):
        with open(fp) as f:
            data = json.load(f)
        categories[fp.stem] = data["queries"]
    return categories


# ── Shared query runner ────────────────────────────────────────────────────
def run_query(client, query, *, skip_fts_count=False):
    """Execute a single query via _search and assert expected results.

    Two modes, determined by the presence of ``results`` in expected:

    * sqllogictest mode (results present): extract column values from hits
      as strings, sort if results_mode is "rowsort", then compare
      cell-by-cell against the expected result set computed by DuckDB.

    * Legacy assertion mode (results absent — histogram queries, Q122,
      and a few complex subqueries): check row_count, column presence,
      and run explicit content assertions.

    When *skip_fts_count* is True, full_text_search row_count is skipped
    (Tantivy FTS not available pre-flush).
    """
    qid = query["id"]
    size = 500
    offset = query["time_offset"]
    sql = query["sql"].replace("{stream}", STREAM).replace("{stream2}", STREAM2)
    json_data = search_payload(
        sql,
        start_time=BASE_TS + offset["start_offset"],
        end_time=BASE_TS + offset["end_offset"],
        size=size,
    )
    resp = client.post("_search?type=logs", json=json_data)
    assert resp.status_code == 200, \
        f"{qid}: Expected 200, got {resp.status_code}: {resp.text[:500]}"

    response_data = resp.json()
    assert "hits" in response_data, f"{qid}: Response should contain 'hits'"
    hits = response_data["hits"]

    expected = query.get("expected", {})
    is_fts = query.get("category") == "full_text_search"

    if "results" in expected and not (skip_fts_count and is_fts) and not expected.get("skip_sqllogictest"):
        # ── sqllogictest mode: result-set comparison with float tolerance ─
        cols = expected["columns"]
        actual = sorted(
            [str(hit.get(col, "")) for col in cols] for hit in hits
        )
        want = sorted(expected["results"])

        assert len(actual) == len(want), \
            f"{qid}: row count mismatch — got {len(actual)}, expected {len(want)}"

        for i, (got_row, want_row) in enumerate(zip(actual, want)):
            assert len(got_row) == len(want_row), \
                f"{qid}: column count mismatch at row {i}"
            for j, (gv, wv) in enumerate(zip(got_row, want_row)):
                if not _values_equal(gv, wv):
                    assert False, \
                        f"{qid}: mismatch at row {i}, col {j} ('{cols[j]}'): got {gv!r}, expected {wv!r}"

    else:
        # ── Legacy assertion mode (histogram, broken queries, complex CTEs)
        if expected.get("row_count") is not None:
            if not (skip_fts_count and is_fts) and not expected.get("skip_row_count"):
                assert len(hits) == expected["row_count"], \
                    f"{qid}: Expected {expected['row_count']} rows, got {len(hits)}"

        assert len(hits) < size, \
            f"{qid}: Got {len(hits)} rows at size={size} — result may be truncated"

        if len(hits) > 0 and not expected.get("skip_column_check"):
            for col in expected.get("columns", []):
                assert col in hits[0], f"{qid}: Expected column '{col}' not in response"

        _run_content_assertions(query, sql, hits)

    logging.info("%s passed: %d rows", qid, len(hits))


# ── Legacy Content Assertions (for queries without DuckDB results) ──────────
def _paren_depth(text, pos):
    """Return the nesting depth of parentheses at *pos* within *text*."""
    depth = 0
    for c in text[:pos]:
        if c == "(":
            depth += 1
        elif c == ")":
            depth -= 1
    return depth


def _run_content_assertions(query, sql, hits):
    """Validate result content via explicit assertions in the query JSON.

    Auto-inferred: ORDER BY direction from the outermost SQL ORDER BY clause.
    Explicit (per-query in ``expected.assertions``): row_contains, value_range.
    """
    qid = query["id"]
    expected = query.get("expected", {})

    assertions: list[dict] = list(expected.get("assertions", []))

    # Auto-infer order from the outermost SQL ORDER BY clause.
    order_matches = [
        m for m in re.finditer(
            r"ORDER\s+BY\s+([^\s,]+)\s*(ASC|DESC)?",
            sql, re.IGNORECASE,
        )
        if _paren_depth(sql, m.start()) == 0
    ]
    if order_matches:
        m = order_matches[-1]
        col = m.group(1).strip().strip('"').strip("'")
        direction = (m.group(2) or "ASC").upper()
        assertions.append({"type": "order", "column": col, "direction": direction})

    for i, a in enumerate(assertions):
        a_type = a["type"]
        label = f"{qid}: [{i}] {a_type}"

        if a_type == "order":
            col = a["column"]
            direction = a.get("direction", "ASC")
            values = [hit.get(col) for hit in hits if col in hit]
            if len(values) < 2:
                continue
            ordered = values == sorted(values) if direction == "ASC" \
                else values == sorted(values, reverse=True)
            assert ordered, \
                f"{label}: '{col}' not in {direction} order — {values[:5]}..."

        elif a_type == "row_contains":
            match = a["match"]
            found = any(
                all(hit.get(k) == v for k, v in match.items())
                for hit in hits
            )
            assert found, \
                f"{label}: no row matches {match}"

        elif a_type == "value_range":
            col = a["column"]
            for j, hit in enumerate(hits):
                val = hit.get(col)
                if val is None:
                    continue
                if "min" in a:
                    assert val >= a["min"], \
                        f"{label}: {col}={val} < min={a['min']} in row {j}"
                if "max" in a:
                    assert val <= a["max"], \
                        f"{label}: {col}={val} > max={a['max']} in row {j}"


# ── Session fixture: ingest data (no flush) ────────────────────────────────
@pytest.fixture(scope="session")
def ingest_query_agent_data():
    """Ingest deterministic data into both streams and wait for searchability.

    Does NOT flush — data stays in memtable for the pre-flush phase.
    Stream 2 uses stream_offset=7 for different but join-key-overlapping data.
    """

    client = OpenObserveClient()

    # Stream 1 (primary)
    records1 = build_dataset()
    data1 = json.dumps(records1)
    resp1 = client.post(f"{STREAM}/_json", data=data1,
                        headers={"Content-Type": "application/json"})
    if resp1.status_code == 200:
        logging.info("Stream1: ingested %d records (%d bytes)", len(records1), len(data1))
    else:
        pytest.fail(f"Stream1 ingestion failed: {resp1.status_code} — {resp1.text[:300]}")

    # Stream 2 (secondary, offset=7 for different but overlapping join keys)
    records2 = build_dataset(stream_offset=7)
    data2 = json.dumps(records2)
    resp2 = client.post(f"{STREAM2}/_json", data=data2,
                        headers={"Content-Type": "application/json"})
    if resp2.status_code == 200:
        logging.info("Stream2: ingested %d records (%d bytes)", len(records2), len(data2))
    else:
        pytest.fail(f"Stream2 ingestion failed: {resp2.status_code} — {resp2.text[:300]}")

    # Poll until all records from both streams are searchable.
    # Each stream is checked independently so cross-stream queries
    # never run against a partially-ready stream.
    expected1 = len(records1)
    expected2 = len(records2)
    _max_ts = max(r["_timestamp"] for r in records1 + records2)

    # Vortex servers do not expose WAL/memtable data — searches return
    # NewEmptyExec until data is flushed to vortex files.  Detect this by
    # attempting 3 quick polls; if no data appears, flush once and continue.
    def _quick_count(stream: str) -> int:
        now = datetime.now(UTC)
        end_us = max(int(now.timestamp() * 1_000_000), _max_ts) + 3_600_000_000
        start_us = int((now - timedelta(weeks=4)).timestamp() * 1_000_000)
        r = client.post("_search?type=logs", json={
            "query": {"sql": f'SELECT COUNT(*) AS c FROM "{stream}"',
                      "start_time": start_us, "end_time": end_us,
                      "from": 0, "size": 1}
        })
        if r.status_code != 200:
            return 0
        hits = r.json().get("hits", [])
        return int(hits[0].get("c", 0)) if hits else 0

    _memtable_visible = False
    for _ in range(3):
        if _quick_count(STREAM) > 0:
            _memtable_visible = True
            break
        time.sleep(1.0)

    if not _memtable_visible:
        # No memtable data after 3s — flush so vortex files are created.
        flush_r = client.put("node/flush", prefix="")
        logging.info("Auto-flush after 3s (vortex or slow memtable): %s", flush_r.status_code)
        time.sleep(2.0)

    def _data_is_searchable():
        now = datetime.now(UTC)
        # Records have timestamps from BASE_TS to BASE_TS + ~66 min (future).
        # Extend end_us past _max_ts so the COUNT(*) window covers all records.
        # The vortex search engine enforces end_time strictly; parquet WAL does not.
        end_us = max(int(now.timestamp() * 1_000_000), _max_ts) + 3_600_000_000
        start_us = int((now - timedelta(weeks=4)).timestamp() * 1_000_000)
        # Check stream 1
        r1 = client.post("_search?type=logs", json={
            "query": {
                "sql": f'SELECT COUNT(*) AS c FROM "{STREAM}"',
                "start_time": start_us,
                "end_time": end_us,
                "from": 0,
                "size": 1,
            }
        })
        if r1.status_code != 200:
            return False
        hits1 = r1.json().get("hits", [])
        if not (hits1 and hits1[0].get("c", 0) >= expected1):
            return False
        # Check stream 2
        r2 = client.post("_search?type=logs", json={
            "query": {
                "sql": f'SELECT COUNT(*) AS c FROM "{STREAM2}"',
                "start_time": start_us,
                "end_time": end_us,
                "from": 0,
                "size": 1,
            }
        })
        if r2.status_code != 200:
            return False
        hits2 = r2.json().get("hits", [])
        if not (hits2 and hits2[0].get("c", 0) >= expected2):
            return False
        # Tail poll: verify a record near max_ts is searchable
        tail_r = client.post("_search?type=logs", json={
            "query": {
                "sql": f'SELECT COUNT(*) AS c FROM "{STREAM}"',
                "start_time": _max_ts - 60_000_000,
                "end_time": _max_ts + 3_600_000_000,
                "from": 0,
                "size": 1,
            }
        })
        if tail_r.status_code != 200:
            return False
        tail_hits = tail_r.json().get("hits", [])
        return bool(tail_hits and tail_hits[0].get("c", 0) >= 1)

    wait_until(_data_is_searchable, timeout=300, interval=1.0,
               msg=f"Stream data not searchable ({len(records1)} records)")
    logging.info("Both streams data is searchable (%d + %d records)", len(records1), len(records2))