diff --git a/changes/9643.feature.md b/changes/9643.feature.md new file mode 100644 index 00000000000..fe32a50eb5c --- /dev/null +++ b/changes/9643.feature.md @@ -0,0 +1 @@ +Add GraphQL API for prometheus query preset admin operations and execute query diff --git a/src/ai/backend/manager/api/gql/prometheus_query_preset/__init__.py b/src/ai/backend/manager/api/gql/prometheus_query_preset/__init__.py new file mode 100644 index 00000000000..0b9a7766f63 --- /dev/null +++ b/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", +] diff --git a/src/ai/backend/manager/api/gql/prometheus_query_preset/fetcher/__init__.py b/src/ai/backend/manager/api/gql/prometheus_query_preset/fetcher/__init__.py new file mode 100644 index 00000000000..17867058061 --- /dev/null +++ b/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", +] diff --git a/src/ai/backend/manager/api/gql/prometheus_query_preset/fetcher/preset.py b/src/ai/backend/manager/api/gql/prometheus_query_preset/fetcher/preset.py new file mode 100644 index 00000000000..c652b75dd31 --- /dev/null +++ b/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, + ) diff --git a/src/ai/backend/manager/api/gql/prometheus_query_preset/resolver/__init__.py b/src/ai/backend/manager/api/gql/prometheus_query_preset/resolver/__init__.py new file mode 100644 index 00000000000..a0706ac266a --- /dev/null +++ b/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", +] diff --git a/src/ai/backend/manager/api/gql/prometheus_query_preset/resolver/mutation.py b/src/ai/backend/manager/api/gql/prometheus_query_preset/resolver/mutation.py new file mode 100644 index 00000000000..dbb4dfd2bf1 --- /dev/null +++ b/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) diff --git a/src/ai/backend/manager/api/gql/prometheus_query_preset/resolver/query.py b/src/ai/backend/manager/api/gql/prometheus_query_preset/resolver/query.py new file mode 100644 index 00000000000..ed3c118333d --- /dev/null +++ b/src/ai/backend/manager/api/gql/prometheus_query_preset/resolver/query.py @@ -0,0 +1,113 @@ +"""Prometheus query preset GQL query resolvers.""" + +from __future__ import annotations + +import uuid + +import strawberry +from strawberry import ID, Info + +from ai.backend.common.data.filter_specs import StringMatchSpec +from ai.backend.common.exception import PrometheusQueryPresetNotFound +from ai.backend.manager.api.gql.prometheus_query_preset.fetcher import ( + fetch_admin_prometheus_query_preset, + fetch_admin_prometheus_query_presets, + fetch_prometheus_query_preset_result, +) +from ai.backend.manager.api.gql.prometheus_query_preset.types import ( + MetricLabelEntryInput, + PrometheusQueryPresetConnection, + PrometheusQueryPresetFilter, + PrometheusQueryPresetGQL, + PrometheusQueryPresetOrderBy, + PrometheusQueryResultGQL, + QueryTimeRangeInput, +) +from ai.backend.manager.api.gql.types import StrawberryGQLContext +from ai.backend.manager.api.gql.utils import check_admin_only +from ai.backend.manager.repositories.base import BatchQuerier +from ai.backend.manager.repositories.base.pagination import OffsetPagination +from ai.backend.manager.repositories.prometheus_query_preset.options import ( + PrometheusQueryPresetConditions, +) +from ai.backend.manager.services.prometheus_query_preset.actions import ( + SearchPresetsAction, +) + + +@strawberry.field(description="Get a single prometheus query preset by ID (admin only).") # type: ignore[misc] +async def admin_prometheus_query_preset( + info: Info[StrawberryGQLContext], + id: ID, +) -> PrometheusQueryPresetGQL | None: + check_admin_only() + return await fetch_admin_prometheus_query_preset(info, preset_id=uuid.UUID(id)) + + +@strawberry.field( + description="List prometheus query presets with filtering and pagination (admin only)." +) # type: ignore[misc] +async def 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 | None: + check_admin_only() + return await fetch_admin_prometheus_query_presets( + info, + filter=filter, + order_by=order_by, + before=before, + after=after, + first=first, + last=last, + limit=limit, + offset=offset, + ) + + +@strawberry.field( + description=( + "Execute a prometheus query preset by name and return the query result. " + "Available to all authenticated users." + ) +) # type: ignore[misc] +async def prometheus_query_preset_result( + info: Info[StrawberryGQLContext], + name: str, + time_range: QueryTimeRangeInput, + labels: list[MetricLabelEntryInput] | None = None, + group_labels: list[str] | None = None, + window: str | None = None, +) -> PrometheusQueryResultGQL: + processors = info.context.processors + + # Resolve name → preset_id via search + name_spec = StringMatchSpec(value=name, negated=False, case_insensitive=False) + querier = BatchQuerier( + conditions=[PrometheusQueryPresetConditions.by_name_equals(name_spec)], + orders=[], + pagination=OffsetPagination(limit=1, offset=0), + ) + search_result = await processors.prometheus_query_preset.search_presets.wait_for_complete( + SearchPresetsAction(querier=querier) + ) + if not search_result.items: + raise PrometheusQueryPresetNotFound(f"Prometheus query preset '{name}' not found") + + preset_data = search_result.items[0] + options = MetricLabelEntryInput.to_execute_options(labels, group_labels) + + return await fetch_prometheus_query_preset_result( + info, + preset_id=preset_data.id, + options=options, + window=window, + time_range=time_range.to_query_time_range(), + ) diff --git a/src/ai/backend/manager/api/gql/prometheus_query_preset/types/__init__.py b/src/ai/backend/manager/api/gql/prometheus_query_preset/types/__init__.py new file mode 100644 index 00000000000..d08f4cf18f5 --- /dev/null +++ b/src/ai/backend/manager/api/gql/prometheus_query_preset/types/__init__.py @@ -0,0 +1,53 @@ +"""Prometheus query preset GQL types.""" + +from .filters import ( + PrometheusQueryPresetFilter, + PrometheusQueryPresetOrderBy, + PrometheusQueryPresetOrderField, +) +from .inputs import ( + CreatePrometheusQueryPresetInput, + MetricLabelEntryInput, + ModifyPrometheusQueryPresetInput, + QueryTimeRangeInput, +) +from .node import ( + CreatePrometheusQueryPresetPayload, + ModifyPrometheusQueryPresetPayload, + PrometheusQueryPresetConnection, + PrometheusQueryPresetEdge, + PrometheusQueryPresetGQL, +) +from .payloads import ( + DeletePrometheusQueryPresetPayload, + MetricLabelEntryGQL, + MetricResultGQL, + MetricResultValueGQL, + PrometheusPresetOptionsGQL, + PrometheusQueryResultGQL, +) + +__all__ = [ + # Node types + "PrometheusQueryPresetGQL", + "PrometheusQueryPresetEdge", + "PrometheusQueryPresetConnection", + # Filter and OrderBy + "PrometheusQueryPresetFilter", + "PrometheusQueryPresetOrderBy", + "PrometheusQueryPresetOrderField", + # Input types + "CreatePrometheusQueryPresetInput", + "ModifyPrometheusQueryPresetInput", + "QueryTimeRangeInput", + "MetricLabelEntryInput", + # Payload and result types + "PrometheusPresetOptionsGQL", + "MetricLabelEntryGQL", + "MetricResultValueGQL", + "MetricResultGQL", + "PrometheusQueryResultGQL", + "CreatePrometheusQueryPresetPayload", + "ModifyPrometheusQueryPresetPayload", + "DeletePrometheusQueryPresetPayload", +] diff --git a/src/ai/backend/manager/api/gql/prometheus_query_preset/types/filters.py b/src/ai/backend/manager/api/gql/prometheus_query_preset/types/filters.py new file mode 100644 index 00000000000..562ad22968a --- /dev/null +++ b/src/ai/backend/manager/api/gql/prometheus_query_preset/types/filters.py @@ -0,0 +1,143 @@ +"""Prometheus query preset GQL filter and order-by types.""" + +from __future__ import annotations + +from enum import StrEnum +from typing import override + +import strawberry + +from ai.backend.manager.api.gql.base import ( + DateTimeFilter, + OrderDirection, + StringFilter, +) +from ai.backend.manager.api.gql.types import GQLFilter, GQLOrderBy +from ai.backend.manager.repositories.base import ( + QueryCondition, + QueryOrder, + combine_conditions_or, + negate_conditions, +) +from ai.backend.manager.repositories.prometheus_query_preset.options import ( + PrometheusQueryPresetConditions, + PrometheusQueryPresetOrders, +) + + +@strawberry.input( + name="PrometheusQueryPresetFilter", + description="Filter input for querying prometheus query presets.", +) +class PrometheusQueryPresetFilter(GQLFilter): + name: StringFilter | None = strawberry.field( + default=None, + description="Filter by preset name.", + ) + metric_name: str | None = strawberry.field( + default=None, + description="Filter by exact metric name.", + ) + created_at: DateTimeFilter | None = strawberry.field( + default=None, + description="Filter by creation timestamp.", + ) + + AND: list[PrometheusQueryPresetFilter] | None = strawberry.field( + default=None, + description="Combine multiple filters with AND logic.", + ) + OR: list[PrometheusQueryPresetFilter] | None = strawberry.field( + default=None, + description="Combine multiple filters with OR logic.", + ) + NOT: list[PrometheusQueryPresetFilter] | None = strawberry.field( + default=None, + description="Negate the specified filters.", + ) + + @override + def build_conditions(self) -> list[QueryCondition]: + conditions: list[QueryCondition] = [] + + if self.name: + condition = self.name.build_query_condition( + contains_factory=lambda spec: PrometheusQueryPresetConditions.by_name_contains( + spec + ), + equals_factory=lambda spec: PrometheusQueryPresetConditions.by_name_equals(spec), + starts_with_factory=lambda spec: PrometheusQueryPresetConditions.by_name_starts_with( + spec + ), + ends_with_factory=lambda spec: PrometheusQueryPresetConditions.by_name_ends_with( + spec + ), + ) + if condition: + conditions.append(condition) + + if self.metric_name is not None: + conditions.append( + PrometheusQueryPresetConditions.by_metric_name_equals(self.metric_name) + ) + + if self.created_at: + condition = self.created_at.build_query_condition( + before_factory=lambda dt: PrometheusQueryPresetConditions.by_created_at_before(dt), + after_factory=lambda dt: PrometheusQueryPresetConditions.by_created_at_after(dt), + ) + if condition: + conditions.append(condition) + + if self.AND: + for sub_filter in self.AND: + conditions.extend(sub_filter.build_conditions()) + + if self.OR: + or_sub_conditions: list[QueryCondition] = [] + for sub_filter in self.OR: + or_sub_conditions.extend(sub_filter.build_conditions()) + if or_sub_conditions: + conditions.append(combine_conditions_or(or_sub_conditions)) + + if self.NOT: + not_sub_conditions: list[QueryCondition] = [] + for sub_filter in self.NOT: + not_sub_conditions.extend(sub_filter.build_conditions()) + if not_sub_conditions: + conditions.append(negate_conditions(not_sub_conditions)) + + return conditions + + +@strawberry.enum( + name="PrometheusQueryPresetOrderField", + description="Fields available for ordering prometheus query preset results.", +) +class PrometheusQueryPresetOrderField(StrEnum): + CREATED_AT = "created_at" + UPDATED_AT = "updated_at" + NAME = "name" + + +@strawberry.input( + name="PrometheusQueryPresetOrderBy", + description="Specifies ordering for prometheus query preset results.", +) +class PrometheusQueryPresetOrderBy(GQLOrderBy): + field: PrometheusQueryPresetOrderField = strawberry.field(description="The field to order by.") + direction: OrderDirection = strawberry.field( + default=OrderDirection.DESC, + description="Sort direction.", + ) + + @override + def to_query_order(self) -> QueryOrder: + ascending = self.direction == OrderDirection.ASC + match self.field: + case PrometheusQueryPresetOrderField.CREATED_AT: + return PrometheusQueryPresetOrders.created_at(ascending) + case PrometheusQueryPresetOrderField.UPDATED_AT: + return PrometheusQueryPresetOrders.updated_at(ascending) + case PrometheusQueryPresetOrderField.NAME: + return PrometheusQueryPresetOrders.name(ascending) diff --git a/src/ai/backend/manager/api/gql/prometheus_query_preset/types/inputs.py b/src/ai/backend/manager/api/gql/prometheus_query_preset/types/inputs.py new file mode 100644 index 00000000000..14ca98b4b98 --- /dev/null +++ b/src/ai/backend/manager/api/gql/prometheus_query_preset/types/inputs.py @@ -0,0 +1,132 @@ +"""Prometheus query preset GQL input types.""" + +from __future__ import annotations + +from datetime import datetime +from uuid import UUID + +import strawberry +from strawberry import UNSET + +from ai.backend.common.dto.clients.prometheus.request import QueryTimeRange +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.base.creator import Creator +from ai.backend.manager.repositories.base.updater import Updater +from ai.backend.manager.repositories.prometheus_query_preset.creators import ( + PrometheusQueryPresetCreatorSpec, +) +from ai.backend.manager.repositories.prometheus_query_preset.updaters import ( + PrometheusQueryPresetUpdaterSpec, +) +from ai.backend.manager.types import OptionalState, TriState + + +@strawberry.input(name="PrometheusPresetOptionsInput", description="Options for preset labels.") +class PrometheusPresetOptionsInput: + filter_labels: list[str] = strawberry.field(description="Allowed filter label keys.") + group_labels: list[str] = strawberry.field(description="Allowed group-by label keys.") + + +@strawberry.input( + name="CreatePrometheusQueryPresetInput", + description="Input for creating a new prometheus query preset.", +) +class CreatePrometheusQueryPresetInput: + name: str = strawberry.field(description="Human-readable preset identifier (must be unique).") + metric_name: str = strawberry.field(description="Prometheus metric name.") + query_template: str = strawberry.field( + description="PromQL template with {labels}, {window}, {group_by} placeholders." + ) + time_window: str | None = strawberry.field( + default=None, description="Preset-specific default window." + ) + options: PrometheusPresetOptionsInput = strawberry.field( + description="Preset options including filter and group labels." + ) + + def to_creator(self) -> Creator[PrometheusQueryPresetRow]: + return Creator( + spec=PrometheusQueryPresetCreatorSpec( + name=self.name, + metric_name=self.metric_name, + query_template=self.query_template, + time_window=self.time_window, + filter_labels=self.options.filter_labels, + group_labels=self.options.group_labels, + ) + ) + + +@strawberry.input( + name="ModifyPrometheusQueryPresetInput", + description="Input for modifying an existing prometheus query preset.", +) +class ModifyPrometheusQueryPresetInput: + name: str | None = strawberry.field(default=UNSET, description="New preset name.") + metric_name: str | None = strawberry.field(default=UNSET, description="New metric name.") + query_template: str | None = strawberry.field(default=UNSET, description="New PromQL template.") + time_window: str | None = strawberry.field( + default=UNSET, description="New default time window." + ) + options: PrometheusPresetOptionsInput | None = strawberry.field( + default=UNSET, description="New preset options." + ) + + def to_updater(self, preset_id: UUID) -> Updater[PrometheusQueryPresetRow]: + spec = PrometheusQueryPresetUpdaterSpec() + + if self.name is not UNSET and self.name is not None: + spec.name = OptionalState.update(self.name) + + if self.metric_name is not UNSET and self.metric_name is not None: + spec.metric_name = OptionalState.update(self.metric_name) + + if self.query_template is not UNSET and self.query_template is not None: + spec.query_template = OptionalState.update(self.query_template) + + if self.time_window is not UNSET: + if self.time_window is None: + spec.time_window = TriState.nullify() + else: + spec.time_window = TriState.update(self.time_window) + + if self.options is not UNSET and self.options is not None: + spec.filter_labels = OptionalState.update(self.options.filter_labels) + spec.group_labels = OptionalState.update(self.options.group_labels) + + return Updater(pk_value=preset_id, spec=spec) + + +@strawberry.input(name="QueryTimeRangeInput", description="Time range for Prometheus query.") +class QueryTimeRangeInput: + start: datetime = strawberry.field(description="Start of the time range.") + end: datetime = strawberry.field(description="End of the time range.") + step: str = strawberry.field(description="Query resolution step (e.g., '60s').") + + def to_query_time_range(self) -> QueryTimeRange: + return QueryTimeRange( + start=self.start.isoformat(), + end=self.end.isoformat(), + step=self.step, + ) + + +@strawberry.input(name="MetricLabelEntryInput", description="Key-value label entry for queries.") +class MetricLabelEntryInput: + key: str = strawberry.field(description="Label key.") + value: str = strawberry.field(description="Label value.") + + @staticmethod + def to_execute_options( + labels: list[MetricLabelEntryInput] | None, + group_labels: list[str] | None, + ) -> ExecutePresetOptions: + filter_labels: dict[str, str] = {} + if labels: + for entry in labels: + filter_labels[entry.key] = entry.value + return ExecutePresetOptions( + filter_labels=filter_labels, + group_labels=group_labels or [], + ) diff --git a/src/ai/backend/manager/api/gql/prometheus_query_preset/types/node.py b/src/ai/backend/manager/api/gql/prometheus_query_preset/types/node.py new file mode 100644 index 00000000000..496e6fb62cd --- /dev/null +++ b/src/ai/backend/manager/api/gql/prometheus_query_preset/types/node.py @@ -0,0 +1,80 @@ +"""Prometheus query preset GQL Node, Edge, Connection, and mutation payload types.""" + +from __future__ import annotations + +from datetime import datetime +from typing import TYPE_CHECKING, Any, Self + +import strawberry +from strawberry import ID +from strawberry.relay import Connection, Edge, Node, NodeID + +from .payloads import PrometheusPresetOptionsGQL + +if TYPE_CHECKING: + from ai.backend.manager.data.prometheus_query_preset import PrometheusQueryPresetData + + +@strawberry.type( + name="PrometheusQueryPreset", + description="Prometheus query preset entity implementing Relay Node pattern.", +) +class PrometheusQueryPresetGQL(Node): + id: NodeID[str] = strawberry.field(description="Preset UUID (primary key).") + name: str = strawberry.field(description="Human-readable preset identifier.") + metric_name: str = strawberry.field(description="Prometheus metric name.") + query_template: str = strawberry.field(description="PromQL template with placeholders.") + time_window: str | None = strawberry.field( + description="Preset-specific default window. Falls back to server config if null." + ) + options: PrometheusPresetOptionsGQL = strawberry.field( + description="Preset options including filter and group labels." + ) + created_at: datetime = strawberry.field(description="Creation timestamp.") + updated_at: datetime = strawberry.field(description="Last update timestamp.") + + @classmethod + def from_data(cls, data: PrometheusQueryPresetData) -> Self: + return cls( + id=ID(str(data.id)), + name=data.name, + metric_name=data.metric_name, + query_template=data.query_template, + time_window=data.time_window, + options=PrometheusPresetOptionsGQL( + filter_labels=data.filter_labels, + group_labels=data.group_labels, + ), + created_at=data.created_at, + updated_at=data.updated_at, + ) + + +PrometheusQueryPresetEdge = Edge[PrometheusQueryPresetGQL] + + +@strawberry.type(description="Paginated connection for prometheus query preset records.") +class PrometheusQueryPresetConnection(Connection[PrometheusQueryPresetGQL]): + count: int = strawberry.field( + description="Total number of preset records matching the query criteria." + ) + + def __init__(self, *args: Any, count: int, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self.count = count + + +@strawberry.type( + name="CreatePrometheusQueryPresetPayload", + description="Payload returned after creating a preset.", +) +class CreatePrometheusQueryPresetPayload: + preset: PrometheusQueryPresetGQL + + +@strawberry.type( + name="ModifyPrometheusQueryPresetPayload", + description="Payload returned after modifying a preset.", +) +class ModifyPrometheusQueryPresetPayload: + preset: PrometheusQueryPresetGQL diff --git a/src/ai/backend/manager/api/gql/prometheus_query_preset/types/payloads.py b/src/ai/backend/manager/api/gql/prometheus_query_preset/types/payloads.py new file mode 100644 index 00000000000..cddf10e05ef --- /dev/null +++ b/src/ai/backend/manager/api/gql/prometheus_query_preset/types/payloads.py @@ -0,0 +1,53 @@ +"""Prometheus query preset GQL payload and result types.""" + +from __future__ import annotations + +import strawberry +from strawberry import ID + + +@strawberry.type(name="PrometheusPresetOptions", description="Preset options for label governance.") +class PrometheusPresetOptionsGQL: + filter_labels: list[str] = strawberry.field(description="Allowed filter label keys.") + group_labels: list[str] = strawberry.field(description="Allowed group-by label keys.") + + +@strawberry.type( + name="MetricLabelEntry", description="Key-value label entry from Prometheus result." +) +class MetricLabelEntryGQL: + key: str + value: str + + +@strawberry.type( + name="MetricResultValue", description="Single timestamp-value pair from Prometheus." +) +class MetricResultValueGQL: + timestamp: float + value: str + + +@strawberry.type(name="MetricResult", description="Single metric result from Prometheus query.") +class MetricResultGQL: + metric: list[MetricLabelEntryGQL] = strawberry.field( + description="Metric labels as key-value entries." + ) + values: list[MetricResultValueGQL] = strawberry.field(description="Time-series values.") + + +@strawberry.type( + name="PrometheusQueryResult", description="Result from executing a prometheus query preset." +) +class PrometheusQueryResultGQL: + status: str = strawberry.field(description="Prometheus response status.") + result_type: str = strawberry.field(description="Result type (e.g., matrix).") + result: list[MetricResultGQL] = strawberry.field(description="Metric result entries.") + + +@strawberry.type( + name="DeletePrometheusQueryPresetPayload", + description="Payload returned after deleting a preset.", +) +class DeletePrometheusQueryPresetPayload: + id: ID = strawberry.field(description="ID of the deleted preset.") diff --git a/src/ai/backend/manager/api/gql/schema.py b/src/ai/backend/manager/api/gql/schema.py index 70852ca2dc6..3bef997b711 100644 --- a/src/ai/backend/manager/api/gql/schema.py +++ b/src/ai/backend/manager/api/gql/schema.py @@ -166,6 +166,14 @@ project_domain_v2, project_v2, ) +from .prometheus_query_preset 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 .rbac import ( admin_assign_role, admin_create_permission, @@ -307,6 +315,11 @@ class Query: admin_kernels_v2 = admin_kernels_v2 admin_sessions_v2 = admin_sessions_v2 admin_image_aliases = admin_image_aliases + # Prometheus Query Preset Admin APIs + admin_prometheus_query_preset = admin_prometheus_query_preset + admin_prometheus_query_presets = admin_prometheus_query_presets + # Prometheus Query Preset - Execute (all authenticated users) + prometheus_query_preset_result = prometheus_query_preset_result # RBAC Admin APIs admin_role = admin_role admin_roles = admin_roles @@ -473,6 +486,10 @@ class Mutation: admin_delete_users_v2 = admin_delete_users_v2 admin_purge_user_v2 = admin_purge_user_v2 admin_bulk_purge_users_v2 = admin_bulk_purge_users_v2 + # Prometheus Query Preset - Admin APIs + admin_create_prometheus_query_preset = admin_create_prometheus_query_preset + admin_modify_prometheus_query_preset = admin_modify_prometheus_query_preset + admin_delete_prometheus_query_preset = admin_delete_prometheus_query_preset # RBAC Admin APIs admin_create_role = admin_create_role admin_update_role = admin_update_role diff --git a/src/ai/backend/manager/repositories/prometheus_query_preset/options.py b/src/ai/backend/manager/repositories/prometheus_query_preset/options.py index c3e91e309bc..64ca6eff970 100644 --- a/src/ai/backend/manager/repositories/prometheus_query_preset/options.py +++ b/src/ai/backend/manager/repositories/prometheus_query_preset/options.py @@ -2,6 +2,7 @@ import uuid from collections.abc import Collection +from datetime import datetime import sqlalchemy as sa @@ -79,6 +80,20 @@ def inner() -> sa.sql.expression.ColumnElement[bool]: return inner + @staticmethod + def by_created_at_before(dt: datetime) -> QueryCondition: + def inner() -> sa.sql.expression.ColumnElement[bool]: + return PrometheusQueryPresetRow.created_at < dt + + return inner + + @staticmethod + def by_created_at_after(dt: datetime) -> QueryCondition: + def inner() -> sa.sql.expression.ColumnElement[bool]: + return PrometheusQueryPresetRow.created_at > dt + + return inner + @staticmethod def by_cursor_forward(cursor_id: str) -> QueryCondition: cursor_uuid = uuid.UUID(cursor_id)