Skip to content

Commit 6a324d6

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 b0ba9b9 commit 6a324d6

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
@@ -4,14 +4,18 @@
44
ContainerMetricQueryBuilder,
55
LabelValuesQuery,
66
)
7+
from .preset import LabelMatcher, LabelOperator, MetricPreset
78
from .querier import ContainerMetricQuerier, MetricQuerier
89
from .types import MetricValue, ValueType
910

1011
__all__ = [
1112
"ContainerLiveStatQueryBuilder",
1213
"ContainerMetricQueryBuilder",
1314
"ContainerMetricQuerier",
15+
"LabelMatcher",
16+
"LabelOperator",
1417
"LabelValuesQuery",
18+
"MetricPreset",
1519
"MetricQuerier",
1620
"MetricValue",
1721
"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,
@@ -30,6 +29,7 @@
3029
KernelLiveStatBatchResult,
3130
MetricResultValue,
3231
)
32+
from ai.backend.manager.clients.prometheus.preset import LabelMatcher, MetricPreset
3333

3434
DEFAULT_TIMEOUT_SECONDS: float = 30.0
3535

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,
@@ -16,6 +15,7 @@
1615
ContainerMetricOptionalLabel,
1716
MetricType,
1817
)
18+
from ai.backend.manager.clients.prometheus.preset import LabelMatcher, MetricPreset
1919
from ai.backend.manager.clients.prometheus.querier import ContainerMetricQuerier
2020
from ai.backend.manager.clients.prometheus.types import ValueType
2121

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,13 +7,13 @@
77
from typing import Final, Self
88
from uuid import UUID
99

10-
from ai.backend.common.clients.prometheus.preset import MetricPreset
1110
from ai.backend.common.dto.clients.prometheus.response import (
1211
MetricResponseInfo,
1312
PrometheusResponse,
1413
)
1514
from ai.backend.common.exception import InvalidAPIParameters
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)