suite_logs/Logger/PrometheusMetric.py

from typing import Dict, Iterable, Optional

from prometheus_client import (
    CollectorRegistry,
    Counter,
    Gauge,
    Histogram,
    Summary,
    start_http_server,
)


class PrometheusMetric:
    """
    Helper ligero para exponer métricas de Prometheus con prefijo y etiquetas comunes.
    Inicia un endpoint HTTP para que Prometheus pueda scrapear las métricas.
    """

    def __init__(
        self,
        prefix: str = "suite_logs",
        default_labels: Optional[Dict[str, str]] = None,
        port: int = 8000,
        start_server: bool = True,
        registry: Optional[CollectorRegistry] = None,
    ):
        """
        :param prefix: prefijo que se agregará a todas las métricas (ej: suite_logs_)
        :param default_labels: etiquetas que se agregan a todas las series (ej: {"service": "api"})
        :param port: puerto HTTP donde se expondrán las métricas
        :param start_server: inicia el servidor HTTP automáticamente
        :param registry: permite inyectar un CollectorRegistry (útil para tests)
        """
        self.prefix = prefix.rstrip("_")
        self.default_labels = dict(default_labels or {})
        self.registry = registry or CollectorRegistry()
        self.port = port
        self._metric_cache: Dict[str, object] = {}

        # Se evita iniciar múltiples servidores si se crean varias instancias por error.
        if start_server:
            start_http_server(self.port, registry=self.registry)

    def counter(
        self,
        name: str,
        doc: str,
        labels: Optional[Iterable[str]] = None,
        prefix: Optional[str] = None,
    ):
        full_name = self._metric_name(name, prefix)
        metric = Counter(
            full_name,
            doc,
            labelnames=self._label_names(labels),
            registry=self.registry,
        )
        return _CounterHandle(metric, self.default_labels)

    def gauge(
        self,
        name: str,
        doc: str,
        labels: Optional[Iterable[str]] = None,
        prefix: Optional[str] = None,
    ):
        full_name = self._metric_name(name, prefix)
        metric = Gauge(
            full_name,
            doc,
            labelnames=self._label_names(labels),
            registry=self.registry,
        )
        return _GaugeHandle(metric, self.default_labels)

    def histogram(
        self,
        name: str,
        doc: str,
        labels: Optional[Iterable[str]] = None,
        buckets: Optional[Iterable[float]] = None,
        prefix: Optional[str] = None,
    ):
        full_name = self._metric_name(name, prefix)
        metric = Histogram(
            full_name,
            doc,
            labelnames=self._label_names(labels),
            registry=self.registry,
            buckets=buckets,
        )
        return _ObserveHandle(metric, self.default_labels)

    def summary(
        self,
        name: str,
        doc: str,
        labels: Optional[Iterable[str]] = None,
        prefix: Optional[str] = None,
    ):
        full_name = self._metric_name(name, prefix)
        metric = Summary(
            full_name,
            doc,
            labelnames=self._label_names(labels),
            registry=self.registry,
        )
        return _ObserveHandle(metric, self.default_labels)

    # Métodos rápidos para emitir métricas sin guardar el handle
    def counter_value(
        self,
        name: str,
        amount: float = 1.0,
        labels: Optional[Dict[str, str]] = None,
        doc: str = "",
        prefix: Optional[str] = None,
    ):
        metric = self._get_or_create(Counter, name, doc, labels, prefix)
        _CounterHandle(metric, self.default_labels).inc(amount, **(labels or {}))

    def gauge_value(
        self,
        name: str,
        value: float,
        labels: Optional[Dict[str, str]] = None,
        doc: str = "",
        prefix: Optional[str] = None,
    ):
        metric = self._get_or_create(Gauge, name, doc, labels, prefix)
        _GaugeHandle(metric, self.default_labels).set(value, **(labels or {}))

    def histogram_observe(
        self,
        name: str,
        value: float,
        labels: Optional[Dict[str, str]] = None,
        doc: str = "",
        buckets: Optional[Iterable[float]] = None,
        prefix: Optional[str] = None,
    ):
        metric = self._get_or_create(
            Histogram, name, doc, labels, prefix, buckets=buckets
        )
        _ObserveHandle(metric, self.default_labels).observe(value, **(labels or {}))

    def summary_observe(
        self,
        name: str,
        value: float,
        labels: Optional[Dict[str, str]] = None,
        doc: str = "",
        prefix: Optional[str] = None,
    ):
        metric = self._get_or_create(Summary, name, doc, labels, prefix)
        _ObserveHandle(metric, self.default_labels).observe(value, **(labels or {}))

    def _metric_name(self, name: str, prefix: Optional[str]) -> str:
        base = (prefix or self.prefix).rstrip("_")
        return f"{base}_{name}" if base else name

    def _label_names(self, labels: Optional[Iterable[str] | Dict[str, str]]) -> Iterable[str]:
        names = list(self.default_labels.keys())
        if labels:
            source = labels.keys() if isinstance(labels, dict) else labels
            for label in source:
                if label not in names:
                    names.append(label)
        return names

    def _get_or_create(
        self,
        metric_cls,
        name: str,
        doc: str,
        labels: Optional[Dict[str, str]],
        prefix: Optional[str],
        buckets: Optional[Iterable[float]] = None,
    ):
        full_name = self._metric_name(name, prefix)
        labelnames = self._label_names(labels)
        cache_key = (metric_cls.__name__, full_name, tuple(labelnames))

        if cache_key in self._metric_cache:
            return self._metric_cache[cache_key]

        kwargs: Dict[str, object] = {
            "labelnames": labelnames,
            "registry": self.registry,
        }
        if metric_cls is Histogram and buckets is not None:
            kwargs["buckets"] = buckets

        metric = metric_cls(full_name, doc or full_name, **kwargs)
        self._metric_cache[cache_key] = metric
        return metric


class _BaseHandle:
    def __init__(self, metric, default_labels: Dict[str, str]):
        self.metric = metric
        self.default_labels = default_labels

    def _child(self, labels: Optional[Dict[str, str]]):
        if not self.metric._labelnames:
            return self.metric

        merged = {**self.default_labels, **(labels or {})}
        missing = [l for l in self.metric._labelnames if l not in merged]
        if missing:
            raise ValueError(f"Faltan labels obligatorios: {missing}")
        return self.metric.labels(**merged)


class _CounterHandle(_BaseHandle):
    def inc(self, amount: float = 1.0, **labels):
        self._child(labels).inc(amount)


class _GaugeHandle(_BaseHandle):
    def set(self, value: float, **labels):
        self._child(labels).set(value)

    def inc(self, amount: float = 1.0, **labels):
        self._child(labels).inc(amount)

    def dec(self, amount: float = 1.0, **labels):
        self._child(labels).dec(amount)


class _ObserveHandle(_BaseHandle):
    def observe(self, value: float, **labels):
        self._child(labels).observe(value)