feat(BA-4702): Implement prometheus query preset repository layer

changes/9470.feature.md

@@ -0,0 +1 @@
+Implement prometheus query preset repository layer.

src/ai/backend/common/exception.py

@@ -181,6 +181,7 @@ class ErrorDomain(enum.StrEnum):
     KEYPAIR_RESOURCE_POLICY = "keypair-resource-policy"
     DATABASE = "database"
     USER_RESOURCE_POLICY = "user-resource-policy"
+    PROMETHEUS_QUERY_PRESET = "prometheus-query-preset"
 
     EXTERNAL_SYSTEM = "external-system"  # Errors from external systems
 
@@ -1167,3 +1168,15 @@ def error_code(self) -> ErrorCode:
             operation=ErrorOperation.GENERIC,
             error_detail=ErrorDetail.INTERNAL_ERROR,
         )
+
+
+class PrometheusQueryPresetNotFound(BackendAIError, web.HTTPNotFound):
+    error_type = "https://api.backend.ai/probs/prometheus-query-preset-not-found"
+    error_title = "The prometheus query preset does not exist."
+
+    def error_code(self) -> ErrorCode:
+        return ErrorCode(
+            domain=ErrorDomain.PROMETHEUS_QUERY_PRESET,
+            operation=ErrorOperation.READ,
+            error_detail=ErrorDetail.NOT_FOUND,
+        )

src/ai/backend/common/metrics/metric.py

@@ -423,6 +423,7 @@ class LayerType(enum.StrEnum):
     NOTIFICATION_REPOSITORY = "notification_repository"
     OBJECT_STORAGE_REPOSITORY = "object_storage_repository"
     PERMISSION_CONTROLLER_REPOSITORY = "permission_controller_repository"
+    PROMETHEUS_QUERY_PRESET_REPOSITORY = "prometheus_query_preset_repository"
     PROJECT_RESOURCE_POLICY_REPOSITORY = "project_resource_policy_repository"
     RESERVOIR_REGISTRY_REPOSITORY = "reservoir_registry_repository"
     RESOURCE_PRESET_REPOSITORY = "resource_preset_repository"

src/ai/backend/manager/data/prometheus_query_preset/__init__.py

@@ -1,3 +1,11 @@
-from .types import PrometheusQueryPresetData
+from .types import (
+    PrometheusQueryPresetData,
+    PrometheusQueryPresetListResult,
+    PrometheusQueryPresetModifier,
+)
 
-__all__ = ("PrometheusQueryPresetData",)
+__all__ = (
+    "PrometheusQueryPresetData",
+    "PrometheusQueryPresetListResult",
+    "PrometheusQueryPresetModifier",
+)

src/ai/backend/manager/data/prometheus_query_preset/types.py

@@ -2,10 +2,13 @@
 
 from dataclasses import dataclass, field
 from datetime import datetime
+from typing import Any, override
 from uuid import UUID
 
+from ai.backend.manager.types import OptionalState, PartialModifier, TriState
 
-@dataclass
+
+@dataclass(frozen=True)
 class PrometheusQueryPresetData:
     """Domain model data for prometheus query preset."""
 
@@ -18,3 +21,36 @@ class PrometheusQueryPresetData:
     group_labels: list[str]
     created_at: datetime = field(compare=False)
     updated_at: datetime = field(compare=False)
+
+
+@dataclass
+class PrometheusQueryPresetModifier(PartialModifier):
+    """Modifier for prometheus query preset."""
+
+    name: OptionalState[str] = field(default_factory=OptionalState[str].nop)
+    metric_name: OptionalState[str] = field(default_factory=OptionalState[str].nop)
+    query_template: OptionalState[str] = field(default_factory=OptionalState[str].nop)
+    time_window: TriState[str] = field(default_factory=TriState[str].nop)
+    filter_labels: OptionalState[list[str]] = field(default_factory=OptionalState[list[str]].nop)
+    group_labels: OptionalState[list[str]] = field(default_factory=OptionalState[list[str]].nop)
+
+    @override
+    def fields_to_update(self) -> dict[str, Any]:
+        to_update: dict[str, Any] = {}
+        self.name.update_dict(to_update, "name")
+        self.metric_name.update_dict(to_update, "metric_name")
+        self.query_template.update_dict(to_update, "query_template")
+        self.time_window.update_dict(to_update, "time_window")
+        self.filter_labels.update_dict(to_update, "filter_labels")
+        self.group_labels.update_dict(to_update, "group_labels")
+        return to_update
+
+
+@dataclass
+class PrometheusQueryPresetListResult:
+    """Search result with total count for prometheus query presets."""
+
+    items: list[PrometheusQueryPresetData]
+    total_count: int
+    has_next_page: bool
+    has_previous_page: bool

src/ai/backend/manager/repositories/prometheus_query_preset/__init__.py

@@ -0,0 +1,20 @@
+from __future__ import annotations
+
+from ai.backend.manager.data.prometheus_query_preset import (
+    PrometheusQueryPresetData,
+    PrometheusQueryPresetListResult,
+    PrometheusQueryPresetModifier,
+)
+
+from .creators import PrometheusQueryPresetCreatorSpec
+from .repositories import PrometheusQueryPresetRepositories
+from .repository import PrometheusQueryPresetRepository
+
+__all__ = (
+    "PrometheusQueryPresetCreatorSpec",
+    "PrometheusQueryPresetData",
+    "PrometheusQueryPresetListResult",
+    "PrometheusQueryPresetModifier",
+    "PrometheusQueryPresetRepositories",
+    "PrometheusQueryPresetRepository",
+)

src/ai/backend/manager/repositories/prometheus_query_preset/creators.py

@@ -0,0 +1,37 @@
+"""CreatorSpec implementations for prometheus query preset repository."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import override
+
+from ai.backend.manager.models.prometheus_query_preset.row import (
+    PresetOptions,
+    PrometheusQueryPresetRow,
+)
+from ai.backend.manager.repositories.base import CreatorSpec
+
+
+@dataclass
+class PrometheusQueryPresetCreatorSpec(CreatorSpec[PrometheusQueryPresetRow]):
+    """CreatorSpec for prometheus query preset."""
+
+    name: str
+    metric_name: str
+    query_template: str
+    time_window: str | None
+    filter_labels: list[str]
+    group_labels: list[str]
+
+    @override
+    def build_row(self) -> PrometheusQueryPresetRow:
+        return PrometheusQueryPresetRow(
+            name=self.name,
+            metric_name=self.metric_name,
+            query_template=self.query_template,
+            time_window=self.time_window,
+            options=PresetOptions(
+                filter_labels=self.filter_labels,
+                group_labels=self.group_labels,
+            ),
+        )

src/ai/backend/manager/repositories/prometheus_query_preset/db_source/__init__.py

@@ -0,0 +1,3 @@
+from .db_source import PrometheusQueryPresetDBSource
+
+__all__ = ("PrometheusQueryPresetDBSource",)

src/ai/backend/manager/repositories/prometheus_query_preset/db_source/db_source.py

@@ -0,0 +1,143 @@
+"""Database source for prometheus query preset repository operations."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, cast
+from uuid import UUID
+
+import sqlalchemy as sa
+from sqlalchemy.engine import CursorResult
+
+from ai.backend.common.exception import PrometheusQueryPresetNotFound
+from ai.backend.manager.data.prometheus_query_preset import (
+    PrometheusQueryPresetData,
+    PrometheusQueryPresetListResult,
+)
+from ai.backend.manager.models.prometheus_query_preset import PrometheusQueryPresetRow
+from ai.backend.manager.repositories.base import (
+    BatchQuerier,
+    Creator,
+    execute_batch_querier,
+    execute_creator,
+)
+from ai.backend.manager.repositories.base.updater import Updater, execute_updater
+from ai.backend.manager.repositories.prometheus_query_preset.updaters import (
+    PrometheusQueryPresetUpdaterSpec,
+)
+from ai.backend.manager.types import OptionalState
+
+if TYPE_CHECKING:
+    from sqlalchemy.ext.asyncio import AsyncSession as SASession
+
+    from ai.backend.manager.models.utils import ExtendedAsyncSAEngine
+
+
+__all__ = ("PrometheusQueryPresetDBSource",)
+
+
+class PrometheusQueryPresetDBSource:
+    """Database source for prometheus query preset operations."""
+
+    _db: ExtendedAsyncSAEngine
+
+    def __init__(self, db: ExtendedAsyncSAEngine) -> None:
+        self._db = db
+
+    async def create(
+        self,
+        creator: Creator[PrometheusQueryPresetRow],
+    ) -> PrometheusQueryPresetData:
+        """Creates a new prometheus query preset."""
+        async with self._db.begin_session() as db_sess:
+            result = await execute_creator(db_sess, creator)
+            return result.row.to_data()
+
+    async def _merge_partial_options(
+        self,
+        db_sess: SASession,
+        updater: Updater[PrometheusQueryPresetRow],
+    ) -> Updater[PrometheusQueryPresetRow]:
+        """When only one of filter_labels/group_labels is being updated,
+        fetch the current options to preserve the other field."""
+        updater.spec = cast(PrometheusQueryPresetUpdaterSpec, updater.spec)
+        filter_updating = updater.spec.filter_labels.optional_value() is not None
+        group_updating = updater.spec.group_labels.optional_value() is not None
+
+        # If both are being updated or both are not being updated, no need to merge
+        if filter_updating == group_updating:
+            return updater
+
+        stmt = sa.select(PrometheusQueryPresetRow.options).where(
+            PrometheusQueryPresetRow.id == updater.pk_value
+        )
+        current_options = (await db_sess.execute(stmt)).scalar_one_or_none()
+        if current_options is None:
+            raise PrometheusQueryPresetNotFound(
+                f"Prometheus query preset {updater.pk_value} not found"
+            )
+
+        if filter_updating:
+            updater.spec.group_labels = OptionalState[list[str]].update(
+                list(current_options.group_labels)
+            )
+        if group_updating:
+            updater.spec.filter_labels = OptionalState[list[str]].update(
+                list(current_options.filter_labels)
+            )
+        return updater
+
+    async def update(
+        self,
+        updater: Updater[PrometheusQueryPresetRow],
+    ) -> PrometheusQueryPresetData:
+        """Updates an existing prometheus query preset."""
+        async with self._db.begin_session() as db_sess:
+            updater = await self._merge_partial_options(db_sess, updater)
+            result = await execute_updater(db_sess, updater)
+            if result is None:
+                raise PrometheusQueryPresetNotFound(
+                    f"Prometheus query preset {updater.pk_value} not found"
+                )
+            return result.row.to_data()
+
+    async def delete(self, preset_id: UUID) -> bool:
+        """Deletes a prometheus query preset."""
+        async with self._db.begin_session() as db_sess:
+            stmt = sa.delete(PrometheusQueryPresetRow).where(
+                PrometheusQueryPresetRow.id == preset_id
+            )
+            result = await db_sess.execute(stmt)
+            return cast(CursorResult[Any], result).rowcount > 0
+
+    async def get_by_id(self, preset_id: UUID) -> PrometheusQueryPresetData:
+        """Retrieves a prometheus query preset by ID."""
+        async with self._db.begin_readonly_session_read_committed() as db_sess:
+            row = await db_sess.get(PrometheusQueryPresetRow, preset_id)
+            if row is None:
+                raise PrometheusQueryPresetNotFound(
+                    f"Prometheus query preset {preset_id} not found"
+                )
+            return row.to_data()
+
+    async def search(
+        self,
+        querier: BatchQuerier,
+    ) -> PrometheusQueryPresetListResult:
+        """Searches prometheus query presets with total count."""
+        async with self._db.begin_readonly_session() as db_sess:
+            query = sa.select(PrometheusQueryPresetRow)
+
+            result = await execute_batch_querier(
+                db_sess,
+                query,
+                querier,
+            )
+
+            items = [row.PrometheusQueryPresetRow.to_data() for row in result.rows]
+
+            return PrometheusQueryPresetListResult(
+                items=items,
+                total_count=result.total_count,
+                has_next_page=result.has_next_page,
+                has_previous_page=result.has_previous_page,
+            )