FabledSteward/plugins/host_agent/metrics_query.py

"""Host-metric read helpers (latest snapshot + bucketed history).

Kept separate from routes.py so it imports only the core PluginMetric model — not
the host_agent ORM models — which lets integration tests import these helpers
without tripping the plugin-loader's double-registration ("Table already
defined") guard.
"""
from __future__ import annotations

from datetime import datetime, timedelta, timezone

from sqlalchemy import func, or_, select

from steward.core.time_range import bucket_seconds
from steward.models.metrics import PluginMetric, PluginMetricHourly

SOURCE_MODULE = "host_agent"

# Host-level metrics charted on the detail page (sub-resources are shown as
# current-value lists, not time series, to keep the page readable).
HISTORY_METRICS = (
    "cpu_pct", "mem_used_pct", "disk_used_pct_worst", "load_1m",
    "net_rx_bps", "net_tx_bps", "disk_read_bps", "disk_write_bps",
    "temp_c_max", "psi_mem_some_avg10",
)


async def _latest_metrics_for_host(session, host_name: str) -> dict[str, dict[str, float]]:
    """{resource_name: {metric: value}} — latest sample per (resource, metric) for
    a host + its sub-resources.

    DISTINCT ON picks the newest row per group in one index-ordered pass over
    ix_plugin_metrics_module_resource_metric_recorded, instead of a GROUP-BY-max
    subquery self-joined back to the table (two passes over the whole history).
    """
    rows = (await session.execute(
        select(PluginMetric)
        .where(
            PluginMetric.source_module == SOURCE_MODULE,
            or_(
                PluginMetric.resource_name == host_name,
                PluginMetric.resource_name.like(host_name + ":%"),
            ),
        )
        .distinct(PluginMetric.resource_name, PluginMetric.metric_name)
        .order_by(
            PluginMetric.resource_name,
            PluginMetric.metric_name,
            PluginMetric.recorded_at.desc(),
        )
    )).scalars().all()
    out: dict[str, dict[str, float]] = {}
    for r in rows:
        out.setdefault(r.resource_name, {})[r.metric_name] = r.value
    return out


async def _history_for_host(session, host_name: str, since, *, raw_days: int = 7) -> dict[str, list]:
    """{metric: [[epoch_ms, avg_value], …]} host-level series since `since`.

    Retention rolls raw plugin_metrics older than `raw_days` into hourly averages
    (plugin_metrics_hourly) and deletes the raw rows. So we read the rollup for the
    part of the range older than that boundary and raw (bucket-averaged in SQL,
    capped to ≤1h to match the rollup) for the recent part — never shipping raw
    samples to Python. The two windows don't overlap, so appending hourly-then-raw
    keeps each metric's series time-ordered. Epoch-ms x feeds a linear chart axis.
    """
    now = datetime.now(timezone.utc)
    raw_cutoff = (now - timedelta(days=raw_days)).replace(minute=0, second=0, microsecond=0)
    series: dict[str, list] = {m: [] for m in HISTORY_METRICS}

    # ── Older than the raw window: the hourly rollup ──
    if since < raw_cutoff:
        hourly = (await session.execute(
            select(PluginMetricHourly.metric_name, PluginMetricHourly.bucket,
                   PluginMetricHourly.value_avg)
            .where(
                PluginMetricHourly.source_module == SOURCE_MODULE,
                PluginMetricHourly.resource_name == host_name,
                PluginMetricHourly.metric_name.in_(HISTORY_METRICS),
                PluginMetricHourly.bucket >= since,
                PluginMetricHourly.bucket < raw_cutoff,
            )
            .order_by(PluginMetricHourly.bucket)
        )).all()
        for metric_name, bucket, avg in hourly:
            series[metric_name].append([int(bucket.timestamp() * 1000), round(float(avg), 2)])

    # ── Recent part: raw, bucket-averaged in SQL ──
    raw_since = since if since >= raw_cutoff else raw_cutoff
    width_s = bucket_seconds(raw_since, 120)
    if since < raw_cutoff:
        width_s = min(width_s, 3600)  # ≤ 1h so it matches the hourly portion
    rbucket = func.date_bin(
        func.make_interval(0, 0, 0, 0, 0, 0, width_s),
        PluginMetric.recorded_at,
        func.to_timestamp(0),                            # epoch origin
    ).label("bucket")
    raw = (await session.execute(
        select(PluginMetric.metric_name, rbucket, func.avg(PluginMetric.value))
        .where(
            PluginMetric.source_module == SOURCE_MODULE,
            PluginMetric.resource_name == host_name,
            PluginMetric.metric_name.in_(HISTORY_METRICS),
            PluginMetric.recorded_at >= raw_since,
        )
        .group_by(PluginMetric.metric_name, rbucket)
        .order_by(rbucket)
    )).all()
    for metric_name, b, avg in raw:
        series[metric_name].append([int(b.timestamp() * 1000), round(float(avg), 2)])
    return series