Skip to content

Commit 1b351be

Browse files
seedspiritclaude
andcommitted
refactor(BA-6004): split validate_query_template into DTO validators module
- Extract validate_query_template into common/dto/manager/v2/prometheus_query_preset/validators.py, making it self-contained (no longer constructs MetricPreset). - Move preset.py (LabelMatcher, LabelOperator, MetricPreset) into manager/clients/prometheus/, completing the common->manager move. - Remove the now-empty common/clients/prometheus/ package. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
1 parent 00b04d8 commit 1b351be

18 files changed

Lines changed: 78 additions & 59 deletions

File tree

src/ai/backend/common/clients/prometheus/__init__.py

Lines changed: 0 additions & 8 deletions
This file was deleted.

src/ai/backend/common/dto/manager/prometheus_query_preset/request.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,11 +10,13 @@
1010
from pydantic import Field, field_validator
1111

1212
from ai.backend.common.api_handlers import SENTINEL, BaseRequestModel, Sentinel
13-
from ai.backend.common.clients.prometheus.preset import validate_query_template
1413
from ai.backend.common.dto.clients.prometheus.defs import PROMETHEUS_DURATION_PATTERN
1514
from ai.backend.common.dto.clients.prometheus.request import QueryTimeRange
1615
from ai.backend.common.dto.manager.defs import DEFAULT_PAGE_LIMIT, MAX_PAGE_LIMIT
1716
from ai.backend.common.dto.manager.query import StringFilter
17+
from ai.backend.common.dto.manager.v2.prometheus_query_preset.validators import (
18+
validate_query_template,
19+
)
1820

1921
from .types import QueryDefinitionOrder
2022

src/ai/backend/common/dto/manager/v2/prometheus_query_preset/request.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,11 +11,11 @@
1111
from pydantic import Field, field_validator
1212

1313
from ai.backend.common.api_handlers import SENTINEL, BaseRequestModel, Sentinel
14-
from ai.backend.common.clients.prometheus.preset import validate_query_template
1514
from ai.backend.common.dto.clients.prometheus.defs import PROMETHEUS_DURATION_PATTERN
1615
from ai.backend.common.dto.manager.query import StringFilter, UUIDFilter
1716

1817
from .types import OrderDirection, QueryDefinitionOrderField
18+
from .validators import validate_query_template
1919

2020
__all__ = (
2121
# Options inputs
Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
"""Request DTO validators for prometheus_query_preset templates."""
2+
3+
from __future__ import annotations
4+
5+
import re
6+
7+
from ai.backend.common.exception import InvalidMetricPresetTemplate
8+
9+
__all__ = ("validate_query_template",)
10+
11+
_PLACEHOLDER_NAMES = frozenset({"labels", "window", "group_by"})
12+
_BRACE_BLOCK_RE = re.compile(r"\{([^{}]*)\}")
13+
_UNSUPPORTED_TEMPLATE_VAR_RE = re.compile(r"\$\{[^}]+\}|\$[A-Za-z_][A-Za-z0-9_]*")
14+
15+
16+
def _escape_non_placeholders(template: str) -> str:
17+
def repl(match: re.Match[str]) -> str:
18+
name = match.group(1)
19+
start, end = match.span()
20+
text = match.string
21+
already_wrapped = (
22+
start > 0 and text[start - 1] == "{" and end < len(text) and text[end] == "}"
23+
)
24+
if name not in _PLACEHOLDER_NAMES:
25+
return match.group(0) if already_wrapped else "{{" + name + "}}"
26+
if name != "labels":
27+
return match.group(0)
28+
return match.group(0) if already_wrapped else "{{" + match.group(0) + "}}"
29+
30+
return _BRACE_BLOCK_RE.sub(repl, template)
31+
32+
33+
def validate_query_template(template: str) -> str:
34+
"""Reject empty templates, foreign variables, or malformed braces."""
35+
if not template.strip():
36+
raise InvalidMetricPresetTemplate("Template must not be empty.")
37+
unsupported_vars = _UNSUPPORTED_TEMPLATE_VAR_RE.findall(template)
38+
if unsupported_vars:
39+
placeholders = ", ".join(f"{{{name}}}" for name in sorted(_PLACEHOLDER_NAMES))
40+
raise InvalidMetricPresetTemplate(
41+
f"Unsupported template variables: {unsupported_vars}. "
42+
f"Use placeholders {placeholders} or literal PromQL values."
43+
)
44+
try:
45+
_escape_non_placeholders(template).format(labels="", window="", group_by="")
46+
except (ValueError, KeyError, IndexError) as e:
47+
raise InvalidMetricPresetTemplate(
48+
f"Failed to render PromQL template ({type(e).__name__}: {e}): {template!r}"
49+
) from e
50+
return template

src/ai/backend/manager/clients/prometheus/__init__.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,16 @@
11
from .client import PrometheusClient
22
from .fixed_query_builder import FixedQueryBuilder, LabelValuesQuery
3+
from .preset import LabelMatcher, LabelOperator, MetricPreset
34
from .querier import ContainerMetricQuerier, MetricQuerier
45
from .types import MetricValue, ValueType
56

67
__all__ = [
78
"ContainerMetricQuerier",
89
"FixedQueryBuilder",
10+
"LabelMatcher",
11+
"LabelOperator",
912
"LabelValuesQuery",
13+
"MetricPreset",
1014
"MetricQuerier",
1115
"MetricValue",
1216
"PrometheusClient",

src/ai/backend/manager/clients/prometheus/client.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@
88
ClientKey,
99
ClientPool,
1010
)
11-
from ai.backend.common.clients.prometheus.preset import LabelMatcher, MetricPreset
1211
from ai.backend.common.dto.clients.prometheus.request import QueryTimeRange
1312
from ai.backend.common.dto.clients.prometheus.response import (
1413
LabelValueResponse,
@@ -28,6 +27,7 @@
2827
KernelMetricValuesByKernel,
2928
MetricResultValue,
3029
)
30+
from ai.backend.manager.clients.prometheus.preset import LabelMatcher, MetricPreset
3131

3232
DEFAULT_TIMEOUT_SECONDS: float = 30.0
3333

src/ai/backend/manager/clients/prometheus/fixed_query_builder.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@
33
from dataclasses import dataclass
44
from typing import Final
55

6-
from ai.backend.common.clients.prometheus.preset import LabelMatcher, MetricPreset
76
from ai.backend.common.metrics.types import (
87
CONTAINER_UTILIZATION_METRIC_LABEL_NAME,
98
CONTAINER_UTILIZATION_METRIC_NAME,
@@ -25,6 +24,7 @@
2524
ContainerMetricOptionalLabel,
2625
MetricType,
2726
)
27+
from ai.backend.manager.clients.prometheus.preset import LabelMatcher, MetricPreset
2828
from ai.backend.manager.clients.prometheus.querier import ContainerMetricQuerier
2929
from ai.backend.manager.clients.prometheus.types import ValueType
3030

src/ai/backend/manager/clients/prometheus/metric_types.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,14 +6,14 @@
66
from typing import Final, Self, cast
77
from uuid import UUID
88

9-
from ai.backend.common.clients.prometheus.preset import MetricPreset
109
from ai.backend.common.dto.clients.prometheus.response import (
1110
MetricResponseInfo,
1211
PrometheusResponse,
1312
)
1413
from ai.backend.common.exception import InvalidAPIParameters
1514
from ai.backend.common.metrics.types import CAPACITY_SENTINEL, CAPACITY_SENTINEL_METRICS
1615
from ai.backend.common.types import KernelId
16+
from ai.backend.manager.clients.prometheus.preset import MetricPreset
1717
from ai.backend.manager.clients.prometheus.types import MetricValue, ValueType
1818

1919
__all__ = [

src/ai/backend/common/clients/prometheus/preset.py renamed to src/ai/backend/manager/clients/prometheus/preset.py

Lines changed: 0 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -8,26 +8,6 @@
88

99
_PLACEHOLDER_NAMES = frozenset({"labels", "window", "group_by"})
1010
_BRACE_BLOCK_RE = re.compile(r"\{([^{}]*)\}")
11-
# `$ident` / `${ident}` — foreign templating syntax (Grafana, shell, etc.)
12-
# that Backend.AI does not substitute and is almost always unintended in PromQL.
13-
_UNSUPPORTED_TEMPLATE_VAR_RE = re.compile(r"\$\{[^}]+\}|\$[A-Za-z_][A-Za-z0-9_]*")
14-
15-
16-
def validate_query_template(template: str) -> str:
17-
"""Reject empty templates, foreign variables, or malformed braces.
18-
19-
Returns the dry-run rendered template (useful for inspection in tests).
20-
"""
21-
if not template.strip():
22-
raise InvalidMetricPresetTemplate("Template must not be empty.")
23-
unsupported_vars = _UNSUPPORTED_TEMPLATE_VAR_RE.findall(template)
24-
if unsupported_vars:
25-
placeholders = ", ".join(f"{{{name}}}" for name in sorted(_PLACEHOLDER_NAMES))
26-
raise InvalidMetricPresetTemplate(
27-
f"Unsupported template variables: {unsupported_vars}. "
28-
f"Use placeholders {placeholders} or literal PromQL values."
29-
)
30-
return MetricPreset(template=template).render()
3111

3212

3313
def _escape_non_placeholders(template: str) -> str:

src/ai/backend/manager/clients/prometheus/querier.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
from dataclasses import dataclass
44
from uuid import UUID
55

6-
from ai.backend.common.clients.prometheus.preset import LabelMatcher
6+
from ai.backend.manager.clients.prometheus.preset import LabelMatcher
77
from ai.backend.manager.clients.prometheus.types import ValueType
88

99

0 commit comments

Comments
 (0)