From 8523cc80a652726ffc228744f760c512ba11a4ea Mon Sep 17 00:00:00 2001 From: Brian Newsom Date: Mon, 27 Jul 2026 14:26:31 -0600 Subject: [PATCH 1/7] test(intake): pin ClickHouse 26.3 LTS Signed-off-by: Brian Newsom --- services/intake/.clickhouse-version | 1 + services/intake/README.md | 3 +++ services/intake/scripts/spans/run_clickhouse.sh | 15 ++++++++++++++- .../intake/spans/evaluation_session_repository.py | 4 ++-- .../intake/tests/integration/spans/conftest.py | 11 ++++++++++- .../spans/test_clickhouse_bootstrap.py | 11 +++++++++++ 6 files changed, 41 insertions(+), 4 deletions(-) create mode 100644 services/intake/.clickhouse-version diff --git a/services/intake/.clickhouse-version b/services/intake/.clickhouse-version new file mode 100644 index 0000000000..59c8fb1977 --- /dev/null +++ b/services/intake/.clickhouse-version @@ -0,0 +1 @@ +26.3 diff --git a/services/intake/README.md b/services/intake/README.md index fc02964f24..0306eba666 100644 --- a/services/intake/README.md +++ b/services/intake/README.md @@ -54,6 +54,9 @@ environment instead of package-scoped `uv run --package ...` commands. Prerequisite: Docker must be installed and running locally. +Intake is tested and profiled on ClickHouse 26.3 LTS. Other ClickHouse versions +may not be supported. + Start a local ClickHouse container for span and trace storage: ```bash diff --git a/services/intake/scripts/spans/run_clickhouse.sh b/services/intake/scripts/spans/run_clickhouse.sh index 17829ca9a3..bced9e21cf 100755 --- a/services/intake/scripts/spans/run_clickhouse.sh +++ b/services/intake/scripts/spans/run_clickhouse.sh @@ -5,11 +5,12 @@ set -euo pipefail container_name="nmp-intake-clickhouse" -image="${CLICKHOUSE_IMAGE:-clickhouse/clickhouse-server:24.3}" clickhouse_user="${CLICKHOUSE_USER:-default}" clickhouse_password="${CLICKHOUSE_PASSWORD:-}" script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" repo_root="$(cd -- "${script_dir}/../../../.." && pwd)" +clickhouse_version="$(tr -d '[:space:]' < "${script_dir}/../../.clickhouse-version")" +image="${CLICKHOUSE_IMAGE:-clickhouse/clickhouse-server:${clickhouse_version}}" data_dir="${CLICKHOUSE_DATA_DIR:-${repo_root}/tmp/intake-clickhouse}" ensure_host_dirs() { @@ -21,13 +22,25 @@ ensure_tmp_dir() { docker exec "${container_name}" sh -c "mkdir -p /var/lib/clickhouse/tmp && chown clickhouse:clickhouse /var/lib/clickhouse/tmp" >/dev/null } +ensure_expected_image() { + local actual_image + actual_image="$(docker inspect --format "{{.Config.Image}}" "${container_name}")" + if [[ "${actual_image}" != "${image}" ]]; then + echo "${container_name} uses ${actual_image}, but Intake requires ${image}." >&2 + echo "Remove the existing container after preserving any data you need, then rerun this script." >&2 + exit 1 + fi +} + if docker ps --filter "name=^/${container_name}$" --filter "status=running" --format "{{.Names}}" | grep -qx "${container_name}"; then + ensure_expected_image ensure_tmp_dir echo "${container_name} is already running" exit 0 fi if docker ps -a --filter "name=^/${container_name}$" --format "{{.Names}}" | grep -qx "${container_name}"; then + ensure_expected_image ensure_host_dirs docker start "${container_name}" >/dev/null ensure_tmp_dir diff --git a/services/intake/src/nmp/intake/spans/evaluation_session_repository.py b/services/intake/src/nmp/intake/spans/evaluation_session_repository.py index b809c82589..839072b031 100644 --- a/services/intake/src/nmp/intake/spans/evaluation_session_repository.py +++ b/services/intake/src/nmp/intake/spans/evaluation_session_repository.py @@ -171,7 +171,7 @@ async def list_sessions( if needs_pre_metrics: # Two-query path for cost/tokens sorts. # - # ClickHouse 24.3 inlines CTEs (does not materialise them), so a single query that + # ClickHouse inlines regular CTEs (does not materialise them), so a single query that # references `page_sessions` from multiple downstream CTEs (current_page_spans, # session_metrics, session_scores, final SELECT) would re-execute the expensive # all-session span aggregation once per reference. Splitting into two queries @@ -340,7 +340,7 @@ def _metric_sort_page_ids_sql( ORDER BY + LIMIT/OFFSET to return the ordered (workspace, session_id) pairs for the requested page. Only IDs are returned — row hydration is a separate query. - Why separate: ClickHouse 24.3 inlines CTEs rather than materialising them, so a + Why separate: ClickHouse inlines regular CTEs rather than materialising them, so a single query that references `page_sessions` from multiple CTEs would re-execute the expensive all-session span aggregation once per reference. Returning IDs here and hydrating in _hydrate_by_ids_sql ensures the aggregation runs exactly once. diff --git a/services/intake/tests/integration/spans/conftest.py b/services/intake/tests/integration/spans/conftest.py index 65b1001783..88946c6061 100644 --- a/services/intake/tests/integration/spans/conftest.py +++ b/services/intake/tests/integration/spans/conftest.py @@ -9,6 +9,7 @@ from collections.abc import Callable from datetime import datetime, timezone from importlib.util import find_spec +from pathlib import Path from typing import Any from uuid import uuid4 @@ -23,6 +24,9 @@ ) from nmp.testing import create_test_client +_CLICKHOUSE_VERSION_FILE = Path(__file__).resolve().parents[3] / ".clickhouse-version" +CLICKHOUSE_VERSION = _CLICKHOUSE_VERSION_FILE.read_text(encoding="utf-8").strip() + def _run(coro: Any) -> Any: return asyncio.run(coro) @@ -55,7 +59,7 @@ def clickhouse_container(): from testcontainers.clickhouse import ClickHouseContainer with ClickHouseContainer( - "clickhouse/clickhouse-server:24.3", + f"clickhouse/clickhouse-server:{CLICKHOUSE_VERSION}", username="test", password="test", dbname="default", @@ -63,6 +67,11 @@ def clickhouse_container(): yield container +@pytest.fixture(scope="session") +def clickhouse_version() -> str: + return CLICKHOUSE_VERSION + + @pytest.fixture(scope="session") def clickhouse_settings(clickhouse_container) -> ClickHouseSettings: return ClickHouseSettings( diff --git a/services/intake/tests/integration/spans/test_clickhouse_bootstrap.py b/services/intake/tests/integration/spans/test_clickhouse_bootstrap.py index db26401b5b..9d6149b60e 100644 --- a/services/intake/tests/integration/spans/test_clickhouse_bootstrap.py +++ b/services/intake/tests/integration/spans/test_clickhouse_bootstrap.py @@ -11,6 +11,17 @@ from nmp.intake.spans.clickhouse_client import ClickHouseSpanClient, bootstrap_schema +def test_clickhouse_server_matches_supported_lts( + clickhouse_client: ClickHouseSpanClient, + clickhouse_version: str, + run_async, +) -> None: + result = run_async(clickhouse_client.query("SELECT version()")) + + assert len(result.result_rows) == 1 + assert str(result.result_rows[0][0]).startswith(f"{clickhouse_version}.") + + def test_clickhouse_bootstrap_is_idempotent(clickhouse_client: ClickHouseSpanClient, run_async): run_async(bootstrap_schema(clickhouse_client)) run_async(bootstrap_schema(clickhouse_client)) From 02f0e8097232f4ea540b9357c31a2a0ef45d7552 Mon Sep 17 00:00:00 2001 From: Brian Newsom Date: Mon, 27 Jul 2026 15:04:10 -0600 Subject: [PATCH 2/7] refactor(intake): add ClickHouse repository boundary Signed-off-by: Brian Newsom --- .../intake/repository/clickhouse/executor.py | 77 ++++++++++++++++ .../clickhouse/session.py} | 48 +++++----- .../intake/repository/clickhouse/tables.py | 23 +++++ .../src/nmp/intake/repository/session.py | 17 ++++ .../src/nmp/intake/spans/api/dependencies.py | 14 ++- .../intake/src/nmp/intake/spans/service.py | 2 +- .../integration/spans/test_sessions_read.py | 8 +- .../tests/test_clickhouse_architecture.py | 43 +++++++++ .../intake/tests/test_clickhouse_executor.py | 92 +++++++++++++++++++ .../test_sessions_clickhouse_repository.py | 83 ++++++++--------- 10 files changed, 331 insertions(+), 76 deletions(-) create mode 100644 services/intake/src/nmp/intake/repository/clickhouse/executor.py rename services/intake/src/nmp/intake/{spans/session_repository.py => repository/clickhouse/session.py} (73%) create mode 100644 services/intake/src/nmp/intake/repository/clickhouse/tables.py create mode 100644 services/intake/src/nmp/intake/repository/session.py create mode 100644 services/intake/tests/test_clickhouse_architecture.py create mode 100644 services/intake/tests/test_clickhouse_executor.py diff --git a/services/intake/src/nmp/intake/repository/clickhouse/executor.py b/services/intake/src/nmp/intake/repository/clickhouse/executor.py new file mode 100644 index 0000000000..c071a5663d --- /dev/null +++ b/services/intake/src/nmp/intake/repository/clickhouse/executor.py @@ -0,0 +1,77 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Typed runtime query boundary for Intake ClickHouse repositories.""" + +from __future__ import annotations + +import logging +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +from time import perf_counter +from typing import Any + +from clickhouse_connect.driver.exceptions import ClickHouseError +from nmp.intake.repository.clickhouse.tables import ClickHouseTable, qualified_table +from nmp.intake.spans.clickhouse_client import ClickHouseSpanClient + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class ClickHouseQuery: + """One named, parameterized ClickHouse read statement.""" + + name: str + statement: str + parameters: Mapping[str, object] = field(default_factory=dict) + + def bind(self, **parameters: object) -> ClickHouseQuery: + """Return a copy with additional bound parameters.""" + + return ClickHouseQuery( + name=self.name, + statement=self.statement, + parameters={**self.parameters, **parameters}, + ) + + +class ClickHouseQueryError(RuntimeError): + """Raised when a named repository query fails.""" + + def __init__(self, query_name: str) -> None: + self.query_name = query_name + super().__init__(f"ClickHouse query failed: {query_name}") + + +class ClickHouseExecutor: + """Execute named repository queries without exposing the raw driver result.""" + + def __init__(self, client: ClickHouseSpanClient) -> None: + self._client = client + + def table(self, table: ClickHouseTable) -> str: + return qualified_table(self._client.database, table) + + async def fetch_all(self, query: ClickHouseQuery) -> list[dict[str, Any]]: + started_at = perf_counter() + try: + result = await self._client.query( + query.statement, + parameters=dict(query.parameters), + ) + except ClickHouseError as exc: + logger.exception("ClickHouse repository query failed", extra={"query_name": query.name}) + raise ClickHouseQueryError(query.name) from exc + finally: + logger.debug( + "ClickHouse repository query finished", + extra={ + "query_name": query.name, + "duration_ms": (perf_counter() - started_at) * 1000, + }, + ) + + columns: Sequence[str] = result.column_names + rows: Sequence[Sequence[Any]] = result.result_rows + return [dict(zip(columns, row, strict=True)) for row in rows] diff --git a/services/intake/src/nmp/intake/spans/session_repository.py b/services/intake/src/nmp/intake/repository/clickhouse/session.py similarity index 73% rename from services/intake/src/nmp/intake/spans/session_repository.py rename to services/intake/src/nmp/intake/repository/clickhouse/session.py index 3159555216..eb7e5054ed 100644 --- a/services/intake/src/nmp/intake/spans/session_repository.py +++ b/services/intake/src/nmp/intake/repository/clickhouse/session.py @@ -1,50 +1,40 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""ClickHouse implementation of Intake session detail reads.""" +"""ClickHouse implementation of Intake session reads.""" from __future__ import annotations from datetime import datetime from typing import Any -from nmp.intake.spans.clickhouse_client import ClickHouseSpanClient +from nmp.intake.repository.clickhouse.executor import ClickHouseExecutor, ClickHouseQuery +from nmp.intake.repository.clickhouse.tables import ClickHouseTable +from nmp.intake.repository.session import SessionRepository from nmp.intake.spans.domain import IntakeSession -from nmp.intake.spans.span_rollups import METRIC_ATTRIBUTE_FIELDS, metric_aggregate_columns -from nmp.intake.spans.storage import float_or_none, int_or_none, normalize_span_status, result_rows +from nmp.intake.spans.span_rollups import metric_aggregate_columns +from nmp.intake.spans.storage import float_or_none, int_or_none, normalize_span_status -SESSION_COLUMNS = [ - "id", - "workspace", - "started_at", - "ended_at", - "status", - *METRIC_ATTRIBUTE_FIELDS.keys(), - "trace_count", - "span_count", -] - -class SessionRepository: - def __init__(self, client: ClickHouseSpanClient) -> None: - self._client = client +class ClickHouseSessionRepository(SessionRepository): + def __init__(self, executor: ClickHouseExecutor) -> None: + self._executor = executor async def get_session(self, *, workspace: str, session_id: str) -> IntakeSession | None: - query, parameters = session_detail_sql(self._client.table("spans")) - result = await self._client.query( - query, - parameters={**parameters, "workspace": workspace, "session_id": session_id}, + query = _session_detail_query(self._executor.table(ClickHouseTable.SPANS)).bind( + workspace=workspace, + session_id=session_id, ) - rows = result_rows(result) + rows = await self._executor.fetch_all(query) return _row_to_session(rows[0]) if rows else None -def session_detail_sql(table: str) -> tuple[str, dict[str, Any]]: - """Return a primary-key-pruned aggregate over the current rows of one session.""" +def _session_detail_query(table: str) -> ClickHouseQuery: + """Build a primary-key-pruned aggregate over the current rows of one session.""" source_alias = "session_spans" metric_columns, parameters = metric_aggregate_columns(source_alias) - query = f""" + statement = f""" SELECT %(session_id)s AS id, any({source_alias}.workspace) AS workspace, @@ -70,7 +60,11 @@ def session_detail_sql(table: str) -> tuple[str, dict[str, Any]]: WHERE {source_alias}.is_deleted = 0 HAVING span_count > 0 """ - return query, parameters + return ClickHouseQuery( + name="sessions.get", + statement=statement, + parameters=parameters, + ) def _row_to_session(row: dict[str, Any]) -> IntakeSession: diff --git a/services/intake/src/nmp/intake/repository/clickhouse/tables.py b/services/intake/src/nmp/intake/repository/clickhouse/tables.py new file mode 100644 index 0000000000..1aacd0305a --- /dev/null +++ b/services/intake/src/nmp/intake/repository/clickhouse/tables.py @@ -0,0 +1,23 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Closed registry of runtime-owned Intake ClickHouse tables.""" + +from enum import StrEnum + +from nmp.intake.spans.clickhouse_migrations import quote_clickhouse_identifier + + +class ClickHouseTable(StrEnum): + ANNOTATIONS = "annotations" + EVALUATOR_RESULTS = "evaluator_results" + SPANS = "spans" + TRACE_INDEX = "trace_index" + + +def qualified_table(database: str, table: ClickHouseTable) -> str: + """Return a safely quoted table name from the closed runtime registry.""" + + if not isinstance(table, ClickHouseTable): + raise TypeError(f"Expected ClickHouseTable, got {type(table).__name__}") + return f"{quote_clickhouse_identifier(database)}.{quote_clickhouse_identifier(table.value)}" diff --git a/services/intake/src/nmp/intake/repository/session.py b/services/intake/src/nmp/intake/repository/session.py new file mode 100644 index 0000000000..8443e932a3 --- /dev/null +++ b/services/intake/src/nmp/intake/repository/session.py @@ -0,0 +1,17 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Repository interface for Intake session reads.""" + +from abc import ABC, abstractmethod + +from nmp.intake.spans.domain import IntakeSession + + +class SessionRepository(ABC): + """Domain-facing interface for session persistence.""" + + @abstractmethod + async def get_session(self, *, workspace: str, session_id: str) -> IntakeSession | None: + """Return one session, or ``None`` when it has no current spans.""" + pass diff --git a/services/intake/src/nmp/intake/spans/api/dependencies.py b/services/intake/src/nmp/intake/spans/api/dependencies.py index ec682b9bd7..7077dbf895 100644 --- a/services/intake/src/nmp/intake/spans/api/dependencies.py +++ b/services/intake/src/nmp/intake/spans/api/dependencies.py @@ -8,11 +8,13 @@ from fastapi import Depends, HTTPException, Request, status from nemo_platform import AsyncNeMoPlatform from nmp.common.service.dependencies import get_sdk_client +from nmp.intake.repository.clickhouse.executor import ClickHouseExecutor +from nmp.intake.repository.clickhouse.session import ClickHouseSessionRepository +from nmp.intake.repository.session import SessionRepository from nmp.intake.spans.annotations_repository import AnnotationsRepository from nmp.intake.spans.clickhouse_client import ClickHouseSpanClient, get_clickhouse_client from nmp.intake.spans.evaluator_results_repository import EvaluatorResultsRepository from nmp.intake.spans.service import IntakeSpansService -from nmp.intake.spans.session_repository import SessionRepository from nmp.intake.spans.span_repository import SpanRepository from nmp.intake.spans.trace_repository import TraceRepository @@ -58,10 +60,16 @@ def get_trace_repository( return TraceRepository(client) -def get_session_repository( +def get_clickhouse_executor( client: Annotated[ClickHouseSpanClient, Depends(get_clickhouse_client)], +) -> ClickHouseExecutor: + return ClickHouseExecutor(client) + + +def get_session_repository( + executor: Annotated[ClickHouseExecutor, Depends(get_clickhouse_executor)], ) -> SessionRepository: - return SessionRepository(client) + return ClickHouseSessionRepository(executor) def get_evaluator_results_repository( diff --git a/services/intake/src/nmp/intake/spans/service.py b/services/intake/src/nmp/intake/spans/service.py index 2340b029e8..35a09f6e75 100644 --- a/services/intake/src/nmp/intake/spans/service.py +++ b/services/intake/src/nmp/intake/spans/service.py @@ -8,6 +8,7 @@ from datetime import datetime from nmp.common.api.common import PaginatedResult +from nmp.intake.repository.session import SessionRepository from nmp.intake.spans.annotations_repository import AnnotationsRepository from nmp.intake.spans.domain import ( Annotation, @@ -25,7 +26,6 @@ TraceMode, ) from nmp.intake.spans.evaluator_results_repository import EvaluatorResultsRepository -from nmp.intake.spans.session_repository import SessionRepository from nmp.intake.spans.span_repository import SpanRepository from nmp.intake.spans.trace_repository import TraceRepository diff --git a/services/intake/tests/integration/spans/test_sessions_read.py b/services/intake/tests/integration/spans/test_sessions_read.py index d4c4b006ae..ba51999859 100644 --- a/services/intake/tests/integration/spans/test_sessions_read.py +++ b/services/intake/tests/integration/spans/test_sessions_read.py @@ -7,8 +7,8 @@ from decimal import Decimal from fastapi.testclient import TestClient +from nmp.intake.repository.clickhouse.session import _session_detail_query from nmp.intake.spans.clickhouse_client import ClickHouseSpanClient -from nmp.intake.spans.session_repository import session_detail_sql def test_session_detail_rolls_up_all_current_spans( @@ -93,11 +93,11 @@ def test_session_detail_rolls_up_all_current_spans( assert "input" not in session assert "output" not in session - query, parameters = session_detail_sql(clickhouse_client.table("spans")) + query = _session_detail_query(clickhouse_client.table("spans")) plan = run_async( clickhouse_client.query( - f"EXPLAIN indexes = 1 {query}", - parameters={**parameters, "workspace": "default", "session_id": session_id}, + f"EXPLAIN indexes = 1 {query.statement}", + parameters={**query.parameters, "workspace": "default", "session_id": session_id}, ) ) plan_text = "\n".join(str(row[0]) for row in plan.result_rows) diff --git a/services/intake/tests/test_clickhouse_architecture.py b/services/intake/tests/test_clickhouse_architecture.py new file mode 100644 index 0000000000..4d23dff919 --- /dev/null +++ b/services/intake/tests/test_clickhouse_architecture.py @@ -0,0 +1,43 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Import-boundary tests for Intake ClickHouse persistence.""" + +import ast +from pathlib import Path + +_SOURCE_ROOT = Path(__file__).resolve().parents[1] / "src" / "nmp" / "intake" +_RAW_CLIENT_MODULE = "nmp.intake.spans.clickhouse_client" + +# Shrink this allowlist as each legacy repository moves behind ClickHouseExecutor. +_ALLOWED_RAW_CLIENT_IMPORTS = { + "api/v2/experiments/endpoints.py", + "repository/clickhouse/executor.py", + "service.py", + "spans/annotations_repository.py", + "spans/api/dependencies.py", + "spans/evaluation_rollup_repository.py", + "spans/evaluation_session_repository.py", + "spans/evaluator_results_repository.py", + "spans/span_repository.py", + "spans/trace_repository.py", +} + + +def test_raw_clickhouse_client_imports_are_confined_to_approved_modules() -> None: + imports = { + path.relative_to(_SOURCE_ROOT).as_posix() for path in _SOURCE_ROOT.rglob("*.py") if _imports_raw_client(path) + } + + unexpected = imports - _ALLOWED_RAW_CLIENT_IMPORTS + assert not unexpected, f"Use ClickHouseExecutor instead of the raw client in: {sorted(unexpected)}" + + +def _imports_raw_client(path: Path) -> bool: + tree = ast.parse(path.read_text(encoding="utf-8")) + for node in ast.walk(tree): + if isinstance(node, ast.ImportFrom) and node.module == _RAW_CLIENT_MODULE: + return True + if isinstance(node, ast.Import) and any(alias.name == _RAW_CLIENT_MODULE for alias in node.names): + return True + return False diff --git a/services/intake/tests/test_clickhouse_executor.py b/services/intake/tests/test_clickhouse_executor.py new file mode 100644 index 0000000000..4e71ba6d87 --- /dev/null +++ b/services/intake/tests/test_clickhouse_executor.py @@ -0,0 +1,92 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the typed ClickHouse repository executor.""" + +from types import SimpleNamespace +from typing import Any, cast + +import pytest +from clickhouse_connect.driver.exceptions import ClickHouseError +from nmp.intake.repository.clickhouse.executor import ClickHouseExecutor, ClickHouseQuery, ClickHouseQueryError +from nmp.intake.repository.clickhouse.tables import ClickHouseTable, qualified_table +from nmp.intake.spans.clickhouse_client import ClickHouseSpanClient + + +class _Client: + database = "intake" + + def __init__(self, *, error: ClickHouseError | None = None) -> None: + self.error = error + self.statements: list[str] = [] + self.parameters: list[dict[str, Any]] = [] + + async def query(self, statement: str, *, parameters: dict[str, Any]) -> SimpleNamespace: + self.statements.append(statement) + self.parameters.append(parameters) + if self.error is not None: + raise self.error + return SimpleNamespace( + column_names=["session_id", "span_count"], + result_rows=[("session-a", 3)], + ) + + +def _executor(client: _Client) -> ClickHouseExecutor: + return ClickHouseExecutor(cast(ClickHouseSpanClient, client)) + + +def test_query_bind_preserves_base_parameters() -> None: + query = ClickHouseQuery( + name="sessions.get", + statement="SELECT %(session_id)s", + parameters={"metric_key": "cost.total"}, + ) + + bound = query.bind(session_id="session-a") + + assert query.parameters == {"metric_key": "cost.total"} + assert bound.parameters == { + "metric_key": "cost.total", + "session_id": "session-a", + } + + +@pytest.mark.asyncio +async def test_executor_returns_mapped_rows_and_bound_parameters() -> None: + client = _Client() + query = ClickHouseQuery( + name="sessions.get", + statement="SELECT %(session_id)s", + parameters={"session_id": "session-a"}, + ) + + rows = await _executor(client).fetch_all(query) + + assert rows == [{"session_id": "session-a", "span_count": 3}] + assert client.statements == [query.statement] + assert client.parameters == [{"session_id": "session-a"}] + + +@pytest.mark.asyncio +async def test_executor_translates_clickhouse_errors_without_sql_details() -> None: + query = ClickHouseQuery( + name="sessions.get", + statement="SELECT secret FROM hidden", + ) + + with pytest.raises(ClickHouseQueryError) as exc_info: + await _executor(_Client(error=ClickHouseError("driver details"))).fetch_all(query) + + assert exc_info.value.query_name == "sessions.get" + assert str(exc_info.value) == "ClickHouse query failed: sessions.get" + assert query.statement not in str(exc_info.value) + + +def test_table_registry_quotes_known_tables() -> None: + assert qualified_table("intake", ClickHouseTable.SPANS) == "`intake`.`spans`" + + +def test_table_registry_rejects_unregistered_names() -> None: + with pytest.raises(TypeError, match="Expected ClickHouseTable"): + qualified_table("intake", cast(ClickHouseTable, "system.tables")) diff --git a/services/intake/tests/test_sessions_clickhouse_repository.py b/services/intake/tests/test_sessions_clickhouse_repository.py index 6a1361c46e..89520a91bf 100644 --- a/services/intake/tests/test_sessions_clickhouse_repository.py +++ b/services/intake/tests/test_sessions_clickhouse_repository.py @@ -4,53 +4,47 @@ """Session repository tests.""" from datetime import datetime, timedelta, timezone -from typing import cast import pytest -from nmp.intake.spans.clickhouse_client import ClickHouseSpanClient -from nmp.intake.spans.session_repository import SESSION_COLUMNS, SessionRepository, session_detail_sql +from nmp.intake.repository.clickhouse.executor import ClickHouseExecutor, ClickHouseQuery +from nmp.intake.repository.clickhouse.session import ClickHouseSessionRepository, _session_detail_query +from nmp.intake.repository.clickhouse.tables import ClickHouseTable -class _QueryResult: - def __init__(self, rows: list[tuple[object, ...]], columns: list[str] | None = None) -> None: - self.result_rows = rows - self.column_names = columns or [] +class _Executor(ClickHouseExecutor): + def __init__(self, rows: list[dict[str, object]] | None = None) -> None: + self.rows = rows or [] + self.queries: list[ClickHouseQuery] = [] + def table(self, table: ClickHouseTable) -> str: + assert table is ClickHouseTable.SPANS + return "spans" -class _Client: - def __init__(self, query_result: _QueryResult | None = None) -> None: - self.query_result = query_result or _QueryResult([]) - self.queries: list[str] = [] - self.parameters: list[dict[str, object]] = [] - - def table(self, name: str) -> str: - return name - - async def query(self, query: str, *, parameters: dict[str, object]) -> _QueryResult: + async def fetch_all(self, query: ClickHouseQuery) -> list[dict[str, object]]: self.queries.append(query) - self.parameters.append(parameters) - return self.query_result + return self.rows -def _repository(client: _Client) -> SessionRepository: - return SessionRepository(cast(ClickHouseSpanClient, client)) +def _repository(executor: _Executor) -> ClickHouseSessionRepository: + return ClickHouseSessionRepository(executor) def test_session_query_is_primary_key_pruned_and_payload_free() -> None: - query, parameters = session_detail_sql("spans") - - assert "FROM spans AS session_spans FINAL" in query - assert "PREWHERE" in query - assert "session_spans.workspace = %(workspace)s" in query - assert "session_spans.session_id = %(session_id)s" in query - assert "trace_index" not in query - assert "JOIN" not in query - assert "session_spans.input" not in query - assert "session_spans.output" not in query - assert "uniqExact(session_spans.source_format, session_spans.trace_id) AS trace_count" in query - assert "count() AS span_count" in query - assert parameters["input_tokens_key"] == "llm.token_count.prompt" - assert parameters["cost_usd_key"] == "cost.total" + query = _session_detail_query("spans") + + assert query.name == "sessions.get" + assert "FROM spans AS session_spans FINAL" in query.statement + assert "PREWHERE" in query.statement + assert "session_spans.workspace = %(workspace)s" in query.statement + assert "session_spans.session_id = %(session_id)s" in query.statement + assert "trace_index" not in query.statement + assert "JOIN" not in query.statement + assert "session_spans.input" not in query.statement + assert "session_spans.output" not in query.statement + assert "uniqExact(session_spans.source_format, session_spans.trace_id) AS trace_count" in query.statement + assert "count() AS span_count" in query.statement + assert query.parameters["input_tokens_key"] == "llm.token_count.prompt" + assert query.parameters["cost_usd_key"] == "cost.total" @pytest.mark.asyncio @@ -73,10 +67,9 @@ async def test_get_session_maps_aggregate_row() -> None: "trace_count": 2, "span_count": 5, } - row = tuple(values[column] for column in SESSION_COLUMNS) - client = _Client(_QueryResult([row], SESSION_COLUMNS)) + executor = _Executor([values]) - session = await _repository(client).get_session(workspace="workspace-a", session_id="session-a") + session = await _repository(executor).get_session(workspace="workspace-a", session_id="session-a") assert session is not None assert session.id == "session-a" @@ -89,12 +82,20 @@ async def test_get_session_maps_aggregate_row() -> None: assert session.span_count == 5 assert session.total_tokens == 858 assert session.cost_usd == 0.0061 - assert client.parameters[0]["workspace"] == "workspace-a" - assert client.parameters[0]["session_id"] == "session-a" + assert executor.queries[0].parameters["workspace"] == "workspace-a" + assert executor.queries[0].parameters["session_id"] == "session-a" @pytest.mark.asyncio async def test_get_session_returns_none_when_no_current_spans_exist() -> None: - session = await _repository(_Client()).get_session(workspace="workspace-a", session_id="missing") + workspace = "workspace' OR 1 = 1 --" + session_id = "session'); DROP TABLE spans; --" + executor = _Executor() + + session = await _repository(executor).get_session(workspace=workspace, session_id=session_id) assert session is None + assert workspace not in executor.queries[0].statement + assert session_id not in executor.queries[0].statement + assert executor.queries[0].parameters["workspace"] == workspace + assert executor.queries[0].parameters["session_id"] == session_id From 6a62a334d4b32f969f277c9df107d33d17bae15b Mon Sep 17 00:00:00 2001 From: Brian Newsom Date: Mon, 27 Jul 2026 15:16:51 -0600 Subject: [PATCH 3/7] refactor(intake): migrate evaluation rollup queries Signed-off-by: Brian Newsom --- .../intake/api/v2/experiments/endpoints.py | 10 +- .../clickhouse/evaluation_rollup.py} | 70 +++---- .../intake/repository/evaluation_rollup.py | 49 +++++ .../tests/test_clickhouse_architecture.py | 1 - .../test_experiment_rollup_repository.py | 190 ++++++++---------- 5 files changed, 159 insertions(+), 161 deletions(-) rename services/intake/src/nmp/intake/{spans/evaluation_rollup_repository.py => repository/clickhouse/evaluation_rollup.py} (91%) create mode 100644 services/intake/src/nmp/intake/repository/evaluation_rollup.py diff --git a/services/intake/src/nmp/intake/api/v2/experiments/endpoints.py b/services/intake/src/nmp/intake/api/v2/experiments/endpoints.py index 31cc365700..fdc1c391b2 100644 --- a/services/intake/src/nmp/intake/api/v2/experiments/endpoints.py +++ b/services/intake/src/nmp/intake/api/v2/experiments/endpoints.py @@ -43,14 +43,12 @@ # layer uses; only the entity's own field names (e.g. parent_experiment_id) reference Experiment directly. from nmp.intake.entities.experiments import Experiment as Evaluation from nmp.intake.entities.experiments import ExperimentGroup +from nmp.intake.repository.clickhouse.evaluation_rollup import ClickHouseEvaluationRollupRepository +from nmp.intake.repository.clickhouse.executor import ClickHouseExecutor +from nmp.intake.repository.evaluation_rollup import EvaluationRollup, EvaluationRollupRepository, ScoreRollup from nmp.intake.spans.api.dependencies import require_workspace_access, validate_list_query_params from nmp.intake.spans.clickhouse_client import ClickHouseSpanClient from nmp.intake.spans.domain import SpanStatus -from nmp.intake.spans.evaluation_rollup_repository import ( - EvaluationRollup, - EvaluationRollupRepository, - ScoreRollup, -) from nmp.intake.spans.evaluation_session_repository import EvaluationSessionRepository, MetricSortTooLargeError from nmp.intake.spans.storage import make_pagination @@ -111,7 +109,7 @@ def get_evaluation_rollup_repository(request: Request) -> EvaluationRollupReposi # Rollups are enrichment only. Evaluation entity reads should continue when # ClickHouse is disabled or temporarily unavailable. client = _get_clickhouse_client(request) - return EvaluationRollupRepository(client) if client is not None else None + return ClickHouseEvaluationRollupRepository(ClickHouseExecutor(client)) if client is not None else None EvaluationRollupRepositoryDep = Annotated[EvaluationRollupRepository | None, Depends(get_evaluation_rollup_repository)] diff --git a/services/intake/src/nmp/intake/spans/evaluation_rollup_repository.py b/services/intake/src/nmp/intake/repository/clickhouse/evaluation_rollup.py similarity index 91% rename from services/intake/src/nmp/intake/spans/evaluation_rollup_repository.py rename to services/intake/src/nmp/intake/repository/clickhouse/evaluation_rollup.py index e82e5fbc89..f13018c775 100644 --- a/services/intake/src/nmp/intake/spans/evaluation_rollup_repository.py +++ b/services/intake/src/nmp/intake/repository/clickhouse/evaluation_rollup.py @@ -1,49 +1,22 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""ClickHouse rollups for Evaluation read models.""" +"""ClickHouse implementation of Evaluation rollup reads.""" from __future__ import annotations -from dataclasses import dataclass, field from typing import Any -from nmp.intake.spans.clickhouse_client import ClickHouseSpanClient +from nmp.intake.repository.clickhouse.executor import ClickHouseExecutor, ClickHouseQuery +from nmp.intake.repository.clickhouse.tables import ClickHouseTable +from nmp.intake.repository.evaluation_rollup import EvaluationRollup, EvaluationRollupRepository, ScoreRollup from nmp.intake.spans.span_attribute_catalog import COST_SCALE, SpanAttributeField, spec_for_field -from nmp.intake.spans.storage import float_or_none, result_rows +from nmp.intake.spans.storage import float_or_none -@dataclass(frozen=True) -class ScoreRollup: - sum: float | None - mean: float | None - median: float | None - p90: float | None - p95: float | None - p99: float | None - count: int - - -@dataclass -class EvaluationRollup: - evaluation_id: str - run_count: int = 0 - test_case_count: int = 0 - model_names: list[str] = field(default_factory=list) - agent_names: list[str] = field(default_factory=list) - agent_versions: list[str] = field(default_factory=list) - evaluator_scores: dict[str, ScoreRollup] = field(default_factory=dict) - cost_usd: ScoreRollup | None = None - latency_ms: ScoreRollup | None = None - - @property - def evaluator_names(self) -> list[str]: - return sorted(self.evaluator_scores) - - -class EvaluationRollupRepository: - def __init__(self, client: ClickHouseSpanClient) -> None: - self._client = client +class ClickHouseEvaluationRollupRepository(EvaluationRollupRepository): + def __init__(self, executor: ClickHouseExecutor) -> None: + self._executor = executor async def get_rollups(self, *, workspace: str, evaluation_ids: list[str]) -> dict[str, EvaluationRollup]: evaluation_ids = list(dict.fromkeys(evaluation_ids)) @@ -53,22 +26,24 @@ async def get_rollups(self, *, workspace: str, evaluation_ids: list[str]) -> dic evaluation_names_sql, evaluation_parameters = _evaluation_id_parameters(evaluation_ids) parameters = {"workspace": workspace, **evaluation_parameters} - trace_index_table = self._client.table("trace_index") + trace_index_table = self._executor.table(ClickHouseTable.TRACE_INDEX) - for row in result_rows( - await self._client.query( - _run_counts_sql(trace_index_table, evaluation_names_sql), + for row in await self._executor.fetch_all( + ClickHouseQuery( + name="evaluation_rollups.run_counts", + statement=_run_counts_sql(trace_index_table, evaluation_names_sql), parameters=parameters, ) ): rollups[row["evaluation_id"]].run_count = int(row["run_count"]) rollups[row["evaluation_id"]].test_case_count = int(row["test_case_count"]) - for row in result_rows( - await self._client.query( - _score_rollups_sql( + for row in await self._executor.fetch_all( + ClickHouseQuery( + name="evaluation_rollups.scores", + statement=_score_rollups_sql( trace_index_table=trace_index_table, - evaluator_results_table=self._client.table("evaluator_results"), + evaluator_results_table=self._executor.table(ClickHouseTable.EVALUATOR_RESULTS), evaluation_names_sql=evaluation_names_sql, ), parameters=parameters, @@ -84,11 +59,12 @@ async def get_rollups(self, *, workspace: str, evaluation_ids: list[str]) -> dic count=int(row["count"]), ) - for row in result_rows( - await self._client.query( - _metric_rollups_sql( + for row in await self._executor.fetch_all( + ClickHouseQuery( + name="evaluation_rollups.metrics", + statement=_metric_rollups_sql( trace_index_table=trace_index_table, - spans_table=self._client.table("spans"), + spans_table=self._executor.table(ClickHouseTable.SPANS), evaluation_names_sql=evaluation_names_sql, ), parameters={ diff --git a/services/intake/src/nmp/intake/repository/evaluation_rollup.py b/services/intake/src/nmp/intake/repository/evaluation_rollup.py new file mode 100644 index 0000000000..e2782f7781 --- /dev/null +++ b/services/intake/src/nmp/intake/repository/evaluation_rollup.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Repository interface and read models for Evaluation rollups.""" + +from abc import ABC, abstractmethod +from dataclasses import dataclass, field + + +@dataclass(frozen=True) +class ScoreRollup: + sum: float | None + mean: float | None + median: float | None + p90: float | None + p95: float | None + p99: float | None + count: int + + +@dataclass +class EvaluationRollup: + evaluation_id: str + run_count: int = 0 + test_case_count: int = 0 + model_names: list[str] = field(default_factory=list) + agent_names: list[str] = field(default_factory=list) + agent_versions: list[str] = field(default_factory=list) + evaluator_scores: dict[str, ScoreRollup] = field(default_factory=dict) + cost_usd: ScoreRollup | None = None + latency_ms: ScoreRollup | None = None + + @property + def evaluator_names(self) -> list[str]: + return sorted(self.evaluator_scores) + + +class EvaluationRollupRepository(ABC): + """Domain-facing interface for Evaluation rollup reads.""" + + @abstractmethod + async def get_rollups( + self, + *, + workspace: str, + evaluation_ids: list[str], + ) -> dict[str, EvaluationRollup]: + """Return rollups keyed by Evaluation ID.""" + pass diff --git a/services/intake/tests/test_clickhouse_architecture.py b/services/intake/tests/test_clickhouse_architecture.py index 4d23dff919..2e53a49a17 100644 --- a/services/intake/tests/test_clickhouse_architecture.py +++ b/services/intake/tests/test_clickhouse_architecture.py @@ -16,7 +16,6 @@ "service.py", "spans/annotations_repository.py", "spans/api/dependencies.py", - "spans/evaluation_rollup_repository.py", "spans/evaluation_session_repository.py", "spans/evaluator_results_repository.py", "spans/span_repository.py", diff --git a/services/intake/tests/test_experiment_rollup_repository.py b/services/intake/tests/test_experiment_rollup_repository.py index d4c202e54b..f95b733784 100644 --- a/services/intake/tests/test_experiment_rollup_repository.py +++ b/services/intake/tests/test_experiment_rollup_repository.py @@ -3,97 +3,72 @@ """Evaluation rollup repository tests.""" -from typing import cast - import pytest -from nmp.intake.spans.clickhouse_client import ClickHouseSpanClient -from nmp.intake.spans.evaluation_rollup_repository import EvaluationRollupRepository - - -class _QueryResult: - def __init__(self, rows: list[tuple[object, ...]], columns: list[str]) -> None: - self.result_rows = rows - self.column_names = columns +from nmp.intake.repository.clickhouse.evaluation_rollup import ClickHouseEvaluationRollupRepository +from nmp.intake.repository.clickhouse.executor import ClickHouseExecutor, ClickHouseQuery +from nmp.intake.repository.clickhouse.tables import ClickHouseTable -class _Client: - def __init__(self, query_results: list[_QueryResult]) -> None: - self.queries: list[str] = [] - self.parameters: list[dict[str, object]] = [] +class _Executor(ClickHouseExecutor): + def __init__(self, query_results: list[list[dict[str, object]]]) -> None: + self.queries: list[ClickHouseQuery] = [] self.query_results = query_results - def table(self, name: str) -> str: - return name + def table(self, table: ClickHouseTable) -> str: + return table.value - async def query(self, query: str, *, parameters: dict[str, object]) -> _QueryResult: + async def fetch_all(self, query: ClickHouseQuery) -> list[dict[str, object]]: self.queries.append(query) - self.parameters.append(parameters) return self.query_results.pop(0) -def _repository(client: _Client) -> EvaluationRollupRepository: - return EvaluationRollupRepository(cast(ClickHouseSpanClient, client)) +def _repository(executor: _Executor) -> ClickHouseEvaluationRollupRepository: + return ClickHouseEvaluationRollupRepository(executor) @pytest.mark.asyncio async def test_evaluation_rollups_anchor_on_root_session_membership(): - client = _Client( + executor = _Executor( [ - _QueryResult( - [("exp-a", 3, 2)], - ["evaluation_id", "run_count", "test_case_count"], - ), - _QueryResult( - [("exp-a", "reward", 3.0, 0.75, 0.8, 1.0, 1.0, 1.0, 4)], - ["evaluation_id", "evaluator_name", "sum", "mean", "median", "p90", "p95", "p99", "count"], - ), - _QueryResult( - [ - ( - "exp-a", - ["model-b", "model-a"], - ["agent-a"], - ["1.0.0", "1.0.1"], - 0.65, - 0.1625, - 0.2, - 0.3, - 0.3, - 0.3, - 4, - 7000.0, - 1750.0, - 2000.0, - 3000.0, - 3000.0, - 3000.0, - 4, - ) - ], - [ - "evaluation_id", - "model_names", - "agent_names", - "agent_versions", - "cost_sum", - "cost_mean", - "cost_median", - "cost_p90", - "cost_p95", - "cost_p99", - "cost_count", - "latency_sum", - "latency_mean", - "latency_median", - "latency_p90", - "latency_p95", - "latency_p99", - "latency_count", - ], - ), + [{"evaluation_id": "exp-a", "run_count": 3, "test_case_count": 2}], + [ + { + "evaluation_id": "exp-a", + "evaluator_name": "reward", + "sum": 3.0, + "mean": 0.75, + "median": 0.8, + "p90": 1.0, + "p95": 1.0, + "p99": 1.0, + "count": 4, + } + ], + [ + { + "evaluation_id": "exp-a", + "model_names": ["model-b", "model-a"], + "agent_names": ["agent-a"], + "agent_versions": ["1.0.0", "1.0.1"], + "cost_sum": 0.65, + "cost_mean": 0.1625, + "cost_median": 0.2, + "cost_p90": 0.3, + "cost_p95": 0.3, + "cost_p99": 0.3, + "cost_count": 4, + "latency_sum": 7000.0, + "latency_mean": 1750.0, + "latency_median": 2000.0, + "latency_p90": 3000.0, + "latency_p95": 3000.0, + "latency_p99": 3000.0, + "latency_count": 4, + } + ], ] ) - repository = _repository(client) + repository = _repository(executor) rollups = await repository.get_rollups(workspace="default", evaluation_ids=["exp-a"]) @@ -128,44 +103,45 @@ async def test_evaluation_rollups_anchor_on_root_session_membership(): assert rollup.latency_ms.p99 == 3000 assert rollup.latency_ms.count == 4 - assert len(client.queries) == 3 - assert "FROM trace_index FINAL" in client.queries[0] - assert "count() AS run_count" in client.queries[0] - assert "uniqExactIf(test_case_id, test_case_id != '')" in client.queries[0] - assert "AS test_case_count" in client.queries[0] - assert "evaluation_id IN (%(evaluation_id_0)s)" in client.queries[0] - assert "ORDER BY root_started_at ASC, root_span_id ASC" in client.queries[0] - assert "FROM evaluator_results FINAL" in client.queries[1] - assert "quantileExact(0.5)(value) AS median" in client.queries[1] - assert "quantileExact(0.99)(value) AS p99" in client.queries[1] - assert "AND (workspace, session_id) IN (" in client.queries[1] - assert "sessions.session_id = results.session_id" in client.queries[1] + assert [query.name for query in executor.queries] == [ + "evaluation_rollups.run_counts", + "evaluation_rollups.scores", + "evaluation_rollups.metrics", + ] + statements = [query.statement for query in executor.queries] + assert "FROM trace_index FINAL" in statements[0] + assert "count() AS run_count" in statements[0] + assert "uniqExactIf(test_case_id, test_case_id != '')" in statements[0] + assert "AS test_case_count" in statements[0] + assert "evaluation_id IN (%(evaluation_id_0)s)" in statements[0] + assert "ORDER BY root_started_at ASC, root_span_id ASC" in statements[0] + assert "FROM evaluator_results FINAL" in statements[1] + assert "quantileExact(0.5)(value) AS median" in statements[1] + assert "quantileExact(0.99)(value) AS p99" in statements[1] + assert "AND (workspace, session_id) IN (" in statements[1] + assert "sessions.session_id = results.session_id" in statements[1] # Scores are reduced to one value per (session, evaluator), then averaged per test case before the # distribution rollup, so the mean is test-case-weighted and count tracks test cases. - assert ( - "GROUP BY sessions.evaluation_id, sessions.session_id, sessions.test_case_id, results.name" in client.queries[1] - ) - assert "test_case_scores AS" in client.queries[1] - assert "WHERE sessions.test_case_id != ''" in client.queries[1] - assert "test_case_metrics AS" in client.queries[2] - assert "current_session_spans AS" in client.queries[2] - assert ( - "(workspace, session_id) IN (SELECT DISTINCT workspace, session_id FROM scoped_sessions)" in client.queries[2] - ) - assert "LEFT JOIN current_session_spans AS spans" in client.queries[2] - assert "sessions.session_id = spans.dedup_session_id" in client.queries[2] - assert "arraySort(arrayDistinct(arrayFlatten(groupArray(model_names)))) AS model_names" in client.queries[2] - assert "quantileExactIf(0.5)" in client.queries[2] - assert "cost_median" in client.queries[2] - assert "quantileExactIf(0.99)" in client.queries[2] - assert "latency_p99" in client.queries[2] - assert "sessions.trace_id = spans.trace_id" not in client.queries[2] - assert client.parameters[0]["evaluation_id_0"] == "exp-a" - assert client.parameters[2]["model_key"] == "gen_ai.request.model" + assert "GROUP BY sessions.evaluation_id, sessions.session_id, sessions.test_case_id, results.name" in statements[1] + assert "test_case_scores AS" in statements[1] + assert "WHERE sessions.test_case_id != ''" in statements[1] + assert "test_case_metrics AS" in statements[2] + assert "current_session_spans AS" in statements[2] + assert "(workspace, session_id) IN (SELECT DISTINCT workspace, session_id FROM scoped_sessions)" in statements[2] + assert "LEFT JOIN current_session_spans AS spans" in statements[2] + assert "sessions.session_id = spans.dedup_session_id" in statements[2] + assert "arraySort(arrayDistinct(arrayFlatten(groupArray(model_names)))) AS model_names" in statements[2] + assert "quantileExactIf(0.5)" in statements[2] + assert "cost_median" in statements[2] + assert "quantileExactIf(0.99)" in statements[2] + assert "latency_p99" in statements[2] + assert "sessions.trace_id = spans.trace_id" not in statements[2] + assert executor.queries[0].parameters["evaluation_id_0"] == "exp-a" + assert executor.queries[2].parameters["model_key"] == "gen_ai.request.model" def test_score_rollup_cte_builders_compose_the_pipeline(): - from nmp.intake.spans.evaluation_rollup_repository import ( + from nmp.intake.repository.clickhouse.evaluation_rollup import ( _evaluators_cte, _session_scores_cte, _test_case_scores_cte, From c2dcecbf717fd58be5474e39c8fe23f453c24048 Mon Sep 17 00:00:00 2001 From: Brian Newsom Date: Mon, 27 Jul 2026 15:41:31 -0600 Subject: [PATCH 4/7] refactor(intake): migrate evaluation session queries Signed-off-by: Brian Newsom --- .../intake/api/v2/experiments/endpoints.py | 7 +- .../nmp/intake/api/v2/experiments/schemas.py | 2 +- .../clickhouse/evaluation_session.py} | 132 +++---- .../intake/repository/clickhouse/executor.py | 6 + .../intake/repository/evaluation_session.py | 68 ++++ .../spans/test_experiment_sessions.py | 28 +- .../tests/test_clickhouse_architecture.py | 1 - .../intake/tests/test_clickhouse_executor.py | 12 + ...valuation_session_clickhouse_repository.py | 363 ++++++++---------- .../tests/test_experiment_session_schemas.py | 2 +- 10 files changed, 342 insertions(+), 279 deletions(-) rename services/intake/src/nmp/intake/{spans/evaluation_session_repository.py => repository/clickhouse/evaluation_session.py} (89%) create mode 100644 services/intake/src/nmp/intake/repository/evaluation_session.py diff --git a/services/intake/src/nmp/intake/api/v2/experiments/endpoints.py b/services/intake/src/nmp/intake/api/v2/experiments/endpoints.py index fdc1c391b2..966fa1c088 100644 --- a/services/intake/src/nmp/intake/api/v2/experiments/endpoints.py +++ b/services/intake/src/nmp/intake/api/v2/experiments/endpoints.py @@ -44,12 +44,13 @@ from nmp.intake.entities.experiments import Experiment as Evaluation from nmp.intake.entities.experiments import ExperimentGroup from nmp.intake.repository.clickhouse.evaluation_rollup import ClickHouseEvaluationRollupRepository +from nmp.intake.repository.clickhouse.evaluation_session import ClickHouseEvaluationSessionRepository from nmp.intake.repository.clickhouse.executor import ClickHouseExecutor from nmp.intake.repository.evaluation_rollup import EvaluationRollup, EvaluationRollupRepository, ScoreRollup +from nmp.intake.repository.evaluation_session import EvaluationSessionRepository, MetricSortTooLargeError from nmp.intake.spans.api.dependencies import require_workspace_access, validate_list_query_params from nmp.intake.spans.clickhouse_client import ClickHouseSpanClient from nmp.intake.spans.domain import SpanStatus -from nmp.intake.spans.evaluation_session_repository import EvaluationSessionRepository, MetricSortTooLargeError from nmp.intake.spans.storage import make_pagination logger = logging.getLogger(__name__) @@ -72,7 +73,7 @@ def _sanitize_for_log(value: str) -> str: _ENTITY_SORT_FIELDS = frozenset({"name", "created_at", "updated_at", "pinned_at"}) # Sessions are sorted in ClickHouse (ORDER BY before LIMIT/OFFSET) so sort composes # correctly with pagination. These are the allowed field names; each maps to an SQL -# expression in the repository - see _list_sql in evaluation_session_repository.py. +# expression in the ClickHouse repository. _SESSION_SORT_FIELDS = frozenset( { "test_case_id", @@ -117,7 +118,7 @@ def get_evaluation_rollup_repository(request: Request) -> EvaluationRollupReposi def get_evaluation_session_repository(request: Request) -> EvaluationSessionRepository | None: client = _get_clickhouse_client(request) - return EvaluationSessionRepository(client) if client is not None else None + return ClickHouseEvaluationSessionRepository(ClickHouseExecutor(client)) if client is not None else None EvaluationSessionRepositoryDep = Annotated[ diff --git a/services/intake/src/nmp/intake/api/v2/experiments/schemas.py b/services/intake/src/nmp/intake/api/v2/experiments/schemas.py index 5cf110b019..7f7b20e621 100644 --- a/services/intake/src/nmp/intake/api/v2/experiments/schemas.py +++ b/services/intake/src/nmp/intake/api/v2/experiments/schemas.py @@ -14,12 +14,12 @@ from nmp.common.entities.values import DatetimeFilter, Filter, NumberFilter, map_entity_field from nmp.intake.entities.experiments import Experiment, ExperimentGroup, ParetoConfig +from nmp.intake.repository.evaluation_session import EvaluationSessionRow from nmp.intake.spans.domain import ( INTAKE_PREVIEW_PAYLOAD_CHAR_LIMIT, IntakeResponseMode, SpanStatus, ) -from nmp.intake.spans.evaluation_session_repository import EvaluationSessionRow from nmp.intake.spans.storage import text_for_mode from pydantic import AnyUrl, BaseModel, ConfigDict, Field, computed_field, model_validator diff --git a/services/intake/src/nmp/intake/spans/evaluation_session_repository.py b/services/intake/src/nmp/intake/repository/clickhouse/evaluation_session.py similarity index 89% rename from services/intake/src/nmp/intake/spans/evaluation_session_repository.py rename to services/intake/src/nmp/intake/repository/clickhouse/evaluation_session.py index 839072b031..8427e6f0a8 100644 --- a/services/intake/src/nmp/intake/spans/evaluation_session_repository.py +++ b/services/intake/src/nmp/intake/repository/clickhouse/evaluation_session.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""ClickHouse repository for per-session rows of an Evaluation. +"""ClickHouse implementation of per-session Evaluation reads. Returns one row per ingested session (test case execution), using ``trace_index`` for root/session membership and per-session aggregates from all spans (tokens + @@ -10,18 +10,22 @@ from __future__ import annotations -from dataclasses import dataclass, field -from datetime import datetime from typing import Any -from nmp.intake.spans.clickhouse_client import ClickHouseSpanClient +from nmp.intake.repository.clickhouse.executor import ClickHouseExecutor, ClickHouseQuery +from nmp.intake.repository.clickhouse.tables import ClickHouseTable +from nmp.intake.repository.evaluation_session import ( + EvaluationSessionPage, + EvaluationSessionRepository, + EvaluationSessionRow, + MetricSortTooLargeError, +) from nmp.intake.spans.domain import IntakeResponseMode, SpanStatus from nmp.intake.spans.span_attribute_catalog import COST_SCALE, SpanAttributeField, spec_for_field from nmp.intake.spans.storage import ( float_or_none, int_or_none, normalize_span_status, - result_rows, str_or_none, text_query_parameters, text_select_for_mode, @@ -77,51 +81,9 @@ _MAX_METRIC_SORT_SESSIONS = 10_000 -class MetricSortTooLargeError(Exception): - """Raised when a cost/tokens sort is requested on more sessions than the pre-metrics cap allows. - - This is a domain exception (not HTTPException) so the repository stays HTTP-agnostic. - The endpoint catches it and converts it to 413. - """ - - def __init__(self, total: int, limit: int) -> None: - self.total = total - self.limit = limit - super().__init__(f"Metric sort requested on {total} sessions, limit is {limit}") - - -@dataclass(frozen=True) -class EvaluationSessionRow: - """One ingested session of an Evaluation.""" - - workspace: str - evaluation_name: str - session_id: str - test_case_id: str | None - trace_id: str - root_span_id: str - started_at: datetime - ended_at: datetime | None - latency_ms: float | None - status: SpanStatus - input: str | None - output: str | None - input_tokens: int | None - output_tokens: int | None - cached_tokens: int | None - cost_total_usd: float | None - evaluator_scores: dict[str, float] = field(default_factory=dict) - - -@dataclass(frozen=True) -class EvaluationSessionPage: - rows: list[EvaluationSessionRow] - total: int - - -class EvaluationSessionRepository: - def __init__(self, client: ClickHouseSpanClient) -> None: - self._client = client +class ClickHouseEvaluationSessionRepository(EvaluationSessionRepository): + def __init__(self, executor: ClickHouseExecutor) -> None: + self._executor = executor async def list_sessions( self, @@ -135,9 +97,9 @@ async def list_sessions( mode: IntakeResponseMode, sort_keys: list[tuple[str, bool]] | None = None, ) -> EvaluationSessionPage: - trace_index_table = self._client.table("trace_index") - spans_table = self._client.table("spans") - evaluator_results_table = self._client.table("evaluator_results") + trace_index_table = self._executor.table(ClickHouseTable.TRACE_INDEX) + spans_table = self._executor.table(ClickHouseTable.SPANS) + evaluator_results_table = self._executor.table(ClickHouseTable.EVALUATOR_RESULTS) scoped_filter_sql, scoped_filter_parameters = _scoped_filter(test_case_id=test_case_id, status=status) @@ -154,11 +116,14 @@ async def list_sessions( trace_index_table=trace_index_table, scoped_filter_sql=scoped_filter_sql, ) - count_result = await self._client.query( - count_sql, - parameters={**base_parameters, **scoped_filter_parameters}, + count_value = await self._executor.fetch_scalar( + ClickHouseQuery( + name="evaluation_sessions.count", + statement=count_sql, + parameters={**base_parameters, **scoped_filter_parameters}, + ) ) - total = int(count_result.result_rows[0][0]) if count_result.result_rows else 0 + total = int(count_value) if count_value is not None else 0 if total == 0: return EvaluationSessionPage(rows=[], total=0) @@ -169,6 +134,7 @@ async def list_sessions( offset = (page - 1) * page_size if needs_pre_metrics: + assert sort_keys is not None # Two-query path for cost/tokens sorts. # # ClickHouse inlines regular CTEs (does not materialise them), so a single query that @@ -183,18 +149,21 @@ async def list_sessions( trace_index_table=trace_index_table, spans_table=spans_table, scoped_filter_sql=scoped_filter_sql, - sort_keys=sort_keys, # type: ignore[arg-type] # guaranteed non-None here + sort_keys=sort_keys, ) - page_ids_result = await self._client.query( - page_ids_sql, - parameters={ - **base_parameters, - **scoped_filter_parameters, - "limit": page_size, - "offset": offset, - }, + page_id_rows = await self._executor.fetch_all( + ClickHouseQuery( + name="evaluation_sessions.metric_sort.page_ids", + statement=page_ids_sql, + parameters={ + **base_parameters, + **scoped_filter_parameters, + "limit": page_size, + "offset": offset, + }, + ) ) - ordered_ids = [record["session_id"] for record in result_rows(page_ids_result)] + ordered_ids = [record["session_id"] for record in page_id_rows] if not ordered_ids: return EvaluationSessionPage(rows=[], total=total) @@ -208,15 +177,18 @@ async def list_sessions( evaluator_results_table=evaluator_results_table, mode=mode, ) - hydrate_result = await self._client.query( - hydrate_sql, - parameters={ - **base_parameters, - **text_query_parameters(mode), - "session_ids": ordered_ids, - }, + hydrate_rows = await self._executor.fetch_all( + ClickHouseQuery( + name="evaluation_sessions.metric_sort.hydrate", + statement=hydrate_sql, + parameters={ + **base_parameters, + **text_query_parameters(mode), + "session_ids": ordered_ids, + }, + ) ) - rows_by_id = {record["session_id"]: _row(record) for record in result_rows(hydrate_result)} + rows_by_id = {record["session_id"]: _row(record) for record in hydrate_rows} # Restore the order from Query 2. `rows_by_id` may be missing a session_id if a # race caused the trace_index to disagree between queries, so guard with `if sid in`. rows = [rows_by_id[sid] for sid in ordered_ids if sid in rows_by_id] @@ -239,8 +211,14 @@ async def list_sessions( "limit": page_size, "offset": offset, } - list_result = await self._client.query(list_sql, parameters=list_parameters) - rows = [_row(record) for record in result_rows(list_result)] + list_rows = await self._executor.fetch_all( + ClickHouseQuery( + name="evaluation_sessions.list", + statement=list_sql, + parameters=list_parameters, + ) + ) + rows = [_row(record) for record in list_rows] return EvaluationSessionPage(rows=rows, total=total) diff --git a/services/intake/src/nmp/intake/repository/clickhouse/executor.py b/services/intake/src/nmp/intake/repository/clickhouse/executor.py index c071a5663d..849fc22483 100644 --- a/services/intake/src/nmp/intake/repository/clickhouse/executor.py +++ b/services/intake/src/nmp/intake/repository/clickhouse/executor.py @@ -75,3 +75,9 @@ async def fetch_all(self, query: ClickHouseQuery) -> list[dict[str, Any]]: columns: Sequence[str] = result.column_names rows: Sequence[Sequence[Any]] = result.result_rows return [dict(zip(columns, row, strict=True)) for row in rows] + + async def fetch_scalar(self, query: ClickHouseQuery) -> Any | None: + """Return the first column of the first row, or ``None`` when no row exists.""" + + rows = await self.fetch_all(query) + return next(iter(rows[0].values())) if rows else None diff --git a/services/intake/src/nmp/intake/repository/evaluation_session.py b/services/intake/src/nmp/intake/repository/evaluation_session.py new file mode 100644 index 0000000000..019dccae14 --- /dev/null +++ b/services/intake/src/nmp/intake/repository/evaluation_session.py @@ -0,0 +1,68 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Repository interface and read models for Evaluation sessions.""" + +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from datetime import datetime + +from nmp.intake.spans.domain import IntakeResponseMode, SpanStatus + + +class MetricSortTooLargeError(Exception): + """Raised when a cost/tokens sort exceeds the bounded pre-metrics path.""" + + def __init__(self, total: int, limit: int) -> None: + self.total = total + self.limit = limit + super().__init__(f"Metric sort requested on {total} sessions, limit is {limit}") + + +@dataclass(frozen=True) +class EvaluationSessionRow: + """One ingested session of an Evaluation.""" + + workspace: str + evaluation_name: str + session_id: str + test_case_id: str | None + trace_id: str + root_span_id: str + started_at: datetime + ended_at: datetime | None + latency_ms: float | None + status: SpanStatus + input: str | None + output: str | None + input_tokens: int | None + output_tokens: int | None + cached_tokens: int | None + cost_total_usd: float | None + evaluator_scores: dict[str, float] = field(default_factory=dict) + + +@dataclass(frozen=True) +class EvaluationSessionPage: + rows: list[EvaluationSessionRow] + total: int + + +class EvaluationSessionRepository(ABC): + """Domain-facing interface for Evaluation session reads.""" + + @abstractmethod + async def list_sessions( + self, + *, + workspace: str, + evaluation_name: str, + status: SpanStatus | None = None, + test_case_id: str | None = None, + page: int, + page_size: int, + mode: IntakeResponseMode, + sort_keys: list[tuple[str, bool]] | None = None, + ) -> EvaluationSessionPage: + """Return one page of sessions for an Evaluation.""" + pass diff --git a/services/intake/tests/integration/spans/test_experiment_sessions.py b/services/intake/tests/integration/spans/test_experiment_sessions.py index 46365a06b8..b770360227 100644 --- a/services/intake/tests/integration/spans/test_experiment_sessions.py +++ b/services/intake/tests/integration/spans/test_experiment_sessions.py @@ -116,6 +116,27 @@ def test_list_evaluation_sessions_returns_joined_session_rows(client: TestClient assert paged_body["data"][0]["test_case_id"] == "case-b" assert paged_body["data"][0]["evaluator_scores"] == {"reward": pytest.approx(0.5)} + latency_sorted = client.get( + f"{EVALUATIONS}/{evaluation_name}/sessions", + params={"sort": "-latency_ms", "page_size": 3}, + ) + assert latency_sorted.status_code == 200, latency_sorted.text + assert [row["test_case_id"] for row in latency_sorted.json()["data"]] == ["case-c", "case-b", "case-a"] + + cost_sorted = client.get( + f"{EVALUATIONS}/{evaluation_name}/sessions", + params={"sort": "-cost_total_usd", "page_size": 3}, + ) + assert cost_sorted.status_code == 200, cost_sorted.text + assert [row["test_case_id"] for row in cost_sorted.json()["data"]] == ["case-c", "case-b", "case-a"] + + tokens_sorted = client.get( + f"{EVALUATIONS}/{evaluation_name}/sessions", + params={"sort": "tokens", "page_size": 3}, + ) + assert tokens_sorted.status_code == 200, tokens_sorted.text + assert [row["test_case_id"] for row in tokens_sorted.json()["data"]] == ["case-a", "case-b", "case-c"] + def test_list_evaluation_sessions_filter_by_test_case(client: TestClient) -> None: evaluation_name = "sessions-filter-exp" @@ -132,7 +153,8 @@ def test_list_evaluation_sessions_filter_by_test_case(client: TestClient) -> Non assert created.status_code == 201, created.text started_at = datetime.now(timezone.utc).replace(microsecond=0) - for index, test_case_id in enumerate(["alpha", "beta"]): + adversarial_test_case_id = "alpha') OR 1 = 1 --" + for index, test_case_id in enumerate([adversarial_test_case_id, "beta"]): response = client.post( ATIF_INGEST, json=_atif_body( @@ -152,13 +174,13 @@ def test_list_evaluation_sessions_filter_by_test_case(client: TestClient) -> Non filtered = client.get( f"{EVALUATIONS}/{evaluation_name}/sessions", - params={"filter[test_case_id]": "alpha"}, + params={"filter[test_case_id]": adversarial_test_case_id}, ) assert filtered.status_code == 200, filtered.text body = filtered.json() assert body["pagination"]["total_results"] == 1 assert len(body["data"]) == 1 - assert body["data"][0]["test_case_id"] == "alpha" + assert body["data"][0]["test_case_id"] == adversarial_test_case_id def test_list_evaluation_sessions_filter_by_status(client: TestClient) -> None: diff --git a/services/intake/tests/test_clickhouse_architecture.py b/services/intake/tests/test_clickhouse_architecture.py index 2e53a49a17..6f9f52c0c1 100644 --- a/services/intake/tests/test_clickhouse_architecture.py +++ b/services/intake/tests/test_clickhouse_architecture.py @@ -16,7 +16,6 @@ "service.py", "spans/annotations_repository.py", "spans/api/dependencies.py", - "spans/evaluation_session_repository.py", "spans/evaluator_results_repository.py", "spans/span_repository.py", "spans/trace_repository.py", diff --git a/services/intake/tests/test_clickhouse_executor.py b/services/intake/tests/test_clickhouse_executor.py index 4e71ba6d87..1b2c08f7db 100644 --- a/services/intake/tests/test_clickhouse_executor.py +++ b/services/intake/tests/test_clickhouse_executor.py @@ -68,6 +68,18 @@ async def test_executor_returns_mapped_rows_and_bound_parameters() -> None: assert client.parameters == [{"session_id": "session-a"}] +@pytest.mark.asyncio +async def test_executor_returns_first_scalar() -> None: + scalar = await _executor(_Client()).fetch_scalar( + ClickHouseQuery( + name="sessions.first_id", + statement="SELECT session_id", + ) + ) + + assert scalar == "session-a" + + @pytest.mark.asyncio async def test_executor_translates_clickhouse_errors_without_sql_details() -> None: query = ClickHouseQuery( diff --git a/services/intake/tests/test_evaluation_session_clickhouse_repository.py b/services/intake/tests/test_evaluation_session_clickhouse_repository.py index fd5a4f1a3e..5f7c3d36c3 100644 --- a/services/intake/tests/test_evaluation_session_clickhouse_repository.py +++ b/services/intake/tests/test_evaluation_session_clickhouse_repository.py @@ -1,212 +1,189 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Evaluation session ClickHouse query tests.""" +"""Evaluation session ClickHouse repository tests.""" -from nmp.intake.spans.evaluation_session_repository import ( +from datetime import datetime, timezone +from typing import Any + +import pytest +from nmp.intake.repository.clickhouse.evaluation_session import ( + _MAX_METRIC_SORT_SESSIONS, _SORT_EXPR_FINAL, _SORT_EXPR_PAGE, + ClickHouseEvaluationSessionRepository, _build_order_by, - _count_sql, - _hydrate_by_ids_sql, - _list_sql, - _metric_sort_page_ids_sql, ) - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - -def _make_list_sql(**kwargs) -> str: - """Call _list_sql with fixed table names; only pass what the test cares about.""" - return _list_sql( - trace_index_table="trace_index", - spans_table="spans", - evaluator_results_table="evaluator_results", - scoped_filter_sql="", - **kwargs, +from nmp.intake.repository.clickhouse.executor import ClickHouseExecutor, ClickHouseQuery +from nmp.intake.repository.clickhouse.tables import ClickHouseTable +from nmp.intake.repository.evaluation_session import MetricSortTooLargeError +from nmp.intake.spans.domain import SpanStatus + + +class _Executor(ClickHouseExecutor): + def __init__(self, query_results: list[list[dict[str, Any]]]) -> None: + self.queries: list[ClickHouseQuery] = [] + self.query_results = query_results + self.tables: list[ClickHouseTable] = [] + + def table(self, table: ClickHouseTable) -> str: + self.tables.append(table) + return table.value + + async def fetch_all(self, query: ClickHouseQuery) -> list[dict[str, Any]]: + self.queries.append(query) + return self.query_results.pop(0) + + +def _repository(executor: _Executor) -> ClickHouseEvaluationSessionRepository: + return ClickHouseEvaluationSessionRepository(executor) + + +def _session_record(session_id: str) -> dict[str, Any]: + now = datetime(2026, 1, 1, tzinfo=timezone.utc) + return { + "workspace": "default", + "evaluation_id": "evaluation-a", + "session_id": session_id, + "test_case_id": "case-a", + "trace_id": f"trace-{session_id}", + "root_span_id": f"root-{session_id}", + "start_time": now, + "end_time": now, + "latency_ms": 125.0, + "root_span_status": "success", + "input": "input", + "output": "output", + "input_tokens": 10, + "output_tokens": 5, + "cached_tokens": 2, + "cost_total_usd": 0.01, + "evaluator_scores": {"reward": 0.9}, + } + + +@pytest.mark.asyncio +async def test_list_sessions_maps_rows_and_binds_all_request_values() -> None: + workspace = "workspace' OR 1 = 1 --" + evaluation_name = "evaluation'); DROP TABLE trace_index; --" + test_case_id = "case' UNION ALL SELECT secret --" + executor = _Executor([[{"count()": 1}], [_session_record("session-a")]]) + + page = await _repository(executor).list_sessions( + workspace=workspace, + evaluation_name=evaluation_name, + status=SpanStatus.ERROR, + test_case_id=test_case_id, + page=2, + page_size=5, + mode="preview", + sort_keys=[("latency_ms", True)], ) - -# --------------------------------------------------------------------------- -# Existing payload-mode tests — updated to pass the now-required sort_keys arg -# --------------------------------------------------------------------------- - - -def test_session_count_query_does_not_read_root_payloads() -> None: - query = _count_sql(trace_index_table="trace_index", scoped_filter_sql="") - - assert "root_input" not in query - assert "root_output" not in query - - -def test_session_preview_query_truncates_input_and_output_in_clickhouse() -> None: - query = _make_list_sql(mode="preview", sort_keys=[]) - - assert "substringUTF8(root_input, 1, %(payload_char_limit)s) AS input" in query - assert "substringUTF8(root_output, 1, %(payload_char_limit)s) AS output" in query - assert "root_input AS input" not in query - assert "root_output AS output" not in query - - -def test_session_summary_query_omits_input_and_output_columns() -> None: - query = _make_list_sql(mode="summary", sort_keys=[]) - - assert "root_input" not in query - assert "root_output" not in query - assert "'' AS input" in query - assert "'' AS output" in query - - -def test_session_detailed_query_reads_full_input_and_output() -> None: - query = _make_list_sql(mode="detailed", sort_keys=[]) - - assert "root_input AS input" in query - assert "root_output AS output" in query - assert "substringUTF8(root_input" not in query - - -# --------------------------------------------------------------------------- -# Sort: default (no sort_keys) -# The original ORDER BY must be preserved so existing behaviour is unchanged. -# --------------------------------------------------------------------------- - - -def test_default_sort_preserves_original_order() -> None: - # sort_keys=[] means no sort param was sent — fall back to start_time ASC. - query = _make_list_sql(mode="summary", sort_keys=[]) - - # Both page_sessions and the final SELECT must use the default order. - assert "ORDER BY start_time ASC, root_span_id ASC" in query - assert "ORDER BY sessions.start_time ASC, sessions.root_span_id ASC" in query - # No pre_page_metrics CTE should be injected. - assert "pre_page_metrics" not in query - - -# --------------------------------------------------------------------------- -# Sort: single field (trace_index column — no pre-metrics join needed) -# --------------------------------------------------------------------------- - - -def test_single_field_sort_latency_desc() -> None: - query = _make_list_sql(mode="summary", sort_keys=[("latency_ms", True)]) - - # page_sessions ORDER BY should use the scoped_sessions alias (no prefix). - assert "ORDER BY latency_ms DESC NULLS LAST, root_span_id ASC" in query - # Final SELECT ORDER BY uses the sessions. prefix. - assert "ORDER BY sessions.latency_ms DESC NULLS LAST, sessions.root_span_id ASC" in query - # No pre_page_metrics needed for a trace_index column. - assert "pre_page_metrics" not in query - - -def test_single_field_sort_status_asc() -> None: - query = _make_list_sql(mode="summary", sort_keys=[("status", False)]) - - assert "ORDER BY root_span_status ASC NULLS LAST, root_span_id ASC" in query - assert "ORDER BY sessions.root_span_status ASC NULLS LAST, sessions.root_span_id ASC" in query - assert "pre_page_metrics" not in query - - -# --------------------------------------------------------------------------- -# Sort: multi-field -# --------------------------------------------------------------------------- - - -def test_multi_field_sort_applies_keys_in_order() -> None: - # Primary: cost DESC, tie-break: latency ASC. Cost requires the two-query path. - ids_query = _metric_sort_page_ids_sql( - trace_index_table="trace_index", - spans_table="spans", - scoped_filter_sql="", - sort_keys=[("cost_total_usd", True), ("latency_ms", False)], - ) - # The ids query orders by pm. for cost, plain column for latency, s.root_span_id tiebreaker. - assert "pm.cost_total_usd DESC NULLS LAST, latency_ms ASC NULLS LAST, s.root_span_id ASC" in ids_query - # Hydrate query has no ORDER BY — caller sorts in Python. - hydrate_query = _hydrate_by_ids_sql( - trace_index_table="trace_index", - spans_table="spans", - evaluator_results_table="evaluator_results", - mode="summary", + assert page.total == 1 + assert [row.session_id for row in page.rows] == ["session-a"] + assert page.rows[0].evaluator_scores == {"reward": 0.9} + assert executor.tables == [ + ClickHouseTable.TRACE_INDEX, + ClickHouseTable.SPANS, + ClickHouseTable.EVALUATOR_RESULTS, + ] + assert [query.name for query in executor.queries] == [ + "evaluation_sessions.count", + "evaluation_sessions.list", + ] + for query in executor.queries: + assert workspace not in query.statement + assert evaluation_name not in query.statement + assert test_case_id not in query.statement + assert query.parameters["workspace"] == workspace + assert query.parameters["evaluation_name"] == evaluation_name + assert query.parameters["test_case_id"] == test_case_id + assert query.parameters["status"] == "error" + assert executor.queries[1].parameters["limit"] == 5 + assert executor.queries[1].parameters["offset"] == 5 + + +@pytest.mark.asyncio +async def test_metric_sort_uses_bounded_page_then_restores_hydration_order() -> None: + executor = _Executor( + [ + [{"count()": 2}], + [{"session_id": "session-b"}, {"session_id": "session-a"}], + [_session_record("session-a"), _session_record("session-b")], + ] ) - assert "ORDER BY" not in hydrate_query - - -# --------------------------------------------------------------------------- -# Sort: cost_total_usd and tokens go through the two-query path -# --------------------------------------------------------------------------- - -def test_cost_sort_uses_two_query_path() -> None: - ids_query = _metric_sort_page_ids_sql( - trace_index_table="trace_index", - spans_table="spans", - scoped_filter_sql="", - sort_keys=[("cost_total_usd", True)], - ) - # pre_page_metrics must appear before the final SELECT so the ORDER BY can reference pm. - pre_pos = ids_query.index("pre_page_metrics AS (") - select_pos = ids_query.index("SELECT s.workspace, s.session_id") - assert pre_pos < select_pos - assert "LEFT JOIN pre_page_metrics AS pm" in ids_query - # Hydrate uses session_ids IN list, no pre_page_metrics. - hydrate_query = _hydrate_by_ids_sql( - trace_index_table="trace_index", - spans_table="spans", - evaluator_results_table="evaluator_results", + page = await _repository(executor).list_sessions( + workspace="default", + evaluation_name="evaluation-a", + page=1, + page_size=2, mode="summary", + sort_keys=[("cost_total_usd", True)], ) - assert "pre_page_metrics" not in hydrate_query - assert "session_id IN %(session_ids)s" in hydrate_query - -def test_tokens_sort_uses_two_query_path() -> None: - ids_query = _metric_sort_page_ids_sql( - trace_index_table="trace_index", - spans_table="spans", - scoped_filter_sql="", - sort_keys=[("tokens", False)], + assert page.total == 2 + assert [row.session_id for row in page.rows] == ["session-b", "session-a"] + assert [query.name for query in executor.queries] == [ + "evaluation_sessions.count", + "evaluation_sessions.metric_sort.page_ids", + "evaluation_sessions.metric_sort.hydrate", + ] + assert executor.queries[2].parameters["session_ids"] == ["session-b", "session-a"] + + +@pytest.mark.asyncio +async def test_empty_result_stops_after_count() -> None: + executor = _Executor([[{"count()": 0}]]) + + page = await _repository(executor).list_sessions( + workspace="default", + evaluation_name="evaluation-a", + page=1, + page_size=100, + mode="detailed", ) - assert "pre_page_metrics AS (" in ids_query - assert "pm.total_tokens ASC NULLS LAST" in ids_query - # NULL-safe coalesce in the pre_page_metrics total_tokens computation. - assert "coalesce" in ids_query - - -# --------------------------------------------------------------------------- -# Sort: stable tiebreaker is always appended -# --------------------------------------------------------------------------- - - -def test_tiebreaker_always_appended() -> None: - # Even a single-field sort should end with root_span_id for determinism. - query = _make_list_sql(mode="summary", sort_keys=[("started_at", False)]) - - assert "root_span_id ASC" in query - - -# --------------------------------------------------------------------------- -# _build_order_by unit tests -# --------------------------------------------------------------------------- - - -def test_build_order_by_single_asc() -> None: - result = _build_order_by([("latency_ms", False)], _SORT_EXPR_PAGE, "root_span_id ASC") - assert result == "latency_ms ASC NULLS LAST, root_span_id ASC" - - -def test_build_order_by_single_desc() -> None: - result = _build_order_by([("cost_total_usd", True)], _SORT_EXPR_PAGE, "s.root_span_id ASC") - assert result == "pm.cost_total_usd DESC NULLS LAST, s.root_span_id ASC" - -def test_build_order_by_multi() -> None: - result = _build_order_by( - [("cost_total_usd", True), ("latency_ms", False)], - _SORT_EXPR_FINAL, - "sessions.root_span_id ASC", + assert page.total == 0 + assert page.rows == [] + assert [query.name for query in executor.queries] == ["evaluation_sessions.count"] + + +@pytest.mark.asyncio +async def test_metric_sort_rejects_unbounded_session_set_after_count() -> None: + executor = _Executor([[{"count()": _MAX_METRIC_SORT_SESSIONS + 1}]]) + + with pytest.raises(MetricSortTooLargeError) as exc_info: + await _repository(executor).list_sessions( + workspace="default", + evaluation_name="evaluation-a", + page=1, + page_size=100, + mode="summary", + sort_keys=[("tokens", False)], + ) + + assert exc_info.value.total == _MAX_METRIC_SORT_SESSIONS + 1 + assert exc_info.value.limit == _MAX_METRIC_SORT_SESSIONS + assert [query.name for query in executor.queries] == ["evaluation_sessions.count"] + + +def test_build_order_by_uses_registered_expressions_in_key_order() -> None: + assert ( + _build_order_by( + [("cost_total_usd", True), ("latency_ms", False)], + _SORT_EXPR_FINAL, + "sessions.root_span_id ASC", + ) + == "metrics.cost_total_usd DESC NULLS LAST, sessions.latency_ms ASC NULLS LAST, " + "sessions.root_span_id ASC" ) - assert result == ( - "metrics.cost_total_usd DESC NULLS LAST, sessions.latency_ms ASC NULLS LAST, sessions.root_span_id ASC" + assert ( + _build_order_by( + [("tokens", False)], + _SORT_EXPR_PAGE, + "s.root_span_id ASC", + ) + == "pm.total_tokens ASC NULLS LAST, s.root_span_id ASC" ) diff --git a/services/intake/tests/test_experiment_session_schemas.py b/services/intake/tests/test_experiment_session_schemas.py index a8d2533056..dc3eaa4178 100644 --- a/services/intake/tests/test_experiment_session_schemas.py +++ b/services/intake/tests/test_experiment_session_schemas.py @@ -4,8 +4,8 @@ from datetime import datetime, timezone from nmp.intake.api.v2.experiments.schemas import EvaluationSessionResponse +from nmp.intake.repository.evaluation_session import EvaluationSessionRow from nmp.intake.spans.domain import SpanStatus -from nmp.intake.spans.evaluation_session_repository import EvaluationSessionRow def test_evaluation_session_from_row_preserves_detailed_payloads() -> None: From 7feb4d45906abff347be650c550e914e99b60a52 Mon Sep 17 00:00:00 2001 From: Brian Newsom Date: Mon, 27 Jul 2026 16:14:45 -0600 Subject: [PATCH 5/7] refactor(intake): add evaluation read service Signed-off-by: Brian Newsom --- .../intake/api/v2/experiments/dependencies.py | 61 +++++ .../intake/api/v2/experiments/endpoints.py | 230 ++++++------------ .../nmp/intake/experiments/read_service.py | 205 ++++++++++++++++ services/intake/tests/conftest.py | 2 +- .../integration/test_experiments_crud.py | 2 +- .../tests/test_clickhouse_architecture.py | 2 +- .../tests/test_evaluation_read_service.py | 206 ++++++++++++++++ .../tests/test_experiment_sort_endpoint.py | 2 +- 8 files changed, 555 insertions(+), 155 deletions(-) create mode 100644 services/intake/src/nmp/intake/api/v2/experiments/dependencies.py create mode 100644 services/intake/src/nmp/intake/experiments/read_service.py create mode 100644 services/intake/tests/test_evaluation_read_service.py diff --git a/services/intake/src/nmp/intake/api/v2/experiments/dependencies.py b/services/intake/src/nmp/intake/api/v2/experiments/dependencies.py new file mode 100644 index 0000000000..d56db74624 --- /dev/null +++ b/services/intake/src/nmp/intake/api/v2/experiments/dependencies.py @@ -0,0 +1,61 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Dependencies for the Evaluations API.""" + +from typing import Annotated + +from fastapi import Depends, Request +from nmp.common.entities.client import EntityClient +from nmp.common.service.dependencies import get_entity_client +from nmp.intake.experiments.read_service import EvaluationReadService +from nmp.intake.repository.clickhouse.evaluation_rollup import ClickHouseEvaluationRollupRepository +from nmp.intake.repository.clickhouse.evaluation_session import ClickHouseEvaluationSessionRepository +from nmp.intake.repository.clickhouse.executor import ClickHouseExecutor +from nmp.intake.repository.evaluation_rollup import EvaluationRollupRepository +from nmp.intake.repository.evaluation_session import EvaluationSessionRepository +from nmp.intake.spans.clickhouse_client import ClickHouseSpanClient + +EntityClientDep = Annotated[EntityClient, Depends(get_entity_client)] + + +def _get_clickhouse_client(request: Request) -> ClickHouseSpanClient | None: + service = getattr(request.app.state, "intake_service", None) or getattr(request.app.state, "service", None) + if service is None: + return None + return getattr(service, "clickhouse_client", None) + + +def get_evaluation_rollup_repository(request: Request) -> EvaluationRollupRepository | None: + # Rollups are enrichment only. Evaluation entity reads should continue when + # ClickHouse is disabled or temporarily unavailable. + client = _get_clickhouse_client(request) + return ClickHouseEvaluationRollupRepository(ClickHouseExecutor(client)) if client is not None else None + + +EvaluationRollupRepositoryDep = Annotated[EvaluationRollupRepository | None, Depends(get_evaluation_rollup_repository)] + + +def get_evaluation_session_repository(request: Request) -> EvaluationSessionRepository | None: + client = _get_clickhouse_client(request) + return ClickHouseEvaluationSessionRepository(ClickHouseExecutor(client)) if client is not None else None + + +EvaluationSessionRepositoryDep = Annotated[ + EvaluationSessionRepository | None, Depends(get_evaluation_session_repository) +] + + +def get_evaluation_read_service( + entity_client: EntityClientDep, + rollup_repository: EvaluationRollupRepositoryDep, + session_repository: EvaluationSessionRepositoryDep, +) -> EvaluationReadService: + return EvaluationReadService( + entity_client=entity_client, + rollup_repository=rollup_repository, + session_repository=session_repository, + ) + + +EvaluationReadServiceDep = Annotated[EvaluationReadService, Depends(get_evaluation_read_service)] diff --git a/services/intake/src/nmp/intake/api/v2/experiments/endpoints.py b/services/intake/src/nmp/intake/api/v2/experiments/endpoints.py index 966fa1c088..b54c898819 100644 --- a/services/intake/src/nmp/intake/api/v2/experiments/endpoints.py +++ b/services/intake/src/nmp/intake/api/v2/experiments/endpoints.py @@ -3,10 +3,10 @@ """Create, list, get, and delete endpoints for Evaluations and ExperimentGroups. -Entity-store (Postgres) operations are wired directly onto ``EntityClient``, -following the inline pattern used by the core services. PUT updates only the -mutable fields; an Evaluation's identity and the dataset/agent it ran against -are fixed. Rollup fields on read models are hydrated from ClickHouse. +Entity-store writes use ``EntityClient`` directly. Evaluation reads use an +application service that composes Entities with ClickHouse-backed rollups and +sessions. PUT updates only mutable fields; an Evaluation's identity and the +dataset/agent it ran against are fixed. """ from __future__ import annotations @@ -23,7 +23,7 @@ from nmp.common.api.parsed_filter import ParsedFilter, make_filter_dep from nmp.common.api.utils import generate_openapi_extra_params from nmp.common.entities.client import EntityClient, EntityConflictError, EntityNotFoundError -from nmp.common.service.dependencies import get_entity_client +from nmp.intake.api.v2.experiments.dependencies import EntityClientDep, EvaluationReadServiceDep from nmp.intake.api.v2.experiments.schemas import ( EvaluationFilter, EvaluationPatchRequest, @@ -43,23 +43,22 @@ # layer uses; only the entity's own field names (e.g. parent_experiment_id) reference Experiment directly. from nmp.intake.entities.experiments import Experiment as Evaluation from nmp.intake.entities.experiments import ExperimentGroup -from nmp.intake.repository.clickhouse.evaluation_rollup import ClickHouseEvaluationRollupRepository -from nmp.intake.repository.clickhouse.evaluation_session import ClickHouseEvaluationSessionRepository -from nmp.intake.repository.clickhouse.executor import ClickHouseExecutor -from nmp.intake.repository.evaluation_rollup import EvaluationRollup, EvaluationRollupRepository, ScoreRollup -from nmp.intake.repository.evaluation_session import EvaluationSessionRepository, MetricSortTooLargeError +from nmp.intake.experiments.read_service import ( + EvaluationNotFoundError, + EvaluationRead, + EvaluationReadLimitExceededError, + EvaluationTelemetryUnavailableError, + InvalidEvaluationSessionStatusError, +) +from nmp.intake.repository.evaluation_rollup import EvaluationRollup, ScoreRollup +from nmp.intake.repository.evaluation_session import MetricSortTooLargeError from nmp.intake.spans.api.dependencies import require_workspace_access, validate_list_query_params -from nmp.intake.spans.clickhouse_client import ClickHouseSpanClient from nmp.intake.spans.domain import SpanStatus from nmp.intake.spans.storage import make_pagination logger = logging.getLogger(__name__) -def _sanitize_for_log(value: str) -> str: - return value.replace("\r", "").replace("\n", "") - - router = APIRouter(dependencies=[Depends(require_workspace_access)]) GROUPS_TAG = "Experiment Groups" @@ -93,39 +92,11 @@ def _sanitize_for_log(value: str) -> str: EntityT = TypeVar("EntityT", Evaluation, ExperimentGroup) -EntityClientDep = Annotated[EntityClient, Depends(get_entity_client)] ExperimentGroupFilterDep = Annotated[ParsedFilter, Depends(make_filter_dep(ExperimentGroupFilter))] EvaluationFilterDep = Annotated[ParsedFilter, Depends(make_filter_dep(EvaluationFilter))] EvaluationSessionFilterDep = Annotated[ParsedFilter, Depends(make_filter_dep(EvaluationSessionFilter))] -def _get_clickhouse_client(request: Request) -> ClickHouseSpanClient | None: - service = getattr(request.app.state, "intake_service", None) or getattr(request.app.state, "service", None) - if service is None: - return None - return getattr(service, "clickhouse_client", None) - - -def get_evaluation_rollup_repository(request: Request) -> EvaluationRollupRepository | None: - # Rollups are enrichment only. Evaluation entity reads should continue when - # ClickHouse is disabled or temporarily unavailable. - client = _get_clickhouse_client(request) - return ClickHouseEvaluationRollupRepository(ClickHouseExecutor(client)) if client is not None else None - - -EvaluationRollupRepositoryDep = Annotated[EvaluationRollupRepository | None, Depends(get_evaluation_rollup_repository)] - - -def get_evaluation_session_repository(request: Request) -> EvaluationSessionRepository | None: - client = _get_clickhouse_client(request) - return ClickHouseEvaluationSessionRepository(ClickHouseExecutor(client)) if client is not None else None - - -EvaluationSessionRepositoryDep = Annotated[ - EvaluationSessionRepository | None, Depends(get_evaluation_session_repository) -] - - @router.post( "/v2/workspaces/{workspace}/experiment-groups", response_model=ExperimentGroupResponse, @@ -405,8 +376,7 @@ async def create_evaluation( async def list_evaluations( workspace: str, request: Request, - entity_client: EntityClientDep, - rollup_repository: EvaluationRollupRepositoryDep, + read_service: EvaluationReadServiceDep, parsed: EvaluationFilterDep, page: int = Query(default=1, ge=1, description="Page number."), page_size: int = Query(default=100, ge=1, le=1000, description="Page size."), @@ -447,34 +417,31 @@ async def list_evaluations( # Compute-on-read: fetch the whole (entity-filtered) group, hydrate every rollup, then filter, sort, # and paginate in memory so a single request can sort/filter by a ClickHouse metric that lives # outside the entity store. Bounded to hundreds of evaluations per group (see _MAX_GROUP_EVALUATIONS). - result = await entity_client.list( - Evaluation, - workspace=workspace, - filter_operation=entity_operation, - page=1, - page_size=_MAX_GROUP_EVALUATIONS, - ) - responses = [EvaluationResponse.from_entity(e) for e in result.data] - total_selected = result.pagination.total_results - if total_selected > _MAX_GROUP_EVALUATIONS: + try: + result = await read_service.list_evaluations( + workspace=workspace, + filter_operation=entity_operation, + limit=_MAX_GROUP_EVALUATIONS, + ) + except EvaluationReadLimitExceededError as exc: # The whole filtered set is sorted in memory; anything past the fetch cap can't be sorted, so a # returned page would be silently incomplete. Fail loudly and tell the caller how to scope the # query instead (or denormalize rollup metrics for entity-store sorting once groups grow this big). logger.warning( "Evaluation list selected %d evaluations, over the %d-row in-memory sort cap; refusing " "to return a partially sorted result.", - total_selected, - _MAX_GROUP_EVALUATIONS, + exc.selected, + exc.limit, ) raise HTTPException( status_code=status.HTTP_413_REQUEST_ENTITY_TOO_LARGE, detail=( - f"This query selects {total_selected} evaluations, exceeding the maximum of " - f"{_MAX_GROUP_EVALUATIONS} that can be sorted in one request. Narrow the result with a " + f"This query selects {exc.selected} evaluations, exceeding the maximum of " + f"{exc.limit} that can be sorted in one request. Narrow the result with a " "filter (e.g. experiment_group_id)." ), - ) - hydrated = await _hydrate_rollups(workspace=workspace, responses=responses, rollup_repository=rollup_repository) + ) from exc + responses = [_to_evaluation_response(evaluation) for evaluation in result.evaluations] # A metric-backed sort or filter is meaningless without rollups: if hydration was skipped (ClickHouse # disabled or down) every metric value would be unset, so a metric sort would silently collapse to # name order and a metric filter would drop everything. Reject the request instead of returning a @@ -482,7 +449,7 @@ async def list_evaluations( # An explicit metric sort or metric filter genuinely can't be served without rollups → 503. A # default sort degrades gracefully instead: its metric values come back unset, so the appended # -created_at key orders the list (the documented fallback), no error. - if not hydrated and (explicit_metric_sort or metric_predicates): + if not result.rollups_available and (explicit_metric_sort or metric_predicates): raise HTTPException( status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="Cannot sort or filter evaluations by a rollup metric: the telemetry store is unavailable.", @@ -511,20 +478,13 @@ async def list_evaluations( async def get_evaluation( workspace: str, name: str, - entity_client: EntityClientDep, - rollup_repository: EvaluationRollupRepositoryDep, + read_service: EvaluationReadServiceDep, ) -> EvaluationResponse: - entity = await _get_or_404( - entity_client, - Evaluation, - workspace=workspace, - name=name, - label="Evaluation", - ) - _reject_if_deleted(entity, workspace=workspace, name=name, label="Evaluation") - response = EvaluationResponse.from_entity(entity) - await _hydrate_rollups(workspace=workspace, responses=[response], rollup_repository=rollup_repository) - return response + try: + evaluation = await read_service.get_evaluation(workspace=workspace, name=name) + except EvaluationNotFoundError as exc: + raise _evaluation_not_found_http_error(exc) from exc + return _to_evaluation_response(evaluation) # Identity and the dataset it was run against are fixed for the life of an @@ -547,7 +507,7 @@ async def update_evaluation( name: str, body: EvaluationRequest, entity_client: EntityClientDep, - rollup_repository: EvaluationRollupRepositoryDep, + read_service: EvaluationReadServiceDep, ) -> EvaluationResponse: existing = await _get_or_404( entity_client, @@ -580,9 +540,7 @@ async def update_evaluation( existing.status = body.status existing.root_cause = body.root_cause updated = await entity_client.update(existing) - response = EvaluationResponse.from_entity(updated) - await _hydrate_rollups(workspace=workspace, responses=[response], rollup_repository=rollup_repository) - return response + return await _evaluation_response_with_rollup(read_service, workspace=workspace, evaluation=updated) @router.patch( @@ -599,7 +557,7 @@ async def patch_evaluation( name: str, body: EvaluationPatchRequest, entity_client: EntityClientDep, - rollup_repository: EvaluationRollupRepositoryDep, + read_service: EvaluationReadServiceDep, ) -> EvaluationResponse: """Partially update an evaluation: only fields present in the request are changed. @@ -639,9 +597,7 @@ async def patch_evaluation( existing.root_cause = body.root_cause updated = await entity_client.update(existing) - response = EvaluationResponse.from_entity(updated) - await _hydrate_rollups(workspace=workspace, responses=[response], rollup_repository=rollup_repository) - return response + return await _evaluation_response_with_rollup(read_service, workspace=workspace, evaluation=updated) @router.delete( @@ -676,7 +632,7 @@ async def pin_evaluation( workspace: str, name: str, entity_client: EntityClientDep, - rollup_repository: EvaluationRollupRepositoryDep, + read_service: EvaluationReadServiceDep, ) -> EvaluationResponse: """Pin an evaluation to the top of the list (workspace-shared). @@ -693,9 +649,7 @@ async def pin_evaluation( _reject_if_deleted(entity, workspace=workspace, name=name, label="Evaluation") entity.pinned_at = datetime.now(timezone.utc) updated = await entity_client.update(entity) - response = EvaluationResponse.from_entity(updated) - await _hydrate_rollups(workspace=workspace, responses=[response], rollup_repository=rollup_repository) - return response + return await _evaluation_response_with_rollup(read_service, workspace=workspace, evaluation=updated) @router.delete( @@ -708,7 +662,7 @@ async def unpin_evaluation( workspace: str, name: str, entity_client: EntityClientDep, - rollup_repository: EvaluationRollupRepositoryDep, + read_service: EvaluationReadServiceDep, ) -> EvaluationResponse: """Unpin an evaluation. Idempotent: unpinning an already-unpinned evaluation is a no-op.""" entity = await _get_or_404( @@ -721,9 +675,7 @@ async def unpin_evaluation( _reject_if_deleted(entity, workspace=workspace, name=name, label="Evaluation") entity.pinned_at = None updated = await entity_client.update(entity) - response = EvaluationResponse.from_entity(updated) - await _hydrate_rollups(workspace=workspace, responses=[response], rollup_repository=rollup_repository) - return response + return await _evaluation_response_with_rollup(read_service, workspace=workspace, evaluation=updated) @router.get( @@ -744,8 +696,7 @@ async def list_evaluation_sessions( workspace: str, name: str, request: Request, - entity_client: EntityClientDep, - session_repository: EvaluationSessionRepositoryDep, + read_service: EvaluationReadServiceDep, parsed: EvaluationSessionFilterDep, page: int = Query(default=1, ge=1, description="Page number."), page_size: int = Query(default=100, ge=1, le=1000, description="Page size."), @@ -768,39 +719,26 @@ async def list_evaluation_sessions( ) -> Page[EvaluationSessionResponse]: validate_list_query_params(request, additional_params={"mode"}) sort_keys = _parse_session_sort_keys(sort) if sort is not None else None - evaluation = await _get_or_404( - entity_client, - Evaluation, - workspace=workspace, - name=name, - label="Evaluation", - ) - _reject_if_deleted(evaluation, workspace=workspace, name=name, label="Evaluation") - if session_repository is None: - raise HTTPException( - status_code=status.HTTP_503_SERVICE_UNAVAILABLE, - detail="ClickHouse is unavailable; per-session reads require telemetry storage.", - ) test_case_id: str | None = parsed.extract("test_case_id") status_raw: str | None = parsed.extract("status") try: - status_filter = SpanStatus(status_raw) if status_raw is not None else None - except ValueError: - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, - detail=f"Invalid status '{status_raw}'. Valid values: {[s.value for s in SpanStatus]}", - ) - try: - result = await session_repository.list_sessions( + result = await read_service.list_sessions( workspace=workspace, evaluation_name=name, - status=status_filter, + status=status_raw, test_case_id=test_case_id, page=page, page_size=page_size, mode=mode, sort_keys=sort_keys, ) + except EvaluationNotFoundError as exc: + raise _evaluation_not_found_http_error(exc) from exc + except InvalidEvaluationSessionStatusError as exc: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"Invalid status '{exc.value}'. Valid values: {[s.value for s in SpanStatus]}", + ) from exc except MetricSortTooLargeError as exc: raise HTTPException( status_code=status.HTTP_413_REQUEST_ENTITY_TOO_LARGE, @@ -811,18 +749,15 @@ async def list_evaluation_sessions( "different field (started_at, latency_ms, status, test_case_id)." ), ) from exc - except Exception as exc: - # Sessions are the response payload (not enrichment), so we can't silently degrade like - # _hydrate_rollups does. Convert backend failures (ClickHouse connection drop, query - # timeout, etc.) to a deterministic 503 instead of letting them bubble as 500s. - logger.exception( - "Per-session read failed for workspace=%s evaluation=%s", - _sanitize_for_log(workspace), - _sanitize_for_log(name), + except EvaluationTelemetryUnavailableError as exc: + detail = ( + "Telemetry store unavailable." + if exc.configured + else "ClickHouse is unavailable; per-session reads require telemetry storage." ) raise HTTPException( status_code=status.HTTP_503_SERVICE_UNAVAILABLE, - detail="Telemetry store unavailable.", + detail=detail, ) from exc data = [EvaluationSessionResponse.from_row(row, mode=mode) for row in result.rows] return Page( @@ -1370,35 +1305,28 @@ def _sort_evaluations( return ordered -async def _hydrate_rollups( +def _to_evaluation_response(evaluation: EvaluationRead) -> EvaluationResponse: + response = EvaluationResponse.from_entity(evaluation.entity) + if evaluation.rollup is not None: + _apply_rollup(response, evaluation.rollup) + return response + + +async def _evaluation_response_with_rollup( + read_service: EvaluationReadServiceDep, *, workspace: str, - responses: list[EvaluationResponse], - rollup_repository: EvaluationRollupRepository | None, -) -> bool: - """Enrich responses with ClickHouse rollups in place. - - Returns True when hydration completed (including the no-op empty-list case) and False when it was - skipped because the rollup store is unavailable (repository absent or query failed). Callers that - sort by a rollup metric use the flag to reject the request rather than silently degrade; callers - that only display metrics can ignore it. - """ - if not responses: - return True - if rollup_repository is None: - return False - try: - rollups = await rollup_repository.get_rollups( - workspace=workspace, evaluation_ids=[response.name for response in responses] - ) - except Exception: - logger.exception("Skipping evaluation rollup hydration because ClickHouse is unavailable") - return False - for response in responses: - rollup = rollups.get(response.name) - if rollup is not None: - _apply_rollup(response, rollup) - return True + evaluation: Evaluation, +) -> EvaluationResponse: + batch = await read_service.attach_rollups(workspace=workspace, evaluations=[evaluation]) + return _to_evaluation_response(batch.evaluations[0]) + + +def _evaluation_not_found_http_error(exc: EvaluationNotFoundError) -> HTTPException: + return HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Evaluation '{exc.workspace}/{exc.name}' not found.", + ) def _apply_rollup(response: EvaluationResponse, rollup: EvaluationRollup) -> None: diff --git a/services/intake/src/nmp/intake/experiments/read_service.py b/services/intake/src/nmp/intake/experiments/read_service.py new file mode 100644 index 0000000000..31503b0d28 --- /dev/null +++ b/services/intake/src/nmp/intake/experiments/read_service.py @@ -0,0 +1,205 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Cross-store read composition for Evaluations.""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass + +from nmp.common.api.filter import FilterOperation +from nmp.common.entities.client import EntityClient, EntityNotFoundError +from nmp.intake.entities.experiments import Experiment as Evaluation +from nmp.intake.repository.evaluation_rollup import EvaluationRollup, EvaluationRollupRepository +from nmp.intake.repository.evaluation_session import ( + EvaluationSessionPage, + EvaluationSessionRepository, + MetricSortTooLargeError, +) +from nmp.intake.spans.domain import IntakeResponseMode, SpanStatus + +logger = logging.getLogger(__name__) + + +class EvaluationNotFoundError(Exception): + def __init__(self, workspace: str, name: str) -> None: + super().__init__(f"Evaluation {workspace}/{name} not found") + self.workspace = workspace + self.name = name + + +class EvaluationReadLimitExceededError(Exception): + def __init__(self, selected: int, limit: int) -> None: + super().__init__(f"Evaluation read selected {selected} rows, limit is {limit}") + self.selected = selected + self.limit = limit + + +class EvaluationTelemetryUnavailableError(Exception): + """Raised when telemetry is required to serve an Evaluation read.""" + + def __init__(self, *, configured: bool) -> None: + self.configured = configured + super().__init__("Evaluation telemetry store unavailable") + + +class InvalidEvaluationSessionStatusError(Exception): + def __init__(self, value: str) -> None: + self.value = value + super().__init__(f"Invalid Evaluation session status: {value}") + + +@dataclass(frozen=True) +class EvaluationRead: + entity: Evaluation + rollup: EvaluationRollup | None + + +@dataclass(frozen=True) +class EvaluationReadBatch: + evaluations: list[EvaluationRead] + total: int + rollups_available: bool + + +class EvaluationReadService: + """Compose Evaluation entities with their ClickHouse-backed read data.""" + + def __init__( + self, + *, + entity_client: EntityClient, + rollup_repository: EvaluationRollupRepository | None, + session_repository: EvaluationSessionRepository | None, + ) -> None: + self._entities = entity_client + self._rollups = rollup_repository + self._sessions = session_repository + + async def list_evaluations( + self, + *, + workspace: str, + filter_operation: FilterOperation | None, + limit: int, + ) -> EvaluationReadBatch: + result = await self._entities.list( + Evaluation, + workspace=workspace, + filter_operation=filter_operation, + page=1, + page_size=limit, + ) + total = result.pagination.total_results + if total > limit: + raise EvaluationReadLimitExceededError(total, limit) + return await self.attach_rollups( + workspace=workspace, + evaluations=result.data, + total=total, + ) + + async def get_evaluation(self, *, workspace: str, name: str) -> EvaluationRead: + evaluation = await self._get_live_evaluation(workspace=workspace, name=name) + batch = await self.attach_rollups(workspace=workspace, evaluations=[evaluation]) + return batch.evaluations[0] + + async def attach_rollups( + self, + *, + workspace: str, + evaluations: list[Evaluation], + total: int | None = None, + ) -> EvaluationReadBatch: + """Attach rollups to entities already loaded by an Evaluation workflow.""" + + if not evaluations: + return EvaluationReadBatch(evaluations=[], total=total or 0, rollups_available=True) + + rollups, available = await self._get_rollups( + workspace=workspace, + evaluation_names=[evaluation.name for evaluation in evaluations], + ) + return EvaluationReadBatch( + evaluations=[ + EvaluationRead(entity=evaluation, rollup=rollups.get(evaluation.name)) for evaluation in evaluations + ], + total=len(evaluations) if total is None else total, + rollups_available=available, + ) + + async def list_sessions( + self, + *, + workspace: str, + evaluation_name: str, + status: str | None, + test_case_id: str | None, + page: int, + page_size: int, + mode: IntakeResponseMode, + sort_keys: list[tuple[str, bool]] | None, + ) -> EvaluationSessionPage: + await self._get_live_evaluation(workspace=workspace, name=evaluation_name) + if self._sessions is None: + raise EvaluationTelemetryUnavailableError(configured=False) + status_filter = None + if status is not None: + try: + status_filter = SpanStatus(status) + except ValueError as exc: + raise InvalidEvaluationSessionStatusError(status) from exc + try: + return await self._sessions.list_sessions( + workspace=workspace, + evaluation_name=evaluation_name, + status=status_filter, + test_case_id=test_case_id, + page=page, + page_size=page_size, + mode=mode, + sort_keys=sort_keys, + ) + except MetricSortTooLargeError: + raise + except Exception as exc: + logger.exception( + "Per-session read failed for workspace=%s evaluation=%s", + _sanitize_for_log(workspace), + _sanitize_for_log(evaluation_name), + ) + raise EvaluationTelemetryUnavailableError(configured=True) from exc + + async def _get_live_evaluation(self, *, workspace: str, name: str) -> Evaluation: + try: + evaluation = await self._entities.get(Evaluation, workspace=workspace, name=name) + except EntityNotFoundError as exc: + raise EvaluationNotFoundError(workspace, name) from exc + if evaluation.is_deleted: + raise EvaluationNotFoundError(workspace, name) + return evaluation + + async def _get_rollups( + self, + *, + workspace: str, + evaluation_names: list[str], + ) -> tuple[dict[str, EvaluationRollup], bool]: + if self._rollups is None: + return {}, False + try: + return ( + await self._rollups.get_rollups( + workspace=workspace, + evaluation_ids=evaluation_names, + ), + True, + ) + except Exception: + logger.exception("Skipping evaluation rollup hydration because ClickHouse is unavailable") + return {}, False + + +def _sanitize_for_log(value: str) -> str: + return value.replace("\r", "").replace("\n", "") diff --git a/services/intake/tests/conftest.py b/services/intake/tests/conftest.py index 2ccfd3f1c8..9b516fb66c 100644 --- a/services/intake/tests/conftest.py +++ b/services/intake/tests/conftest.py @@ -5,7 +5,7 @@ import pytest from fastapi.testclient import TestClient -from nmp.intake.api.v2.experiments.endpoints import get_evaluation_rollup_repository +from nmp.intake.api.v2.experiments.dependencies import get_evaluation_rollup_repository from nmp.intake.service import IntakeService from nmp.testing import create_test_client diff --git a/services/intake/tests/integration/test_experiments_crud.py b/services/intake/tests/integration/test_experiments_crud.py index c63bc9c795..0d09903009 100644 --- a/services/intake/tests/integration/test_experiments_crud.py +++ b/services/intake/tests/integration/test_experiments_crud.py @@ -9,7 +9,7 @@ from fastapi import FastAPI from fastapi.testclient import TestClient -from nmp.intake.api.v2.experiments.endpoints import get_evaluation_rollup_repository +from nmp.intake.api.v2.experiments.dependencies import get_evaluation_rollup_repository GROUPS = "/apis/intake/v2/workspaces/default/experiment-groups" EVALUATIONS = "/apis/intake/v2/workspaces/default/evaluations" diff --git a/services/intake/tests/test_clickhouse_architecture.py b/services/intake/tests/test_clickhouse_architecture.py index 6f9f52c0c1..4225c0eb2a 100644 --- a/services/intake/tests/test_clickhouse_architecture.py +++ b/services/intake/tests/test_clickhouse_architecture.py @@ -11,7 +11,7 @@ # Shrink this allowlist as each legacy repository moves behind ClickHouseExecutor. _ALLOWED_RAW_CLIENT_IMPORTS = { - "api/v2/experiments/endpoints.py", + "api/v2/experiments/dependencies.py", "repository/clickhouse/executor.py", "service.py", "spans/annotations_repository.py", diff --git a/services/intake/tests/test_evaluation_read_service.py b/services/intake/tests/test_evaluation_read_service.py new file mode 100644 index 0000000000..f68917dda9 --- /dev/null +++ b/services/intake/tests/test_evaluation_read_service.py @@ -0,0 +1,206 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Behavioral tests for Evaluation cross-store read composition.""" + +from types import SimpleNamespace +from typing import cast +from unittest.mock import AsyncMock + +import pytest +from nmp.common.entities.client import EntityClient, EntityNotFoundError +from nmp.intake.entities.experiments import Experiment as Evaluation +from nmp.intake.experiments.read_service import ( + EvaluationNotFoundError, + EvaluationReadLimitExceededError, + EvaluationReadService, + EvaluationTelemetryUnavailableError, +) +from nmp.intake.repository.evaluation_rollup import EvaluationRollup, EvaluationRollupRepository +from nmp.intake.repository.evaluation_session import EvaluationSessionPage, EvaluationSessionRepository +from nmp.intake.spans.domain import IntakeResponseMode, SpanStatus + + +class _RollupRepository(EvaluationRollupRepository): + def __init__(self, rollups: dict[str, EvaluationRollup] | Exception) -> None: + self._rollups = rollups + self.calls: list[tuple[str, list[str]]] = [] + + async def get_rollups( + self, + *, + workspace: str, + evaluation_ids: list[str], + ) -> dict[str, EvaluationRollup]: + self.calls.append((workspace, evaluation_ids)) + if isinstance(self._rollups, Exception): + raise self._rollups + return self._rollups + + +class _SessionRepository(EvaluationSessionRepository): + def __init__(self, result: EvaluationSessionPage | Exception) -> None: + self._result = result + + async def list_sessions( + self, + *, + workspace: str, + evaluation_name: str, + status: SpanStatus | None = None, + test_case_id: str | None = None, + page: int, + page_size: int, + mode: IntakeResponseMode, + sort_keys: list[tuple[str, bool]] | None = None, + ) -> EvaluationSessionPage: + if isinstance(self._result, Exception): + raise self._result + return self._result + + +def _evaluation(name: str, *, deleted: bool = False) -> Evaluation: + evaluation = Evaluation( + workspace="default", + name=name, + experiment_ids=["experiment-group-1"], + dataset_name="dataset", + ) + evaluation.is_deleted = deleted + return evaluation + + +def _entity_client() -> tuple[EntityClient, AsyncMock, AsyncMock]: + list_mock = AsyncMock() + get_mock = AsyncMock() + client = cast(EntityClient, SimpleNamespace(list=list_mock, get=get_mock)) + return client, list_mock, get_mock + + +@pytest.mark.asyncio +async def test_list_evaluations_batches_rollup_enrichment() -> None: + first = _evaluation("eval-1") + second = _evaluation("eval-2") + client, list_mock, _ = _entity_client() + list_mock.return_value = SimpleNamespace( + data=[first, second], + pagination=SimpleNamespace(total_results=2), + ) + first_rollup = EvaluationRollup(evaluation_id=first.name, run_count=4) + rollups = _RollupRepository({first.name: first_rollup}) + service = EvaluationReadService( + entity_client=client, + rollup_repository=rollups, + session_repository=None, + ) + + result = await service.list_evaluations( + workspace="default", + filter_operation=None, + limit=1000, + ) + + assert result.total == 2 + assert result.rollups_available is True + assert result.evaluations[0].entity is first + assert result.evaluations[0].rollup is first_rollup + assert result.evaluations[1].entity is second + assert result.evaluations[1].rollup is None + assert list_mock.await_count == 1 + assert rollups.calls == [("default", ["eval-1", "eval-2"])] + + +@pytest.mark.asyncio +async def test_list_evaluations_degrades_when_rollups_fail() -> None: + evaluation = _evaluation("eval-1") + client, list_mock, _ = _entity_client() + list_mock.return_value = SimpleNamespace( + data=[evaluation], + pagination=SimpleNamespace(total_results=1), + ) + service = EvaluationReadService( + entity_client=client, + rollup_repository=_RollupRepository(RuntimeError("ClickHouse unavailable")), + session_repository=None, + ) + + result = await service.list_evaluations( + workspace="default", + filter_operation=None, + limit=1000, + ) + + assert result.rollups_available is False + assert result.evaluations[0].entity is evaluation + assert result.evaluations[0].rollup is None + + +@pytest.mark.asyncio +async def test_list_evaluations_rejects_over_limit_before_rollup_lookup() -> None: + evaluation = _evaluation("eval-1") + client, list_mock, _ = _entity_client() + list_mock.return_value = SimpleNamespace( + data=[evaluation], + pagination=SimpleNamespace(total_results=1001), + ) + rollups = _RollupRepository({}) + service = EvaluationReadService( + entity_client=client, + rollup_repository=rollups, + session_repository=None, + ) + + with pytest.raises(EvaluationReadLimitExceededError) as exc_info: + await service.list_evaluations( + workspace="default", + filter_operation=None, + limit=1000, + ) + + assert exc_info.value.selected == 1001 + assert rollups.calls == [] + + +@pytest.mark.asyncio +async def test_get_evaluation_rejects_missing_and_deleted_entities() -> None: + client, _, get_mock = _entity_client() + service = EvaluationReadService( + entity_client=client, + rollup_repository=_RollupRepository({}), + session_repository=None, + ) + get_mock.side_effect = EntityNotFoundError("missing") + + with pytest.raises(EvaluationNotFoundError): + await service.get_evaluation(workspace="default", name="missing") + + get_mock.side_effect = None + get_mock.return_value = _evaluation("deleted", deleted=True) + with pytest.raises(EvaluationNotFoundError): + await service.get_evaluation(workspace="default", name="deleted") + + +@pytest.mark.asyncio +async def test_session_failures_are_translated_after_entity_validation() -> None: + client, _, get_mock = _entity_client() + get_mock.return_value = _evaluation("eval-1") + service = EvaluationReadService( + entity_client=client, + rollup_repository=None, + session_repository=_SessionRepository(RuntimeError("ClickHouse unavailable")), + ) + + with pytest.raises(EvaluationTelemetryUnavailableError) as exc_info: + await service.list_sessions( + workspace="default", + evaluation_name="eval-1", + status=None, + test_case_id=None, + page=1, + page_size=100, + mode="detailed", + sort_keys=None, + ) + + assert exc_info.value.configured is True + assert get_mock.await_count == 1 diff --git a/services/intake/tests/test_experiment_sort_endpoint.py b/services/intake/tests/test_experiment_sort_endpoint.py index 4d50833224..ca9b38aac1 100644 --- a/services/intake/tests/test_experiment_sort_endpoint.py +++ b/services/intake/tests/test_experiment_sort_endpoint.py @@ -3,7 +3,7 @@ """Endpoint-level guards for the evaluations list sort (rollups unavailable / bad field). -The shared ``client`` fixture overrides ``get_evaluation_rollup_repository`` to return ``None``, +The shared ``client`` fixture disables the Evaluation rollup repository, which is exactly the "ClickHouse disabled / unavailable" condition. A metric-backed sort cannot be computed without rollups, so it must fail loudly rather than silently degrade to name order. """ From 3b264612024f25f9f6cf92f0cbbfc3592ae14ef5 Mon Sep 17 00:00:00 2001 From: Brian Newsom Date: Mon, 27 Jul 2026 22:56:02 -0600 Subject: [PATCH 6/7] refactor(intake): migrate annotation and evaluator result storage Signed-off-by: Brian Newsom --- .../src/nmp/intake/repository/annotations.py | 36 +++++++ .../clickhouse/annotations.py} | 94 +++++++++++------- .../clickhouse/evaluator_results.py} | 99 ++++++++++++------- .../intake/repository/clickhouse/executor.py | 46 ++++++++- .../intake/repository/evaluator_results.py | 36 +++++++ .../src/nmp/intake/spans/api/dependencies.py | 14 +-- .../intake/src/nmp/intake/spans/service.py | 4 +- .../tests/test_clickhouse_architecture.py | 2 - .../intake/tests/test_clickhouse_executor.py | 71 ++++++++++++- 9 files changed, 318 insertions(+), 84 deletions(-) create mode 100644 services/intake/src/nmp/intake/repository/annotations.py rename services/intake/src/nmp/intake/{spans/annotations_repository.py => repository/clickhouse/annotations.py} (71%) rename services/intake/src/nmp/intake/{spans/evaluator_results_repository.py => repository/clickhouse/evaluator_results.py} (64%) create mode 100644 services/intake/src/nmp/intake/repository/evaluator_results.py diff --git a/services/intake/src/nmp/intake/repository/annotations.py b/services/intake/src/nmp/intake/repository/annotations.py new file mode 100644 index 0000000000..cf953c004d --- /dev/null +++ b/services/intake/src/nmp/intake/repository/annotations.py @@ -0,0 +1,36 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Repository interface for Intake annotations.""" + +from abc import ABC, abstractmethod + +from nmp.common.api.common import PaginatedResult +from nmp.intake.spans.domain import Annotation, AnnotationListFilter + + +class AnnotationsRepository(ABC): + """Domain-facing interface for annotation persistence.""" + + @abstractmethod + async def save_annotations(self, annotations: list[Annotation]) -> None: + pass + + @abstractmethod + async def get_annotation(self, *, workspace: str, annotation_id: str) -> Annotation | None: + pass + + @abstractmethod + async def list_annotations( + self, + *, + filters: AnnotationListFilter, + page: int, + page_size: int, + sort: str, + ) -> PaginatedResult[Annotation]: + pass + + @abstractmethod + async def soft_delete_annotation(self, *, annotation: Annotation) -> None: + pass diff --git a/services/intake/src/nmp/intake/spans/annotations_repository.py b/services/intake/src/nmp/intake/repository/clickhouse/annotations.py similarity index 71% rename from services/intake/src/nmp/intake/spans/annotations_repository.py rename to services/intake/src/nmp/intake/repository/clickhouse/annotations.py index b5bac91ed5..96aacad84e 100644 --- a/services/intake/src/nmp/intake/spans/annotations_repository.py +++ b/services/intake/src/nmp/intake/repository/clickhouse/annotations.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""ClickHouse implementation of Intake annotation storage.""" +"""ClickHouse implementation of Intake annotation persistence.""" from __future__ import annotations @@ -10,9 +10,11 @@ from typing import Any from nmp.common.api.common import PaginatedResult -from nmp.intake.spans.clickhouse_client import ClickHouseSpanClient +from nmp.intake.repository.annotations import AnnotationsRepository +from nmp.intake.repository.clickhouse.executor import ClickHouseExecutor, ClickHouseInsert, ClickHouseQuery +from nmp.intake.repository.clickhouse.tables import ClickHouseTable from nmp.intake.spans.domain import Annotation, AnnotationKind, AnnotationListFilter -from nmp.intake.spans.storage import dict_to_row, make_pagination, result_rows +from nmp.intake.spans.storage import dict_to_row, make_pagination ANNOTATION_COLUMNS = [ "annotation_id", @@ -36,29 +38,38 @@ } -class AnnotationsRepository: - def __init__(self, client: ClickHouseSpanClient) -> None: - self._client = client +class ClickHouseAnnotationsRepository(AnnotationsRepository): + def __init__(self, executor: ClickHouseExecutor) -> None: + self._executor = executor async def save_annotations(self, annotations: list[Annotation]) -> None: if not annotations: return rows = [dict_to_row(_annotation_to_row(item), ANNOTATION_COLUMNS) for item in annotations] - await self._client.insert("annotations", rows, column_names=ANNOTATION_COLUMNS) + await self._executor.insert( + ClickHouseInsert( + name="annotations.save", + table=ClickHouseTable.ANNOTATIONS, + rows=rows, + column_names=ANNOTATION_COLUMNS, + ) + ) async def get_annotation(self, *, workspace: str, annotation_id: str) -> Annotation | None: - result = await self._client.query( - f""" - SELECT * - FROM {self._client.table("annotations")} FINAL - WHERE workspace = %(workspace)s - AND annotation_id = %(annotation_id)s - AND is_deleted = 0 - LIMIT 1 - """, - parameters={"workspace": workspace, "annotation_id": annotation_id}, + rows = await self._executor.fetch_all( + ClickHouseQuery( + name="annotations.get", + statement=f""" + SELECT * + FROM {self._executor.table(ClickHouseTable.ANNOTATIONS)} FINAL + WHERE workspace = %(workspace)s + AND annotation_id = %(annotation_id)s + AND is_deleted = 0 + LIMIT 1 + """, + parameters={"workspace": workspace, "annotation_id": annotation_id}, + ) ) - rows = result_rows(result) if not rows: return None return _row_to_annotation(rows[0]) @@ -72,24 +83,32 @@ async def list_annotations( sort: str, ) -> PaginatedResult[Annotation]: where_sql, parameters = _annotation_where(filters) - table = self._client.table("annotations") - total_result = await self._client.query( - f"SELECT count() FROM {table} FINAL WHERE {where_sql}", - parameters=parameters, + table = self._executor.table(ClickHouseTable.ANNOTATIONS) + total_results = int( + await self._executor.fetch_scalar( + ClickHouseQuery( + name="annotations.list.count", + statement=f"SELECT count() FROM {table} FINAL WHERE {where_sql}", + parameters=parameters, + ) + ) + or 0 ) - total_results = int(total_result.result_rows[0][0]) offset = (page - 1) * page_size - rows_result = await self._client.query( - f""" - SELECT * - FROM {table} FINAL - WHERE {where_sql} - ORDER BY {_annotation_order_by(sort)} - LIMIT %(limit)s OFFSET %(offset)s - """, - parameters={**parameters, "limit": page_size, "offset": offset}, + rows = await self._executor.fetch_all( + ClickHouseQuery( + name="annotations.list.rows", + statement=f""" + SELECT * + FROM {table} FINAL + WHERE {where_sql} + ORDER BY {_annotation_order_by(sort)} + LIMIT %(limit)s OFFSET %(offset)s + """, + parameters={**parameters, "limit": page_size, "offset": offset}, + ) ) - annotations = [_row_to_annotation(row) for row in result_rows(rows_result)] + annotations = [_row_to_annotation(row) for row in rows] return PaginatedResult( data=annotations, pagination=make_pagination( @@ -111,7 +130,14 @@ async def soft_delete_annotation(self, *, annotation: Annotation) -> None: row = _annotation_to_row(annotation, is_deleted=True) row["ingested_at"] = datetime.now(timezone.utc) rows = [dict_to_row(row, ANNOTATION_COLUMNS)] - await self._client.insert("annotations", rows, column_names=ANNOTATION_COLUMNS) + await self._executor.insert( + ClickHouseInsert( + name="annotations.soft_delete", + table=ClickHouseTable.ANNOTATIONS, + rows=rows, + column_names=ANNOTATION_COLUMNS, + ) + ) def _annotation_where(filters: AnnotationListFilter) -> tuple[str, dict[str, Any]]: diff --git a/services/intake/src/nmp/intake/spans/evaluator_results_repository.py b/services/intake/src/nmp/intake/repository/clickhouse/evaluator_results.py similarity index 64% rename from services/intake/src/nmp/intake/spans/evaluator_results_repository.py rename to services/intake/src/nmp/intake/repository/clickhouse/evaluator_results.py index 21dab1eab4..e1eb668e85 100644 --- a/services/intake/src/nmp/intake/spans/evaluator_results_repository.py +++ b/services/intake/src/nmp/intake/repository/clickhouse/evaluator_results.py @@ -6,13 +6,15 @@ from typing import Any from nmp.common.api.common import PaginatedResult -from nmp.intake.spans.clickhouse_client import ClickHouseSpanClient +from nmp.intake.repository.clickhouse.executor import ClickHouseExecutor, ClickHouseInsert, ClickHouseQuery +from nmp.intake.repository.clickhouse.tables import ClickHouseTable +from nmp.intake.repository.evaluator_results import EvaluatorResultsRepository from nmp.intake.spans.domain import ( EvaluatorResult, EvaluatorResultDataType, EvaluatorResultListFilter, ) -from nmp.intake.spans.storage import dict_to_row, make_pagination, result_rows +from nmp.intake.spans.storage import dict_to_row, make_pagination EVALUATOR_RESULT_COLUMNS = [ "evaluator_result_id", @@ -35,27 +37,36 @@ } -class EvaluatorResultsRepository: - def __init__(self, client: ClickHouseSpanClient) -> None: - self._client = client +class ClickHouseEvaluatorResultsRepository(EvaluatorResultsRepository): + def __init__(self, executor: ClickHouseExecutor) -> None: + self._executor = executor async def save_evaluator_results(self, results: list[EvaluatorResult]) -> None: if not results: return rows = [dict_to_row(_evaluator_result_to_row(result), EVALUATOR_RESULT_COLUMNS) for result in results] - await self._client.insert("evaluator_results", rows, column_names=EVALUATOR_RESULT_COLUMNS) + await self._executor.insert( + ClickHouseInsert( + name="evaluator_results.save", + table=ClickHouseTable.EVALUATOR_RESULTS, + rows=rows, + column_names=EVALUATOR_RESULT_COLUMNS, + ) + ) async def get_evaluator_result(self, *, workspace: str, evaluator_result_id: str) -> EvaluatorResult | None: - result = await self._client.query( - f""" - SELECT * - FROM {self._client.table("evaluator_results")} FINAL - WHERE workspace = %(workspace)s AND evaluator_result_id = %(evaluator_result_id)s - LIMIT 1 - """, - parameters={"workspace": workspace, "evaluator_result_id": evaluator_result_id}, + rows = await self._executor.fetch_all( + ClickHouseQuery( + name="evaluator_results.get", + statement=f""" + SELECT * + FROM {self._executor.table(ClickHouseTable.EVALUATOR_RESULTS)} FINAL + WHERE workspace = %(workspace)s AND evaluator_result_id = %(evaluator_result_id)s + LIMIT 1 + """, + parameters={"workspace": workspace, "evaluator_result_id": evaluator_result_id}, + ) ) - rows = result_rows(result) if not rows: return None return _row_to_evaluator_result(rows[0]) @@ -69,23 +80,32 @@ async def list_evaluator_results( sort: str, ) -> PaginatedResult[EvaluatorResult]: where_sql, parameters = _evaluator_result_where(filters) - table = self._client.table("evaluator_results") - total_result = await self._client.query( - f"SELECT count() FROM {table} FINAL WHERE {where_sql}", parameters=parameters + table = self._executor.table(ClickHouseTable.EVALUATOR_RESULTS) + total_results = int( + await self._executor.fetch_scalar( + ClickHouseQuery( + name="evaluator_results.list.count", + statement=f"SELECT count() FROM {table} FINAL WHERE {where_sql}", + parameters=parameters, + ) + ) + or 0 ) - total_results = int(total_result.result_rows[0][0]) offset = (page - 1) * page_size - rows_result = await self._client.query( - f""" - SELECT * - FROM {table} FINAL - WHERE {where_sql} - ORDER BY {_evaluator_result_order_by(sort)} - LIMIT %(limit)s OFFSET %(offset)s - """, - parameters={**parameters, "limit": page_size, "offset": offset}, + rows = await self._executor.fetch_all( + ClickHouseQuery( + name="evaluator_results.list.rows", + statement=f""" + SELECT * + FROM {table} FINAL + WHERE {where_sql} + ORDER BY {_evaluator_result_order_by(sort)} + LIMIT %(limit)s OFFSET %(offset)s + """, + parameters={**parameters, "limit": page_size, "offset": offset}, + ) ) - results = [_row_to_evaluator_result(row) for row in result_rows(rows_result)] + results = [_row_to_evaluator_result(row) for row in rows] return PaginatedResult( data=results, pagination=make_pagination( @@ -94,16 +114,19 @@ async def list_evaluator_results( ) async def list_evaluator_results_for_span(self, *, workspace: str, span_id: str) -> list[EvaluatorResult]: - result = await self._client.query( - f""" - SELECT * - FROM {self._client.table("evaluator_results")} FINAL - WHERE workspace = %(workspace)s AND span_id = %(span_id)s - ORDER BY created_at ASC, evaluator_result_id ASC - """, - parameters={"workspace": workspace, "span_id": span_id}, + rows = await self._executor.fetch_all( + ClickHouseQuery( + name="evaluator_results.list_for_span", + statement=f""" + SELECT * + FROM {self._executor.table(ClickHouseTable.EVALUATOR_RESULTS)} FINAL + WHERE workspace = %(workspace)s AND span_id = %(span_id)s + ORDER BY created_at ASC, evaluator_result_id ASC + """, + parameters={"workspace": workspace, "span_id": span_id}, + ) ) - return [_row_to_evaluator_result(row) for row in result_rows(result)] + return [_row_to_evaluator_result(row) for row in rows] def _evaluator_result_where(filters: EvaluatorResultListFilter) -> tuple[str, dict[str, Any]]: diff --git a/services/intake/src/nmp/intake/repository/clickhouse/executor.py b/services/intake/src/nmp/intake/repository/clickhouse/executor.py index 849fc22483..ef6b8a3386 100644 --- a/services/intake/src/nmp/intake/repository/clickhouse/executor.py +++ b/services/intake/src/nmp/intake/repository/clickhouse/executor.py @@ -36,6 +36,16 @@ def bind(self, **parameters: object) -> ClickHouseQuery: ) +@dataclass(frozen=True) +class ClickHouseInsert: + """One named insert into a registered runtime table.""" + + name: str + table: ClickHouseTable + rows: Sequence[Sequence[Any]] + column_names: Sequence[str] + + class ClickHouseQueryError(RuntimeError): """Raised when a named repository query fails.""" @@ -44,8 +54,16 @@ def __init__(self, query_name: str) -> None: super().__init__(f"ClickHouse query failed: {query_name}") +class ClickHouseInsertError(RuntimeError): + """Raised when a named repository insert fails.""" + + def __init__(self, insert_name: str) -> None: + self.insert_name = insert_name + super().__init__(f"ClickHouse insert failed: {insert_name}") + + class ClickHouseExecutor: - """Execute named repository queries without exposing the raw driver result.""" + """Execute named repository operations without exposing the raw driver.""" def __init__(self, client: ClickHouseSpanClient) -> None: self._client = client @@ -81,3 +99,29 @@ async def fetch_scalar(self, query: ClickHouseQuery) -> Any | None: rows = await self.fetch_all(query) return next(iter(rows[0].values())) if rows else None + + async def insert(self, insert: ClickHouseInsert) -> None: + if not insert.rows: + return + if not isinstance(insert.table, ClickHouseTable): + raise TypeError(f"Expected ClickHouseTable, got {type(insert.table).__name__}") + + started_at = perf_counter() + try: + await self._client.insert( + insert.table.value, + insert.rows, + column_names=insert.column_names, + ) + except ClickHouseError as exc: + logger.exception("ClickHouse repository insert failed", extra={"insert_name": insert.name}) + raise ClickHouseInsertError(insert.name) from exc + finally: + logger.debug( + "ClickHouse repository insert finished", + extra={ + "insert_name": insert.name, + "row_count": len(insert.rows), + "duration_ms": (perf_counter() - started_at) * 1000, + }, + ) diff --git a/services/intake/src/nmp/intake/repository/evaluator_results.py b/services/intake/src/nmp/intake/repository/evaluator_results.py new file mode 100644 index 0000000000..4417dd5894 --- /dev/null +++ b/services/intake/src/nmp/intake/repository/evaluator_results.py @@ -0,0 +1,36 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Repository interface for Intake evaluator results.""" + +from abc import ABC, abstractmethod + +from nmp.common.api.common import PaginatedResult +from nmp.intake.spans.domain import EvaluatorResult, EvaluatorResultListFilter + + +class EvaluatorResultsRepository(ABC): + """Domain-facing interface for evaluator-result persistence.""" + + @abstractmethod + async def save_evaluator_results(self, results: list[EvaluatorResult]) -> None: + pass + + @abstractmethod + async def get_evaluator_result(self, *, workspace: str, evaluator_result_id: str) -> EvaluatorResult | None: + pass + + @abstractmethod + async def list_evaluator_results( + self, + *, + filters: EvaluatorResultListFilter, + page: int, + page_size: int, + sort: str, + ) -> PaginatedResult[EvaluatorResult]: + pass + + @abstractmethod + async def list_evaluator_results_for_span(self, *, workspace: str, span_id: str) -> list[EvaluatorResult]: + pass diff --git a/services/intake/src/nmp/intake/spans/api/dependencies.py b/services/intake/src/nmp/intake/spans/api/dependencies.py index 7077dbf895..9b2b01bac6 100644 --- a/services/intake/src/nmp/intake/spans/api/dependencies.py +++ b/services/intake/src/nmp/intake/spans/api/dependencies.py @@ -8,12 +8,14 @@ from fastapi import Depends, HTTPException, Request, status from nemo_platform import AsyncNeMoPlatform from nmp.common.service.dependencies import get_sdk_client +from nmp.intake.repository.annotations import AnnotationsRepository +from nmp.intake.repository.clickhouse.annotations import ClickHouseAnnotationsRepository +from nmp.intake.repository.clickhouse.evaluator_results import ClickHouseEvaluatorResultsRepository from nmp.intake.repository.clickhouse.executor import ClickHouseExecutor from nmp.intake.repository.clickhouse.session import ClickHouseSessionRepository +from nmp.intake.repository.evaluator_results import EvaluatorResultsRepository from nmp.intake.repository.session import SessionRepository -from nmp.intake.spans.annotations_repository import AnnotationsRepository from nmp.intake.spans.clickhouse_client import ClickHouseSpanClient, get_clickhouse_client -from nmp.intake.spans.evaluator_results_repository import EvaluatorResultsRepository from nmp.intake.spans.service import IntakeSpansService from nmp.intake.spans.span_repository import SpanRepository from nmp.intake.spans.trace_repository import TraceRepository @@ -73,15 +75,15 @@ def get_session_repository( def get_evaluator_results_repository( - client: Annotated[ClickHouseSpanClient, Depends(get_clickhouse_client)], + executor: Annotated[ClickHouseExecutor, Depends(get_clickhouse_executor)], ) -> EvaluatorResultsRepository: - return EvaluatorResultsRepository(client) + return ClickHouseEvaluatorResultsRepository(executor) def get_annotations_repository( - client: Annotated[ClickHouseSpanClient, Depends(get_clickhouse_client)], + executor: Annotated[ClickHouseExecutor, Depends(get_clickhouse_executor)], ) -> AnnotationsRepository: - return AnnotationsRepository(client) + return ClickHouseAnnotationsRepository(executor) def get_spans_service( diff --git a/services/intake/src/nmp/intake/spans/service.py b/services/intake/src/nmp/intake/spans/service.py index 35a09f6e75..ba26ee2061 100644 --- a/services/intake/src/nmp/intake/spans/service.py +++ b/services/intake/src/nmp/intake/spans/service.py @@ -8,8 +8,9 @@ from datetime import datetime from nmp.common.api.common import PaginatedResult +from nmp.intake.repository.annotations import AnnotationsRepository +from nmp.intake.repository.evaluator_results import EvaluatorResultsRepository from nmp.intake.repository.session import SessionRepository -from nmp.intake.spans.annotations_repository import AnnotationsRepository from nmp.intake.spans.domain import ( Annotation, AnnotationListFilter, @@ -25,7 +26,6 @@ TraceListFilter, TraceMode, ) -from nmp.intake.spans.evaluator_results_repository import EvaluatorResultsRepository from nmp.intake.spans.span_repository import SpanRepository from nmp.intake.spans.trace_repository import TraceRepository diff --git a/services/intake/tests/test_clickhouse_architecture.py b/services/intake/tests/test_clickhouse_architecture.py index 4225c0eb2a..4415bd5173 100644 --- a/services/intake/tests/test_clickhouse_architecture.py +++ b/services/intake/tests/test_clickhouse_architecture.py @@ -14,9 +14,7 @@ "api/v2/experiments/dependencies.py", "repository/clickhouse/executor.py", "service.py", - "spans/annotations_repository.py", "spans/api/dependencies.py", - "spans/evaluator_results_repository.py", "spans/span_repository.py", "spans/trace_repository.py", } diff --git a/services/intake/tests/test_clickhouse_executor.py b/services/intake/tests/test_clickhouse_executor.py index 1b2c08f7db..14fe3ed8ba 100644 --- a/services/intake/tests/test_clickhouse_executor.py +++ b/services/intake/tests/test_clickhouse_executor.py @@ -8,7 +8,13 @@ import pytest from clickhouse_connect.driver.exceptions import ClickHouseError -from nmp.intake.repository.clickhouse.executor import ClickHouseExecutor, ClickHouseQuery, ClickHouseQueryError +from nmp.intake.repository.clickhouse.executor import ( + ClickHouseExecutor, + ClickHouseInsert, + ClickHouseInsertError, + ClickHouseQuery, + ClickHouseQueryError, +) from nmp.intake.repository.clickhouse.tables import ClickHouseTable, qualified_table from nmp.intake.spans.clickhouse_client import ClickHouseSpanClient @@ -20,6 +26,7 @@ def __init__(self, *, error: ClickHouseError | None = None) -> None: self.error = error self.statements: list[str] = [] self.parameters: list[dict[str, Any]] = [] + self.inserts: list[tuple[str, list[tuple[object, ...]], list[str]]] = [] async def query(self, statement: str, *, parameters: dict[str, Any]) -> SimpleNamespace: self.statements.append(statement) @@ -31,6 +38,17 @@ async def query(self, statement: str, *, parameters: dict[str, Any]) -> SimpleNa result_rows=[("session-a", 3)], ) + async def insert( + self, + table: str, + rows: list[tuple[object, ...]], + *, + column_names: list[str], + ) -> None: + if self.error is not None: + raise self.error + self.inserts.append((table, rows, column_names)) + def _executor(client: _Client) -> ClickHouseExecutor: return ClickHouseExecutor(cast(ClickHouseSpanClient, client)) @@ -95,6 +113,57 @@ async def test_executor_translates_clickhouse_errors_without_sql_details() -> No assert query.statement not in str(exc_info.value) +@pytest.mark.asyncio +async def test_executor_inserts_rows_into_registered_table() -> None: + client = _Client() + insert = ClickHouseInsert( + name="annotations.save", + table=ClickHouseTable.ANNOTATIONS, + rows=[("annotation-a", "default")], + column_names=["annotation_id", "workspace"], + ) + + await _executor(client).insert(insert) + + assert client.inserts == [ + ( + "annotations", + [("annotation-a", "default")], + ["annotation_id", "workspace"], + ) + ] + + +@pytest.mark.asyncio +async def test_executor_translates_insert_errors_without_row_details() -> None: + insert = ClickHouseInsert( + name="annotations.save", + table=ClickHouseTable.ANNOTATIONS, + rows=[("secret-value",)], + column_names=["value"], + ) + + with pytest.raises(ClickHouseInsertError) as exc_info: + await _executor(_Client(error=ClickHouseError("driver details"))).insert(insert) + + assert exc_info.value.insert_name == "annotations.save" + assert str(exc_info.value) == "ClickHouse insert failed: annotations.save" + assert "secret-value" not in str(exc_info.value) + + +@pytest.mark.asyncio +async def test_executor_rejects_unregistered_insert_table() -> None: + insert = ClickHouseInsert( + name="invalid.save", + table=cast(ClickHouseTable, "system.tables"), + rows=[("value",)], + column_names=["value"], + ) + + with pytest.raises(TypeError, match="Expected ClickHouseTable"): + await _executor(_Client()).insert(insert) + + def test_table_registry_quotes_known_tables() -> None: assert qualified_table("intake", ClickHouseTable.SPANS) == "`intake`.`spans`" From 39ef316e7cc90df14400034fab2bf588c0e710ec Mon Sep 17 00:00:00 2001 From: Brian Newsom Date: Mon, 27 Jul 2026 23:10:38 -0600 Subject: [PATCH 7/7] refactor(intake): migrate span and trace storage Signed-off-by: Brian Newsom --- .../clickhouse/evaluation_session.py | 2 +- .../intake/repository/clickhouse/executor.py | 34 ++++- .../clickhouse/span.py} | 111 ++++++++++------ .../clickhouse/trace.py} | 125 ++++++++++-------- .../intake/src/nmp/intake/repository/span.py | 45 +++++++ .../intake/src/nmp/intake/repository/trace.py | 39 ++++++ .../src/nmp/intake/spans/api/dependencies.py | 26 ++-- .../intake/src/nmp/intake/spans/service.py | 4 +- .../tests/test_clickhouse_architecture.py | 2 - .../intake/tests/test_clickhouse_executor.py | 35 ++++- .../tests/test_spans_clickhouse_repository.py | 43 +++--- .../test_traces_clickhouse_repository.py | 56 ++++---- 12 files changed, 351 insertions(+), 171 deletions(-) rename services/intake/src/nmp/intake/{spans/span_repository.py => repository/clickhouse/span.py} (79%) rename services/intake/src/nmp/intake/{spans/trace_repository.py => repository/clickhouse/trace.py} (82%) create mode 100644 services/intake/src/nmp/intake/repository/span.py create mode 100644 services/intake/src/nmp/intake/repository/trace.py diff --git a/services/intake/src/nmp/intake/repository/clickhouse/evaluation_session.py b/services/intake/src/nmp/intake/repository/clickhouse/evaluation_session.py index 8427e6f0a8..6a4074e3d6 100644 --- a/services/intake/src/nmp/intake/repository/clickhouse/evaluation_session.py +++ b/services/intake/src/nmp/intake/repository/clickhouse/evaluation_session.py @@ -14,6 +14,7 @@ from nmp.intake.repository.clickhouse.executor import ClickHouseExecutor, ClickHouseQuery from nmp.intake.repository.clickhouse.tables import ClickHouseTable +from nmp.intake.repository.clickhouse.trace import current_spans_sql from nmp.intake.repository.evaluation_session import ( EvaluationSessionPage, EvaluationSessionRepository, @@ -30,7 +31,6 @@ text_query_parameters, text_select_for_mode, ) -from nmp.intake.spans.trace_repository import current_spans_sql # Sort fields that require a pre-pagination spans join to compute. These values live in the # `spans` table (not `trace_index`), so they don't exist until after the session_metrics join. diff --git a/services/intake/src/nmp/intake/repository/clickhouse/executor.py b/services/intake/src/nmp/intake/repository/clickhouse/executor.py index ef6b8a3386..f20914b4e1 100644 --- a/services/intake/src/nmp/intake/repository/clickhouse/executor.py +++ b/services/intake/src/nmp/intake/repository/clickhouse/executor.py @@ -12,12 +12,23 @@ from typing import Any from clickhouse_connect.driver.exceptions import ClickHouseError +from clickhouse_connect.driver.external import ExternalData from nmp.intake.repository.clickhouse.tables import ClickHouseTable, qualified_table from nmp.intake.spans.clickhouse_client import ClickHouseSpanClient logger = logging.getLogger(__name__) +@dataclass(frozen=True) +class ClickHouseExternalData: + """One typed external-data payload used by a repository query.""" + + file_name: str + data: bytes + fmt: str + structure: str + + @dataclass(frozen=True) class ClickHouseQuery: """One named, parameterized ClickHouse read statement.""" @@ -25,6 +36,7 @@ class ClickHouseQuery: name: str statement: str parameters: Mapping[str, object] = field(default_factory=dict) + external_data: ClickHouseExternalData | None = None def bind(self, **parameters: object) -> ClickHouseQuery: """Return a copy with additional bound parameters.""" @@ -33,6 +45,7 @@ def bind(self, **parameters: object) -> ClickHouseQuery: name=self.name, statement=self.statement, parameters={**self.parameters, **parameters}, + external_data=self.external_data, ) @@ -74,10 +87,23 @@ def table(self, table: ClickHouseTable) -> str: async def fetch_all(self, query: ClickHouseQuery) -> list[dict[str, Any]]: started_at = perf_counter() try: - result = await self._client.query( - query.statement, - parameters=dict(query.parameters), - ) + if query.external_data is None: + result = await self._client.query( + query.statement, + parameters=dict(query.parameters), + ) + else: + external_data = query.external_data + result = await self._client.query( + query.statement, + parameters=dict(query.parameters), + external_data=ExternalData( + file_name=external_data.file_name, + data=external_data.data, + fmt=external_data.fmt, + structure=external_data.structure, + ), + ) except ClickHouseError as exc: logger.exception("ClickHouse repository query failed", extra={"query_name": query.name}) raise ClickHouseQueryError(query.name) from exc diff --git a/services/intake/src/nmp/intake/spans/span_repository.py b/services/intake/src/nmp/intake/repository/clickhouse/span.py similarity index 79% rename from services/intake/src/nmp/intake/spans/span_repository.py rename to services/intake/src/nmp/intake/repository/clickhouse/span.py index 111ffb6183..c3651d1336 100644 --- a/services/intake/src/nmp/intake/spans/span_repository.py +++ b/services/intake/src/nmp/intake/repository/clickhouse/span.py @@ -10,7 +10,9 @@ from typing import Any from nmp.common.api.common import PaginatedResult -from nmp.intake.spans.clickhouse_client import ClickHouseSpanClient +from nmp.intake.repository.clickhouse.executor import ClickHouseExecutor, ClickHouseInsert, ClickHouseQuery +from nmp.intake.repository.clickhouse.tables import ClickHouseTable +from nmp.intake.repository.span import SpanRepository from nmp.intake.spans.domain import IntakeResponseMode, IntakeSpan, SpanGroup, SpanListFilter from nmp.intake.spans.span_attribute_catalog import where_clause from nmp.intake.spans.storage import ( @@ -18,7 +20,6 @@ make_pagination, normalize_span_kind, normalize_span_status, - result_rows, text_query_parameters, text_select_for_mode, ) @@ -65,13 +66,20 @@ class _GroupExpression: required_sql: str -class SpanRepository: - def __init__(self, client: ClickHouseSpanClient) -> None: - self._client = client +class ClickHouseSpanRepository(SpanRepository): + def __init__(self, executor: ClickHouseExecutor) -> None: + self._executor = executor async def save_spans(self, spans: list[IntakeSpan]) -> None: rows = [dict_to_row(_span_to_row(span), SPAN_INSERT_COLUMNS) for span in spans] - await self._client.insert("spans", rows, column_names=SPAN_INSERT_COLUMNS) + await self._executor.insert( + ClickHouseInsert( + name="spans.save", + table=ClickHouseTable.SPANS, + rows=rows, + column_names=SPAN_INSERT_COLUMNS, + ) + ) async def list_spans( self, @@ -83,11 +91,17 @@ async def list_spans( mode: IntakeResponseMode, ) -> PaginatedResult[IntakeSpan]: where_sql, parameters = _span_where(filters) - table = self._client.table("spans") - total_result = await self._client.query( - f"SELECT count() FROM {table} FINAL WHERE {where_sql}", parameters=parameters + table = self._executor.table(ClickHouseTable.SPANS) + total_results = int( + await self._executor.fetch_scalar( + ClickHouseQuery( + name="spans.list.count", + statement=f"SELECT count() FROM {table} FINAL WHERE {where_sql}", + parameters=parameters, + ) + ) + or 0 ) - total_results = int(total_result.result_rows[0][0]) offset = (page - 1) * page_size columns_sql = _span_select_columns(mode=mode) rows_parameters: dict[str, Any] = { @@ -96,17 +110,19 @@ async def list_spans( "limit": page_size, "offset": offset, } - rows_result = await self._client.query( - f""" - SELECT {columns_sql} - FROM {table} FINAL - WHERE {where_sql} - ORDER BY {_order_by(sort)} - LIMIT %(limit)s OFFSET %(offset)s - """, - parameters=rows_parameters, + rows = await self._executor.fetch_all( + ClickHouseQuery( + name="spans.list.rows", + statement=f""" + SELECT {columns_sql} + FROM {table} FINAL + WHERE {where_sql} + ORDER BY {_order_by(sort)} + LIMIT %(limit)s OFFSET %(offset)s + """, + parameters=rows_parameters, + ) ) - rows = result_rows(rows_result) spans = _rows_to_spans(rows) return PaginatedResult( data=spans, @@ -131,7 +147,7 @@ async def list_span_groups( if required_sql: where_sql = f"{where_sql} AND {required_sql}" - table = self._client.table("spans") + table = self._executor.table(ClickHouseTable.SPANS) select_sql = ", ".join(expression.select_sql for expression in group_expressions) group_sql = ", ".join(expression.group_sql for expression in group_expressions) grouped_sql = f""" @@ -141,22 +157,29 @@ async def list_span_groups( GROUP BY {group_sql} """ - total_result = await self._client.query( - f"SELECT count() FROM ({grouped_sql}) AS span_groups", - parameters=parameters, + total_results = int( + await self._executor.fetch_scalar( + ClickHouseQuery( + name="spans.list_groups.count", + statement=f"SELECT count() FROM ({grouped_sql}) AS span_groups", + parameters=parameters, + ) + ) + or 0 ) - total_results = int(total_result.result_rows[0][0]) offset = (page - 1) * page_size - rows_result = await self._client.query( - f""" - SELECT * - FROM ({grouped_sql}) AS span_groups - ORDER BY {_group_order_by(sort, group_by)} - LIMIT %(limit)s OFFSET %(offset)s - """, - parameters={**parameters, "limit": page_size, "offset": offset}, + rows = await self._executor.fetch_all( + ClickHouseQuery( + name="spans.list_groups.rows", + statement=f""" + SELECT * + FROM ({grouped_sql}) AS span_groups + ORDER BY {_group_order_by(sort, group_by)} + LIMIT %(limit)s OFFSET %(offset)s + """, + parameters={**parameters, "limit": page_size, "offset": offset}, + ) ) - rows = result_rows(rows_result) groups = [_row_to_group(row, group_by=group_by) for row in rows] return PaginatedResult( data=groups, @@ -167,16 +190,18 @@ async def list_span_groups( async def get_span(self, *, workspace: str, span_id: str) -> IntakeSpan | None: columns_sql = ", ".join(SPAN_COLUMNS) - result = await self._client.query( - f""" - SELECT {columns_sql} - FROM {self._client.table("spans")} FINAL - WHERE workspace = %(workspace)s AND external_span_id = %(span_id)s AND is_deleted = 0 - LIMIT 1 - """, - parameters={"workspace": workspace, "span_id": span_id}, + rows = await self._executor.fetch_all( + ClickHouseQuery( + name="spans.get", + statement=f""" + SELECT {columns_sql} + FROM {self._executor.table(ClickHouseTable.SPANS)} FINAL + WHERE workspace = %(workspace)s AND external_span_id = %(span_id)s AND is_deleted = 0 + LIMIT 1 + """, + parameters={"workspace": workspace, "span_id": span_id}, + ) ) - rows = result_rows(result) if not rows: return None return _row_to_span(rows[0]) diff --git a/services/intake/src/nmp/intake/spans/trace_repository.py b/services/intake/src/nmp/intake/repository/clickhouse/trace.py similarity index 82% rename from services/intake/src/nmp/intake/spans/trace_repository.py rename to services/intake/src/nmp/intake/repository/clickhouse/trace.py index 8a1d75e6a6..6cfbd6bafa 100644 --- a/services/intake/src/nmp/intake/spans/trace_repository.py +++ b/services/intake/src/nmp/intake/repository/clickhouse/trace.py @@ -9,9 +9,10 @@ from datetime import datetime, timezone from typing import Any -from clickhouse_connect.driver.external import ExternalData from nmp.common.api.common import PaginatedResult -from nmp.intake.spans.clickhouse_client import ClickHouseSpanClient +from nmp.intake.repository.clickhouse.executor import ClickHouseExecutor, ClickHouseExternalData, ClickHouseQuery +from nmp.intake.repository.clickhouse.tables import ClickHouseTable +from nmp.intake.repository.trace import TraceRepository from nmp.intake.spans.domain import IntakeTrace, TraceListFilter, TraceMode from nmp.intake.spans.span_attribute_catalog import SpanAttributeField, spec_for_field from nmp.intake.spans.span_rollups import METRIC_ATTRIBUTE_FIELDS, metric_aggregate_columns @@ -20,7 +21,6 @@ int_or_none, make_pagination, normalize_span_status, - result_rows, text_query_parameters, text_select_for_mode, ) @@ -77,9 +77,9 @@ _ZERO_DATETIME = datetime.fromtimestamp(0, tz=timezone.utc) -class TraceRepository: - def __init__(self, client: ClickHouseSpanClient) -> None: - self._client = client +class ClickHouseTraceRepository(TraceRepository): + def __init__(self, executor: ClickHouseExecutor) -> None: + self._executor = executor async def list_traces( self, @@ -90,22 +90,27 @@ async def list_traces( sort: str, mode: TraceMode, ) -> PaginatedResult[IntakeTrace]: - trace_index_table = self._client.table("trace_index") - spans_table = self._client.table("spans") + trace_index_table = self._executor.table(ClickHouseTable.TRACE_INDEX) + spans_table = self._executor.table(ClickHouseTable.SPANS) count_trace_index_sql, parameters = _trace_index_sql( trace_index_table, filters, mode="summary", ) - total_result = await self._client.query( - f""" - SELECT count() - FROM ({count_trace_index_sql}) AS traces - """, - parameters=parameters, + total_results = int( + await self._executor.fetch_scalar( + ClickHouseQuery( + name="traces.list.count", + statement=f""" + SELECT count() + FROM ({count_trace_index_sql}) AS traces + """, + parameters=parameters, + ) + ) + or 0 ) - total_results = int(total_result.result_rows[0][0]) offset = (page - 1) * page_size trace_index_sql, row_parameters = _trace_index_sql( @@ -119,16 +124,18 @@ async def list_traces( mode=mode, sort=sort, ) - rows_result = await self._client.query( - rows_sql, - parameters={ - **row_parameters, - **rows_parameters, - "limit": page_size, - "offset": offset, - }, + rows = await self._executor.fetch_all( + ClickHouseQuery( + name="traces.list.rows", + statement=rows_sql, + parameters={ + **row_parameters, + **rows_parameters, + "limit": page_size, + "offset": offset, + }, + ) ) - rows = result_rows(rows_result) traces = [_row_to_trace(row) for row in rows] return PaginatedResult( data=traces, @@ -164,41 +171,45 @@ async def latest_trace_started_at_by_group( if not pairs: return {} - trace_index_table = self._client.table("trace_index") - result = await self._client.query( - f""" - WITH - refs AS ( - SELECT group_id, trace_id - FROM trace_refs - ), - traces AS ( - SELECT - trace_roots.trace_id AS id, - trace_roots.root_started_at AS started_at - FROM {trace_index_table} AS trace_roots FINAL - WHERE trace_roots.workspace = %(workspace)s - AND trace_roots.is_deleted = 0 - AND trace_roots.trace_id IN (SELECT trace_id FROM refs) - ORDER BY trace_roots.root_started_at ASC, trace_roots.root_span_id ASC - LIMIT 1 BY trace_roots.workspace, trace_roots.source_format, trace_roots.trace_id - ) - SELECT refs.group_id, max(traces.started_at) AS started_at - FROM refs - INNER JOIN traces ON traces.id = refs.trace_id - GROUP BY refs.group_id - """, - parameters={"workspace": workspace}, - external_data=ExternalData( - file_name="trace_refs.jsonl", - data=b"\n".join( - json.dumps({"group_id": group_id, "trace_id": trace_id}).encode() for group_id, trace_id in pairs + trace_index_table = self._executor.table(ClickHouseTable.TRACE_INDEX) + rows = await self._executor.fetch_all( + ClickHouseQuery( + name="traces.latest_started_at_by_group", + statement=f""" + WITH + refs AS ( + SELECT group_id, trace_id + FROM trace_refs ), - fmt="JSONEachRow", - structure="group_id String, trace_id String", - ), + traces AS ( + SELECT + trace_roots.trace_id AS id, + trace_roots.root_started_at AS started_at + FROM {trace_index_table} AS trace_roots FINAL + WHERE trace_roots.workspace = %(workspace)s + AND trace_roots.is_deleted = 0 + AND trace_roots.trace_id IN (SELECT trace_id FROM refs) + ORDER BY trace_roots.root_started_at ASC, trace_roots.root_span_id ASC + LIMIT 1 BY trace_roots.workspace, trace_roots.source_format, trace_roots.trace_id + ) + SELECT refs.group_id, max(traces.started_at) AS started_at + FROM refs + INNER JOIN traces ON traces.id = refs.trace_id + GROUP BY refs.group_id + """, + parameters={"workspace": workspace}, + external_data=ClickHouseExternalData( + file_name="trace_refs.jsonl", + data=b"\n".join( + json.dumps({"group_id": group_id, "trace_id": trace_id}).encode() + for group_id, trace_id in pairs + ), + fmt="JSONEachRow", + structure="group_id String, trace_id String", + ), + ) ) - return {str(group_id): started_at for group_id, started_at in result.result_rows} + return {str(row["group_id"]): row["started_at"] for row in rows} def _trace_rows_sql( diff --git a/services/intake/src/nmp/intake/repository/span.py b/services/intake/src/nmp/intake/repository/span.py new file mode 100644 index 0000000000..9e46b45bec --- /dev/null +++ b/services/intake/src/nmp/intake/repository/span.py @@ -0,0 +1,45 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Repository interface for Intake spans.""" + +from abc import ABC, abstractmethod + +from nmp.common.api.common import PaginatedResult +from nmp.intake.spans.domain import IntakeResponseMode, IntakeSpan, SpanGroup, SpanListFilter + + +class SpanRepository(ABC): + """Domain-facing interface for span persistence.""" + + @abstractmethod + async def save_spans(self, spans: list[IntakeSpan]) -> None: + pass + + @abstractmethod + async def list_spans( + self, + *, + filters: SpanListFilter, + page: int, + page_size: int, + sort: str, + mode: IntakeResponseMode, + ) -> PaginatedResult[IntakeSpan]: + pass + + @abstractmethod + async def list_span_groups( + self, + *, + filters: SpanListFilter, + group_by: list[str], + page: int, + page_size: int, + sort: str, + ) -> PaginatedResult[SpanGroup]: + pass + + @abstractmethod + async def get_span(self, *, workspace: str, span_id: str) -> IntakeSpan | None: + pass diff --git a/services/intake/src/nmp/intake/repository/trace.py b/services/intake/src/nmp/intake/repository/trace.py new file mode 100644 index 0000000000..93556c21c5 --- /dev/null +++ b/services/intake/src/nmp/intake/repository/trace.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Repository interface for Intake traces.""" + +from abc import ABC, abstractmethod +from datetime import datetime + +from nmp.common.api.common import PaginatedResult +from nmp.intake.spans.domain import IntakeTrace, TraceListFilter, TraceMode + + +class TraceRepository(ABC): + """Domain-facing interface for trace reads.""" + + @abstractmethod + async def list_traces( + self, + *, + filters: TraceListFilter, + page: int, + page_size: int, + sort: str, + mode: TraceMode, + ) -> PaginatedResult[IntakeTrace]: + pass + + @abstractmethod + async def get_trace(self, *, workspace: str, trace_id: str, mode: TraceMode) -> IntakeTrace | None: + pass + + @abstractmethod + async def latest_trace_started_at_by_group( + self, + *, + workspace: str, + trace_refs_by_group: dict[str, list[str]], + ) -> dict[str, datetime]: + pass diff --git a/services/intake/src/nmp/intake/spans/api/dependencies.py b/services/intake/src/nmp/intake/spans/api/dependencies.py index 9b2b01bac6..512b3c1ec9 100644 --- a/services/intake/src/nmp/intake/spans/api/dependencies.py +++ b/services/intake/src/nmp/intake/spans/api/dependencies.py @@ -13,12 +13,14 @@ from nmp.intake.repository.clickhouse.evaluator_results import ClickHouseEvaluatorResultsRepository from nmp.intake.repository.clickhouse.executor import ClickHouseExecutor from nmp.intake.repository.clickhouse.session import ClickHouseSessionRepository +from nmp.intake.repository.clickhouse.span import ClickHouseSpanRepository +from nmp.intake.repository.clickhouse.trace import ClickHouseTraceRepository from nmp.intake.repository.evaluator_results import EvaluatorResultsRepository from nmp.intake.repository.session import SessionRepository +from nmp.intake.repository.span import SpanRepository +from nmp.intake.repository.trace import TraceRepository from nmp.intake.spans.clickhouse_client import ClickHouseSpanClient, get_clickhouse_client from nmp.intake.spans.service import IntakeSpansService -from nmp.intake.spans.span_repository import SpanRepository -from nmp.intake.spans.trace_repository import TraceRepository async def require_workspace_access( @@ -50,22 +52,22 @@ def validate_list_query_params(request: Request, additional_params: set[str] | N ) -def get_span_repository( +def get_clickhouse_executor( client: Annotated[ClickHouseSpanClient, Depends(get_clickhouse_client)], +) -> ClickHouseExecutor: + return ClickHouseExecutor(client) + + +def get_span_repository( + executor: Annotated[ClickHouseExecutor, Depends(get_clickhouse_executor)], ) -> SpanRepository: - return SpanRepository(client) + return ClickHouseSpanRepository(executor) def get_trace_repository( - client: Annotated[ClickHouseSpanClient, Depends(get_clickhouse_client)], + executor: Annotated[ClickHouseExecutor, Depends(get_clickhouse_executor)], ) -> TraceRepository: - return TraceRepository(client) - - -def get_clickhouse_executor( - client: Annotated[ClickHouseSpanClient, Depends(get_clickhouse_client)], -) -> ClickHouseExecutor: - return ClickHouseExecutor(client) + return ClickHouseTraceRepository(executor) def get_session_repository( diff --git a/services/intake/src/nmp/intake/spans/service.py b/services/intake/src/nmp/intake/spans/service.py index ba26ee2061..fb536a79ed 100644 --- a/services/intake/src/nmp/intake/spans/service.py +++ b/services/intake/src/nmp/intake/spans/service.py @@ -11,6 +11,8 @@ from nmp.intake.repository.annotations import AnnotationsRepository from nmp.intake.repository.evaluator_results import EvaluatorResultsRepository from nmp.intake.repository.session import SessionRepository +from nmp.intake.repository.span import SpanRepository +from nmp.intake.repository.trace import TraceRepository from nmp.intake.spans.domain import ( Annotation, AnnotationListFilter, @@ -26,8 +28,6 @@ TraceListFilter, TraceMode, ) -from nmp.intake.spans.span_repository import SpanRepository -from nmp.intake.spans.trace_repository import TraceRepository class SpanNotFoundError(Exception): diff --git a/services/intake/tests/test_clickhouse_architecture.py b/services/intake/tests/test_clickhouse_architecture.py index 4415bd5173..5134d29840 100644 --- a/services/intake/tests/test_clickhouse_architecture.py +++ b/services/intake/tests/test_clickhouse_architecture.py @@ -15,8 +15,6 @@ "repository/clickhouse/executor.py", "service.py", "spans/api/dependencies.py", - "spans/span_repository.py", - "spans/trace_repository.py", } diff --git a/services/intake/tests/test_clickhouse_executor.py b/services/intake/tests/test_clickhouse_executor.py index 14fe3ed8ba..0f82d8c96e 100644 --- a/services/intake/tests/test_clickhouse_executor.py +++ b/services/intake/tests/test_clickhouse_executor.py @@ -8,8 +8,10 @@ import pytest from clickhouse_connect.driver.exceptions import ClickHouseError +from clickhouse_connect.driver.external import ExternalData from nmp.intake.repository.clickhouse.executor import ( ClickHouseExecutor, + ClickHouseExternalData, ClickHouseInsert, ClickHouseInsertError, ClickHouseQuery, @@ -26,11 +28,19 @@ def __init__(self, *, error: ClickHouseError | None = None) -> None: self.error = error self.statements: list[str] = [] self.parameters: list[dict[str, Any]] = [] + self.external_data: list[object | None] = [] self.inserts: list[tuple[str, list[tuple[object, ...]], list[str]]] = [] - async def query(self, statement: str, *, parameters: dict[str, Any]) -> SimpleNamespace: + async def query( + self, + statement: str, + *, + parameters: dict[str, Any], + external_data: object | None = None, + ) -> SimpleNamespace: self.statements.append(statement) self.parameters.append(parameters) + self.external_data.append(external_data) if self.error is not None: raise self.error return SimpleNamespace( @@ -98,6 +108,29 @@ async def test_executor_returns_first_scalar() -> None: assert scalar == "session-a" +@pytest.mark.asyncio +async def test_executor_builds_driver_external_data_inside_boundary() -> None: + client = _Client() + query = ClickHouseQuery( + name="traces.latest_started_at_by_group", + statement="SELECT * FROM trace_refs", + external_data=ClickHouseExternalData( + file_name="trace_refs.jsonl", + data=b'{"group_id":"group-a","trace_id":"trace-a"}', + fmt="JSONEachRow", + structure="group_id String, trace_id String", + ), + ) + + await _executor(client).fetch_all(query) + + external_data = cast(ExternalData, client.external_data[0]) + assert external_data.query_params == { + "trace_refs_format": "JSONEachRow", + "trace_refs_structure": "group_id String, trace_id String", + } + + @pytest.mark.asyncio async def test_executor_translates_clickhouse_errors_without_sql_details() -> None: query = ClickHouseQuery( diff --git a/services/intake/tests/test_spans_clickhouse_repository.py b/services/intake/tests/test_spans_clickhouse_repository.py index eec4cffcec..cd83c976ab 100644 --- a/services/intake/tests/test_spans_clickhouse_repository.py +++ b/services/intake/tests/test_spans_clickhouse_repository.py @@ -4,43 +4,50 @@ """Span repository tests.""" from datetime import datetime, timezone -from typing import cast import pytest +from nmp.intake.repository.clickhouse.executor import ClickHouseExecutor, ClickHouseQuery +from nmp.intake.repository.clickhouse.span import ( + SPAN_COLUMNS, + SPAN_GROUP_COLUMN_FIELDS, + ClickHouseSpanRepository, + _order_by, +) +from nmp.intake.repository.clickhouse.tables import ClickHouseTable from nmp.intake.spans.api.spans_schemas import SpanGroupBy -from nmp.intake.spans.clickhouse_client import ClickHouseSpanClient -from nmp.intake.spans.domain import SpanListFilter -from nmp.intake.spans.span_repository import SPAN_COLUMNS, SPAN_GROUP_COLUMN_FIELDS, SpanRepository, _order_by +from nmp.intake.spans.domain import SpanListFilter, SpanStatus from nmp.intake.spans.storage import make_pagination class _QueryResult: def __init__(self, rows: list[tuple[object, ...]], columns: list[str] | None = None) -> None: self.result_rows = rows - self.column_names = columns or [] + self.column_names = columns or (["count()"] if rows and len(rows[0]) == 1 else []) -class _Client: +class _Client(ClickHouseExecutor): def __init__(self, query_results: list[_QueryResult] | None = None) -> None: self.queries: list[str] = [] self.parameters: list[dict[str, object]] = [] self.query_results = query_results or [] - def table(self, name: str) -> str: - return name + def table(self, table: ClickHouseTable) -> str: + return table.value - async def query(self, query: str, *, parameters: dict[str, object]) -> _QueryResult: - self.queries.append(query) - self.parameters.append(parameters) + async def fetch_all(self, query: ClickHouseQuery) -> list[dict[str, object]]: + self.queries.append(query.statement) + self.parameters.append(dict(query.parameters)) if self.query_results: - return self.query_results.pop(0) - if query.lstrip().startswith("SELECT count()"): - return _QueryResult([(0,)]) - return _QueryResult([]) + result = self.query_results.pop(0) + elif query.statement.lstrip().startswith("SELECT count()"): + result = _QueryResult([(0,)], ["count()"]) + else: + result = _QueryResult([]) + return [dict(zip(result.column_names, row, strict=True)) for row in result.result_rows] -def _repository(client: _Client) -> SpanRepository: - return SpanRepository(cast(ClickHouseSpanClient, client)) +def _repository(client: _Client) -> ClickHouseSpanRepository: + return ClickHouseSpanRepository(client) def test_order_by_whitelists_supported_span_sort_keys(): @@ -197,7 +204,7 @@ async def test_list_span_groups_reuses_span_filters(): repository = _repository(client) await repository.list_span_groups( - filters=SpanListFilter(workspace="workspace-a", status="error"), + filters=SpanListFilter(workspace="workspace-a", status=SpanStatus.ERROR), group_by=["trace_id"], page=1, page_size=10, diff --git a/services/intake/tests/test_traces_clickhouse_repository.py b/services/intake/tests/test_traces_clickhouse_repository.py index 7c6782da1f..929ff5fb0b 100644 --- a/services/intake/tests/test_traces_clickhouse_repository.py +++ b/services/intake/tests/test_traces_clickhouse_repository.py @@ -5,49 +5,45 @@ import json from datetime import datetime, timedelta, timezone -from typing import cast import pytest -from nmp.intake.spans.clickhouse_client import ClickHouseSpanClient +from nmp.intake.repository.clickhouse.executor import ClickHouseExecutor, ClickHouseExternalData, ClickHouseQuery +from nmp.intake.repository.clickhouse.tables import ClickHouseTable +from nmp.intake.repository.clickhouse.trace import TRACE_COLUMNS, ClickHouseTraceRepository, _order_by from nmp.intake.spans.domain import TraceListFilter -from nmp.intake.spans.trace_repository import TRACE_COLUMNS, TraceRepository, _order_by class _QueryResult: def __init__(self, rows: list[tuple[object, ...]], columns: list[str] | None = None) -> None: self.result_rows = rows - self.column_names = columns or [] + self.column_names = columns or (["count()"] if rows and len(rows[0]) == 1 else []) -class _Client: +class _Client(ClickHouseExecutor): def __init__(self, query_results: list[_QueryResult] | None = None) -> None: self.queries: list[str] = [] self.parameters: list[dict[str, object]] = [] - self.external_data: list[object | None] = [] + self.external_data: list[ClickHouseExternalData | None] = [] self.query_results = query_results or [] - def table(self, name: str) -> str: - return name - - async def query( - self, - query: str, - *, - parameters: dict[str, object], - external_data: object | None = None, - ) -> _QueryResult: - self.queries.append(query) - self.parameters.append(parameters) - self.external_data.append(external_data) + def table(self, table: ClickHouseTable) -> str: + return table.value + + async def fetch_all(self, query: ClickHouseQuery) -> list[dict[str, object]]: + self.queries.append(query.statement) + self.parameters.append(dict(query.parameters)) + self.external_data.append(query.external_data) if self.query_results: - return self.query_results.pop(0) - if query.lstrip().startswith("SELECT count()"): - return _QueryResult([(0,)]) - return _QueryResult([]) + result = self.query_results.pop(0) + elif query.statement.lstrip().startswith("SELECT count()"): + result = _QueryResult([(0,)], ["count()"]) + else: + result = _QueryResult([]) + return [dict(zip(result.column_names, row, strict=True)) for row in result.result_rows] -def _repository(client: _Client) -> TraceRepository: - return TraceRepository(cast(ClickHouseSpanClient, client)) +def _repository(client: _Client) -> ClickHouseTraceRepository: + return ClickHouseTraceRepository(client) def test_order_by_whitelists_supported_trace_sort_keys(): @@ -95,7 +91,7 @@ async def test_summary_mode_reads_root_spans_without_metric_aggregates(): @pytest.mark.asyncio async def test_latest_trace_started_at_by_group_aggregates_all_references_in_one_query(): latest = datetime(2026, 1, 2, tzinfo=timezone.utc) - client = _Client(query_results=[_QueryResult([("insight-a", latest)])]) + client = _Client(query_results=[_QueryResult([("insight-a", latest)], ["group_id", "started_at"])]) repository = _repository(client) result = await repository.latest_trace_started_at_by_group( @@ -115,11 +111,9 @@ async def test_latest_trace_started_at_by_group_aggregates_all_references_in_one assert client.parameters[0] == {"workspace": "workspace-a"} external_data = client.external_data[0] assert external_data is not None - assert external_data.query_params == { - "trace_refs_format": "JSONEachRow", - "trace_refs_structure": "group_id String, trace_id String", - } - assert [json.loads(line) for line in external_data.form_data["trace_refs"][1].splitlines()] == [ + assert external_data.fmt == "JSONEachRow" + assert external_data.structure == "group_id String, trace_id String" + assert [json.loads(line) for line in external_data.data.splitlines()] == [ {"group_id": "insight-a", "trace_id": "trace-old"}, {"group_id": "insight-a", "trace_id": "trace-new"}, {"group_id": "insight-missing", "trace_id": "trace-missing"},