refactor(intake): complete ClickHouse repository migration

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

services/intake/src/nmp/intake/spans/annotations_repository.py → 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]]:

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.

services/intake/src/nmp/intake/spans/evaluator_results_repository.py → 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]]: