feat(BA-4667): Add GraphQL API for prometheus query preset operations

changes/9643.feature.md

@@ -0,0 +1 @@
+Add GraphQL API for prometheus query preset admin operations and execute query

src/ai/backend/manager/api/gql/prometheus_query_preset/__init__.py

@@ -0,0 +1,60 @@
+"""Prometheus query preset GraphQL API package."""
+
+from .resolver import (
+    admin_create_prometheus_query_preset,
+    admin_delete_prometheus_query_preset,
+    admin_modify_prometheus_query_preset,
+    admin_prometheus_query_preset,
+    admin_prometheus_query_presets,
+    prometheus_query_preset_result,
+)
+from .types import (
+    CreatePrometheusQueryPresetInput,
+    CreatePrometheusQueryPresetPayload,
+    DeletePrometheusQueryPresetPayload,
+    MetricLabelEntryGQL,
+    MetricLabelEntryInput,
+    MetricResultGQL,
+    MetricResultValueGQL,
+    ModifyPrometheusQueryPresetInput,
+    ModifyPrometheusQueryPresetPayload,
+    PrometheusPresetOptionsGQL,
+    PrometheusQueryPresetConnection,
+    PrometheusQueryPresetEdge,
+    PrometheusQueryPresetFilter,
+    PrometheusQueryPresetGQL,
+    PrometheusQueryPresetOrderBy,
+    PrometheusQueryPresetOrderField,
+    PrometheusQueryResultGQL,
+    QueryTimeRangeInput,
+)
+
+__all__ = [
+    # Queries
+    "admin_prometheus_query_preset",
+    "admin_prometheus_query_presets",
+    "prometheus_query_preset_result",
+    # Mutations
+    "admin_create_prometheus_query_preset",
+    "admin_modify_prometheus_query_preset",
+    "admin_delete_prometheus_query_preset",
+    # Types
+    "PrometheusQueryPresetGQL",
+    "PrometheusQueryPresetEdge",
+    "PrometheusQueryPresetConnection",
+    "PrometheusQueryPresetFilter",
+    "PrometheusQueryPresetOrderBy",
+    "PrometheusQueryPresetOrderField",
+    "PrometheusPresetOptionsGQL",
+    "MetricLabelEntryGQL",
+    "MetricResultValueGQL",
+    "MetricResultGQL",
+    "PrometheusQueryResultGQL",
+    "CreatePrometheusQueryPresetInput",
+    "ModifyPrometheusQueryPresetInput",
+    "QueryTimeRangeInput",
+    "MetricLabelEntryInput",
+    "CreatePrometheusQueryPresetPayload",
+    "ModifyPrometheusQueryPresetPayload",
+    "DeletePrometheusQueryPresetPayload",
+]

src/ai/backend/manager/api/gql/prometheus_query_preset/fetcher/__init__.py

@@ -0,0 +1,13 @@
+"""Prometheus query preset GQL fetcher functions."""
+
+from .preset import (
+    fetch_admin_prometheus_query_preset,
+    fetch_admin_prometheus_query_presets,
+    fetch_prometheus_query_preset_result,
+)
+
+__all__ = [
+    "fetch_admin_prometheus_query_preset",
+    "fetch_admin_prometheus_query_presets",
+    "fetch_prometheus_query_preset_result",
+]

src/ai/backend/manager/api/gql/prometheus_query_preset/fetcher/preset.py

@@ -0,0 +1,144 @@
+"""Prometheus query preset GQL data fetcher functions."""
+
+from __future__ import annotations
+
+from functools import lru_cache
+from uuid import UUID
+
+from strawberry import Info
+from strawberry.relay import PageInfo
+
+from ai.backend.common.dto.clients.prometheus.request import QueryTimeRange
+from ai.backend.manager.api.gql.adapter import PaginationOptions, PaginationSpec
+from ai.backend.manager.api.gql.base import encode_cursor
+from ai.backend.manager.api.gql.prometheus_query_preset.types import (
+    MetricLabelEntryGQL,
+    MetricResultGQL,
+    MetricResultValueGQL,
+    PrometheusQueryPresetConnection,
+    PrometheusQueryPresetEdge,
+    PrometheusQueryPresetFilter,
+    PrometheusQueryPresetGQL,
+    PrometheusQueryPresetOrderBy,
+    PrometheusQueryResultGQL,
+)
+from ai.backend.manager.api.gql.types import StrawberryGQLContext
+from ai.backend.manager.data.prometheus_query_preset import ExecutePresetOptions
+from ai.backend.manager.models.prometheus_query_preset import PrometheusQueryPresetRow
+from ai.backend.manager.repositories.prometheus_query_preset.options import (
+    PrometheusQueryPresetConditions,
+    PrometheusQueryPresetOrders,
+)
+from ai.backend.manager.services.prometheus_query_preset.actions import (
+    ExecutePresetAction,
+    GetPresetAction,
+    SearchPresetsAction,
+)
+
+
+@lru_cache(maxsize=1)
+def get_preset_pagination_spec() -> PaginationSpec:
+    return PaginationSpec(
+        forward_order=PrometheusQueryPresetOrders.created_at(ascending=False),
+        backward_order=PrometheusQueryPresetOrders.created_at(ascending=True),
+        forward_condition_factory=PrometheusQueryPresetConditions.by_cursor_forward,
+        backward_condition_factory=PrometheusQueryPresetConditions.by_cursor_backward,
+        tiebreaker_order=PrometheusQueryPresetRow.id.asc(),
+    )
+
+
+async def fetch_admin_prometheus_query_preset(
+    info: Info[StrawberryGQLContext],
+    preset_id: UUID,
+) -> PrometheusQueryPresetGQL:
+    processors = info.context.processors
+    action_result = await processors.prometheus_query_preset.get_preset.wait_for_complete(
+        GetPresetAction(preset_id=preset_id)
+    )
+    return PrometheusQueryPresetGQL.from_data(action_result.preset)
+
+
+async def fetch_admin_prometheus_query_presets(
+    info: Info[StrawberryGQLContext],
+    filter: PrometheusQueryPresetFilter | None = None,
+    order_by: list[PrometheusQueryPresetOrderBy] | None = None,
+    before: str | None = None,
+    after: str | None = None,
+    first: int | None = None,
+    last: int | None = None,
+    limit: int | None = None,
+    offset: int | None = None,
+) -> PrometheusQueryPresetConnection:
+    processors = info.context.processors
+
+    querier = info.context.gql_adapter.build_querier(
+        PaginationOptions(
+            first=first,
+            after=after,
+            last=last,
+            before=before,
+            limit=limit,
+            offset=offset,
+        ),
+        get_preset_pagination_spec(),
+        filter=filter,
+        order_by=order_by,
+        base_conditions=None,
+    )
+
+    action_result = await processors.prometheus_query_preset.search_presets.wait_for_complete(
+        SearchPresetsAction(querier=querier)
+    )
+
+    nodes = [PrometheusQueryPresetGQL.from_data(data) for data in action_result.items]
+    edges = [
+        PrometheusQueryPresetEdge(node=node, cursor=encode_cursor(str(node.id))) for node in nodes
+    ]
+
+    return PrometheusQueryPresetConnection(
+        edges=edges,
+        page_info=PageInfo(
+            has_next_page=action_result.has_next_page,
+            has_previous_page=action_result.has_previous_page,
+            start_cursor=edges[0].cursor if edges else None,
+            end_cursor=edges[-1].cursor if edges else None,
+        ),
+        count=action_result.total_count,
+    )
+
+
+async def fetch_prometheus_query_preset_result(
+    info: Info[StrawberryGQLContext],
+    preset_id: UUID,
+    options: ExecutePresetOptions,
+    window: str | None,
+    time_range: QueryTimeRange,
+) -> PrometheusQueryResultGQL:
+    processors = info.context.processors
+
+    action_result = await processors.prometheus_query_preset.execute_preset.wait_for_complete(
+        ExecutePresetAction(
+            preset_id=preset_id,
+            options=options,
+            window=window,
+            time_range=time_range,
+        )
+    )
+
+    response = action_result.response
+    result_entries: list[MetricResultGQL] = []
+    for metric_response in response.data.result:
+        metric_labels = [
+            MetricLabelEntryGQL(key=k, value=str(v))
+            for k, v in metric_response.metric.model_dump(exclude_none=True).items()
+        ]
+        values = [
+            MetricResultValueGQL(timestamp=ts, value=val) for ts, val in metric_response.values
+        ]
+        result_entries.append(MetricResultGQL(metric=metric_labels, values=values))
+
+    return PrometheusQueryResultGQL(
+        status=response.status,
+        result_type=response.data.result_type,
+        result=result_entries,
+    )

src/ai/backend/manager/api/gql/prometheus_query_preset/resolver/__init__.py

@@ -0,0 +1,23 @@
+"""Prometheus query preset GQL resolvers."""
+
+from .mutation import (
+    admin_create_prometheus_query_preset,
+    admin_delete_prometheus_query_preset,
+    admin_modify_prometheus_query_preset,
+)
+from .query import (
+    admin_prometheus_query_preset,
+    admin_prometheus_query_presets,
+    prometheus_query_preset_result,
+)
+
+__all__ = [
+    # Queries
+    "admin_prometheus_query_preset",
+    "admin_prometheus_query_presets",
+    "prometheus_query_preset_result",
+    # Mutations
+    "admin_create_prometheus_query_preset",
+    "admin_modify_prometheus_query_preset",
+    "admin_delete_prometheus_query_preset",
+]

src/ai/backend/manager/api/gql/prometheus_query_preset/resolver/mutation.py

@@ -0,0 +1,75 @@
+"""Prometheus query preset GQL mutation resolvers."""
+
+from __future__ import annotations
+
+import uuid
+
+import strawberry
+from strawberry import ID, Info
+
+from ai.backend.manager.api.gql.prometheus_query_preset.types import (
+    CreatePrometheusQueryPresetInput,
+    CreatePrometheusQueryPresetPayload,
+    DeletePrometheusQueryPresetPayload,
+    ModifyPrometheusQueryPresetInput,
+    ModifyPrometheusQueryPresetPayload,
+    PrometheusQueryPresetGQL,
+)
+from ai.backend.manager.api.gql.types import StrawberryGQLContext
+from ai.backend.manager.api.gql.utils import check_admin_only
+from ai.backend.manager.services.prometheus_query_preset.actions import (
+    CreatePresetAction,
+    DeletePresetAction,
+    ModifyPresetAction,
+)
+
+
+@strawberry.mutation(description="Create a new prometheus query preset (admin only).")  # type: ignore[misc]
+async def admin_create_prometheus_query_preset(
+    info: Info[StrawberryGQLContext],
+    input: CreatePrometheusQueryPresetInput,
+) -> CreatePrometheusQueryPresetPayload:
+    check_admin_only()
+    processors = info.context.processors
+
+    action_result = await processors.prometheus_query_preset.create_preset.wait_for_complete(
+        CreatePresetAction(creator=input.to_creator())
+    )
+
+    return CreatePrometheusQueryPresetPayload(
+        preset=PrometheusQueryPresetGQL.from_data(action_result.preset)
+    )
+
+
+@strawberry.mutation(description="Modify an existing prometheus query preset (admin only).")  # type: ignore[misc]
+async def admin_modify_prometheus_query_preset(
+    info: Info[StrawberryGQLContext],
+    id: ID,
+    input: ModifyPrometheusQueryPresetInput,
+) -> ModifyPrometheusQueryPresetPayload:
+    check_admin_only()
+    processors = info.context.processors
+
+    preset_id = uuid.UUID(id)
+    action_result = await processors.prometheus_query_preset.modify_preset.wait_for_complete(
+        ModifyPresetAction(preset_id=preset_id, updater=input.to_updater(preset_id))
+    )
+
+    return ModifyPrometheusQueryPresetPayload(
+        preset=PrometheusQueryPresetGQL.from_data(action_result.preset)
+    )
+
+
+@strawberry.mutation(description="Delete a prometheus query preset (admin only).")  # type: ignore[misc]
+async def admin_delete_prometheus_query_preset(
+    info: Info[StrawberryGQLContext],
+    id: ID,
+) -> DeletePrometheusQueryPresetPayload:
+    check_admin_only()
+    processors = info.context.processors
+
+    await processors.prometheus_query_preset.delete_preset.wait_for_complete(
+        DeletePresetAction(preset_id=uuid.UUID(id))
+    )
+
+    return DeletePrometheusQueryPresetPayload(id=id)