karldns.de/karldns/store/metrics.py

from __future__ import annotations

from datetime import UTC, datetime

from .base import StoreConnectionMixin


def _update_metric_increments(
    latency_ms: float,
    *,
    is_success: bool,
    is_dns_failure: bool,
    is_reconcile_failure: bool,
    is_bad_ip: bool,
) -> dict[str, float]:
    increments: dict[str, float] = {
        "update_requests_total": 1,
        "update_latency_ms_total": latency_ms,
        "update_latency_samples_total": 1,
    }
    optional_increments = {
        "update_success_total": is_success,
        "dns_failures_total": is_dns_failure,
        "dns_reconcile_failures_total": is_reconcile_failure,
        "bad_ip_total": is_bad_ip,
    }
    increments.update({name: 1 for name, enabled in optional_increments.items() if enabled})
    return increments


class MetricsMixin(StoreConnectionMixin):
    def increment_metric(self, name: str, amount: float = 1.0) -> None:
        now = datetime.now(UTC).isoformat()
        with self.connect() as conn:
            conn.execute(
                """
                INSERT INTO service_metrics(name, value, updated_at)
                VALUES (?, ?, ?)
                ON CONFLICT(name) DO UPDATE SET
                  value=value + excluded.value,
                  updated_at=excluded.updated_at
                """,
                (name, amount, now),
            )

    def set_metric(self, name: str, value: float) -> None:
        now = datetime.now(UTC).isoformat()
        with self.connect() as conn:
            conn.execute(
                """
                INSERT INTO service_metrics(name, value, updated_at)
                VALUES (?, ?, ?)
                ON CONFLICT(name) DO UPDATE SET
                  value=excluded.value,
                  updated_at=excluded.updated_at
                """,
                (name, value, now),
            )

    def get_service_metrics(self) -> dict[str, float]:
        with self.connect() as conn:
            rows = conn.execute("SELECT name, value FROM service_metrics").fetchall()
        return {str(row["name"]): float(row["value"]) for row in rows}

    def record_update_metrics(
        self,
        latency_ms: float,
        *,
        is_success: bool = False,
        is_dns_failure: bool = False,
        is_reconcile_failure: bool = False,
        is_bad_ip: bool = False,
    ) -> None:
        increments = _update_metric_increments(
            latency_ms,
            is_success=is_success,
            is_dns_failure=is_dns_failure,
            is_reconcile_failure=is_reconcile_failure,
            is_bad_ip=is_bad_ip,
        )

        now = datetime.now(UTC).isoformat()
        with self.connect() as conn:
            conn.executemany(
                """
                INSERT INTO service_metrics(name, value, updated_at)
                VALUES (?, ?, ?)
                ON CONFLICT(name) DO UPDATE SET
                  value=value + excluded.value,
                  updated_at=excluded.updated_at
                """,
                [(name, value, now) for name, value in increments.items()],
            )
            conn.execute(
                """
                INSERT INTO service_metrics(name, value, updated_at)
                VALUES ('last_update_latency_ms', ?, ?)
                ON CONFLICT(name) DO UPDATE SET
                  value=excluded.value,
                  updated_at=excluded.updated_at
                """,
                (latency_ms, now),
            )