"""VMP Workspace Python SDK: Query API + Remote Write with BasicAuth / AK-SK authentication."""
from __future__ import annotations
from collections import OrderedDict
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Dict, List, Optional, Sequence, Tuple
from urllib.parse import parse_qsl, urlparse
from requests.auth import AuthBase
"""VMP workspace access configuration."""
remote_write_url: str = ""
region: Optional[str] = None
access_key: Optional[str] = None
secret_key: Optional[str] = None
username: Optional[str] = None
password: Optional[str] = None
timeout_seconds: int = 30
def from_env(cls) -> "WorkspaceConfig":
"""Load configuration from environment variables."""
query_url=os.environ.get("VMP_QUERY_URL", ""),
remote_write_url=os.environ.get("VMP_REMOTE_WRITE_URL", ""),
auth_mode=os.environ.get("VMP_AUTH_MODE", "aksk").lower(),
region=os.environ.get("VOLCENGINE_REGION") or os.environ.get("VMP_REGION"),
service=os.environ.get("VMP_SERVICE", "vmp"),
access_key=os.environ.get("VOLCENGINE_ACCESS_KEY") or os.environ.get("VMP_ACCESS_KEY"),
secret_key=os.environ.get("VOLCENGINE_SECRET_KEY") or os.environ.get("VMP_SECRET_KEY"),
session_token=os.environ.get("VOLCENGINE_SESSION_TOKEN", ""),
username=os.environ.get("VMP_BASIC_AUTH_USERNAME"),
password=os.environ.get("VMP_BASIC_AUTH_PASSWORD"),
timeout_seconds=int(os.environ.get("VMP_TIMEOUT_SECONDS", "30")),
class BasicAuth(AuthBase):
"""BasicAuth adapter."""
def __init__(self, username: str, password: str) -> None:
token = base64.b64encode(f"{username}:{password}".encode()).decode()
self._header = f"Basic {token}"
def __call__(self, request: requests.PreparedRequest) -> requests.PreparedRequest:
request.headers["Authorization"] = self._header
class VolcSigV4Auth(AuthBase):
"""Volcengine AK/SK SignV4 authentication adapter."""
def __init__(self, access_key: str, secret_key: str, region: str, service: str = "vmp", session_token: str = "") -> None:
from volcengine.Credentials import Credentials
self._credentials = Credentials(access_key, secret_key, service, region, session_token)
def __call__(self, request: requests.PreparedRequest) -> requests.PreparedRequest:
from volcengine.auth.SignerV4 import SignerV4
from volcengine.base.Request import Request as VolcRequest
parsed = urlparse(request.url)
volc_request = VolcRequest()
volc_request.set_shema(parsed.scheme)
volc_request.set_method(request.method.upper())
volc_request.set_connection_timeout(10)
volc_request.set_socket_timeout(10)
volc_request.set_host(parsed.netloc)
volc_request.set_path(parsed.path or "/")
volc_request.set_query(OrderedDict(parse_qsl(parsed.query, keep_blank_values=True)))
sign_headers = OrderedDict()
sign_headers["Host"] = parsed.netloc
content_type = request.headers.get("Content-Type")
sign_headers["Content-Type"] = content_type
volc_request.set_headers(sign_headers)
body = request.body or b""
if isinstance(body, str):
body = body.encode("utf-8")
volc_request.set_body(body)
SignerV4.sign(volc_request, self._credentials)
for key, value in volc_request.headers.items():
request.headers[key] = value
def build_vmp_auth(config: WorkspaceConfig) -> AuthBase:
"""Build auth adapter from config."""
if config.auth_mode == "basic":
if not config.username or not config.password:
raise ValueError("basic auth mode requires username and password")
return BasicAuth(config.username, config.password)
if config.auth_mode == "aksk":
if not config.access_key or not config.secret_key or not config.region:
raise ValueError("aksk auth mode requires access_key, secret_key and region")
access_key=config.access_key,
secret_key=config.secret_key,
session_token=config.session_token,
raise ValueError(f"unsupported auth mode: {config.auth_mode}")
class PrometheusHttpApiClient:
"""Generic Prometheus HTTP Query API client (auth-agnostic)."""
def __init__(self, base_url: str, session: requests.Session, timeout_seconds: int = 30) -> None:
self.base_url = base_url.rstrip("/")
self.timeout = timeout_seconds
def _get(self, path: str, params: Optional[List[Tuple[str, str]]] = None) -> Dict:
response = self.session.get(self.base_url + path, params=params, timeout=self.timeout)
raise requests.HTTPError(
f"GET {path} failed: HTTP {response.status_code}, response: {response.text}",
def query(self, promql: str, ts: Optional[datetime] = None) -> Dict:
params: List[Tuple[str, str]] = [("query", promql)]
params.append(("time", _to_rfc3339(ts)))
return self._get("/api/v1/query", params)
def query_range(self, promql: str, start: datetime, end: datetime, step: str) -> Dict:
params = [("query", promql), ("start", _to_rfc3339(start)), ("end", _to_rfc3339(end)), ("step", step)]
return self._get("/api/v1/query_range", params)
def series(self, matches: Sequence[str], start: Optional[datetime] = None, end: Optional[datetime] = None) -> Dict:
raise ValueError("series requires at least one match selector")
params: List[Tuple[str, str]] = [("match[]", m) for m in matches]
params.append(("start", _to_rfc3339(start)))
params.append(("end", _to_rfc3339(end)))
return self._get("/api/v1/series", params)
def labels(self, start: Optional[datetime] = None, end: Optional[datetime] = None, matches: Optional[Sequence[str]] = None) -> Dict:
params: List[Tuple[str, str]] = []
params.extend(("match[]", m) for m in matches)
params.append(("start", _to_rfc3339(start)))
params.append(("end", _to_rfc3339(end)))
return self._get("/api/v1/labels", params)
def label_values(self, label_name: str, start: Optional[datetime] = None, end: Optional[datetime] = None, matches: Optional[Sequence[str]] = None) -> Dict:
raise ValueError("label_values requires a non-empty label_name")
params: List[Tuple[str, str]] = []
params.extend(("match[]", m) for m in matches)
params.append(("start", _to_rfc3339(start)))
params.append(("end", _to_rfc3339(end)))
return self._get(f"/api/v1/label/{label_name}/values", params)
class VmpWorkspaceClient:
"""VMP workspace client: auth + Query API + Remote Write."""
def __init__(self, config: WorkspaceConfig) -> None:
self.session = requests.Session()
self.session.auth = build_vmp_auth(config)
self.timeout = config.timeout_seconds
self.http_api = PrometheusHttpApiClient(config.query_url, self.session, config.timeout_seconds)
def query(self, promql: str, ts: Optional[datetime] = None) -> Dict:
return self.http_api.query(promql, ts)
def query_range(self, promql: str, start: datetime, end: datetime, step: str) -> Dict:
return self.http_api.query_range(promql, start, end, step)
def series(self, matches: Sequence[str], start: Optional[datetime] = None, end: Optional[datetime] = None) -> Dict:
return self.http_api.series(matches, start=start, end=end)
def labels(self, start: Optional[datetime] = None, end: Optional[datetime] = None, matches: Optional[Sequence[str]] = None) -> Dict:
return self.http_api.labels(start=start, end=end, matches=matches)
def label_values(self, label_name: str, start: Optional[datetime] = None, end: Optional[datetime] = None, matches: Optional[Sequence[str]] = None) -> Dict:
return self.http_api.label_values(label_name, start=start, end=end, matches=matches)
def write_samples(self, series: Sequence[Tuple[Dict[str, str], float, int]]) -> Dict:
"""Write samples via Remote Write. Each element is (labels_dict, value, timestamp_ms)."""
remote_write_pb2 = _get_remote_write_pb2()
payload = remote_write_pb2.WriteRequest()
for labels, value, timestamp_ms in series:
ts = payload.timeseries.add()
for key, label_value in labels.items():
label.value = label_value
sample = ts.samples.add()
sample.value = float(value)
sample.timestamp = int(timestamp_ms)
body = snappy.compress(payload.SerializeToString())
"Content-Type": "application/x-protobuf",
"Content-Encoding": "snappy",
"X-Prometheus-Remote-Write-Version": "0.1.0",
response = self.session.post(self.config.remote_write_url, data=body, headers=headers, timeout=self.timeout)
raise requests.HTTPError(
f"Remote Write failed: HTTP {response.status_code}, response: {response.text}",
return {"status_code": response.status_code, "body": response.text}
def _to_rfc3339(value: datetime) -> str:
if value.tzinfo is None:
value = value.replace(tzinfo=timezone.utc)
return value.astimezone(timezone.utc).isoformat().replace("+00:00", "Z")
_REMOTE_WRITE_PB2 = None
def _get_remote_write_pb2():
global _REMOTE_WRITE_PB2
if _REMOTE_WRITE_PB2 is not None:
return _REMOTE_WRITE_PB2
from google.protobuf import descriptor_pool as _descriptor_pool
from google.protobuf.internal import builder as _builder
DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(
b'\n\x1dprometheus_remote_write.proto\x12\nprometheus\":\n\x0cWriteRequest\x12*\n\ntimeseries\x18\x01 \x03(\x0b\x32\x16.prometheus.TimeSeries\"T\n\nTimeSeries\x12!\n\x06labels\x18\x01 \x03(\x0b\x32\x11.prometheus.Label\x12#\n\x07samples\x18\x02 \x03(\x0b\x32\x12.prometheus.Sample\"$\n\x05Label\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t\"*\n\x06Sample\x12\r\n\x05value\x18\x01 \x01(\x01\x12\x11\n\ttimestamp\x18\x02 \x01(\x03\x62\x06proto3'
_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'prometheus_remote_write_pb2', _globals)
module = types.SimpleNamespace(**_globals)
_REMOTE_WRITE_PB2 = module