manage/backend/src/media_library_viewer_api/integrations/prometheus.py

"""Prometheus service definition."""

from __future__ import annotations

from typing import TYPE_CHECKING, Any

import requests

from media_library_viewer_api.integrations.base import (
    SecretField,
    ServiceBaseUrl,
    ServiceConfigBase,
    ServiceDefinition,
    TestResult,
    WidgetConfigBase,
    translate_connection_error,
    widget_kind,
)

if TYPE_CHECKING:
    from media_library_viewer_api.services.settings_store import SettingsStore


def test_connection(
    config: dict[str, Any],
    secrets: dict[str, str],
    store: SettingsStore,
) -> TestResult:
    """POST {grafana_url}/api/ds/query with expr 'up' via the Grafana gateway."""
    try:
        grafana_url = str(config.get("grafana_url") or "").rstrip("/")
        api_key = str(secrets.get("grafana_api_key") or "")
        datasource_uid = str(config.get("datasource_uid") or "prometheus")
        timeout = int(config.get("timeout_seconds") or 60)
        if not grafana_url:
            return TestResult(ok=False, detail="Grafana gateway URL is required.")
        if not api_key:
            return TestResult(ok=False, detail="Grafana API key is required.")
        body = {
            "queries": [
                {
                    "datasource": {"uid": datasource_uid, "type": "prometheus"},
                    "expr": "up",
                    "format": "time_series",
                    "intervalMs": 15000,
                    "maxDataPoints": 1,
                    "refId": "A",
                }
            ],
            "from": "now-1m",
            "to": "now",
        }
        resp = requests.post(
            f"{grafana_url}/api/ds/query",
            json=body,
            timeout=timeout,
            headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
        )
        resp.raise_for_status()
        return TestResult(
            ok=True,
            detail="Grafana gateway reachable.",
            evidence="Gateway reachable; datasource responded.",
        )
    except requests.HTTPError as exc:
        return translate_connection_error(exc, context="Prometheus via Grafana")
    except Exception as exc:
        return translate_connection_error(exc, context="Prometheus via Grafana")


class PrometheusConfig(ServiceConfigBase):
    """Non-secret Prometheus-via-Grafana gateway config."""

    grafana_url: ServiceBaseUrl
    datasource_uid: str = "prometheus"
    timeout_seconds: int = 60


class PrometheusMetricWidgetConfig(WidgetConfigBase):
    """A PromQL instant query rendered as a metric."""

    promql: str


class PrometheusChartWidgetConfig(WidgetConfigBase):
    """A PromQL range query rendered as a multi-series line chart (SC-101..SC-104)."""

    promql: str
    window: str = "1h"  # one of 1h / 6h / 24h / 7d (see WINDOW_PRESETS)


class PrometheusGaugeWidgetConfig(WidgetConfigBase):
    """A PromQL instant query rendered as a gauge with optional threshold bands (SC-109..SC-111)."""

    promql: str
    warn_at: float | None = None
    crit_at: float | None = None
    min: float | None = None
    max: float | None = None
    unit: str | None = None


class PrometheusMeanWidgetConfig(WidgetConfigBase):
    """A PromQL range query averaged client-side into a single value (SC-112..SC-114)."""

    promql: str
    window: str = "1h"  # one of 1h / 6h / 24h / 7d (see WINDOW_PRESETS)
    unit: str | None = None


DEFINITION = ServiceDefinition(
    service_type="prometheus",
    name="Prometheus",
    description="Metrics storage and PromQL queries.",
    config_model=PrometheusConfig,
    secret_fields=[
        SecretField(
            key="grafana_api_key",
            label="Grafana API key",
            required=True,
            helper="Service account token or API key for the Grafana gateway",
        ),
    ],
    widget_kinds=[
        widget_kind(
            kind="metric",
            name="Metric",
            description="Instant query result rendered as a metric.",
            model_cls=PrometheusMetricWidgetConfig,
            default_config={"promql": ""},
            refresh_interval_ms=30_000,
        ),
        widget_kind(
            kind="chart",
            name="Chart",
            description="Multi-series line chart from a PromQL range query.",
            model_cls=PrometheusChartWidgetConfig,
            default_config={"promql": "", "window": "1h"},
            refresh_interval_ms=60_000,
        ),
        widget_kind(
            kind="gauge",
            name="Gauge",
            description="Instant query rendered as a gauge with optional threshold bands.",
            model_cls=PrometheusGaugeWidgetConfig,
            default_config={"promql": ""},
            refresh_interval_ms=30_000,
        ),
        widget_kind(
            kind="mean",
            name="Mean",
            description="Average value of a PromQL query over a time window.",
            model_cls=PrometheusMeanWidgetConfig,
            default_config={"promql": "", "window": "1h"},
            refresh_interval_ms=60_000,
        ),
    ],
    test_callable=test_connection,
)