hyperhive/swarm-metrics/src/query.rs

//! The one request this binary makes: a read against the metrics store's
//! Prometheus-compatible API on the gateway.
//!
//! The gateway's read locations (`readPaths` in `swarm-victoriametrics.nix`)
//! answer an unauthenticated caller with a bare 401, never a login redirect,
//! so [`run`] treats any non-2xx as a failure with its body rather than as an
//! empty result.

use std::io::Write;

use anyhow::{Context as _, Result, bail};
use reqwest::Url;

use crate::auth::Config;

/// One read against the store, with the store's own parameter names.
#[derive(Debug, PartialEq, Eq)]
pub enum Request {
    /// `/api/v1/query`: a `MetricsQL` expression at one instant.
    Query {
        promql: String,
        time: Option<String>,
    },
    /// `/api/v1/query_range`: a `MetricsQL` expression over a window.
    QueryRange {
        promql: String,
        start: String,
        end: Option<String>,
        step: Option<String>,
    },
    /// `/api/v1/series`: the series matching the selectors.
    Series { matchers: Vec<String>, span: Span },
    /// `/api/v1/labels`: label names.
    Labels { matchers: Vec<String>, span: Span },
    /// `/api/v1/label/<name>/values`: one label's values.
    LabelValues {
        name: String,
        matchers: Vec<String>,
        span: Span,
    },
}

/// The window and result cap shared by the discovery endpoints.
#[derive(Debug, Default, PartialEq, Eq)]
pub struct Span {
    pub start: Option<String>,
    pub end: Option<String>,
    pub limit: Option<u64>,
}

/// The full request URL for `request` against the gateway at `base`.
///
/// Pure so every verb's path and parameters are testable without a store.
pub fn url(base: &str, request: &Request) -> Result<Url> {
    let mut url = Url::parse(base).with_context(|| format!("{base} is not a metrics store URL"))?;
    let mut params: Vec<(&str, &str)> = Vec::new();
    let (path, matchers, span): (Vec<&str>, &[String], Option<&Span>) = match request {
        Request::Query { promql, time } => {
            params.push(("query", promql));
            params.extend(time.as_deref().map(|t| ("time", t)));
            (vec!["query"], &[], None)
        }
        Request::QueryRange {
            promql,
            start,
            end,
            step,
        } => {
            params.push(("query", promql));
            params.push(("start", start));
            params.extend(end.as_deref().map(|t| ("end", t)));
            params.extend(step.as_deref().map(|s| ("step", s)));
            (vec!["query_range"], &[], None)
        }
        Request::Series { matchers, span } => (vec!["series"], matchers, Some(span)),
        Request::Labels { matchers, span } => (vec!["labels"], matchers, Some(span)),
        Request::LabelValues {
            name,
            matchers,
            span,
        } => (vec!["label", name, "values"], matchers, Some(span)),
    };
    params.extend(matchers.iter().map(|m| ("match[]", m.as_str())));
    let limit = span.and_then(|s| s.limit).map(|n| n.to_string());
    if let Some(span) = span {
        params.extend(span.start.as_deref().map(|t| ("start", t)));
        params.extend(span.end.as_deref().map(|t| ("end", t)));
        params.extend(limit.as_deref().map(|n| ("limit", n)));
    }
    // `extend` percent-encodes each segment, `/` included.
    url.path_segments_mut()
        .map_err(|()| anyhow::anyhow!("{base} cannot carry a path"))?
        .pop_if_empty()
        .extend(["api", "v1"])
        .extend(path);
    url.query_pairs_mut().extend_pairs(params);
    Ok(url)
}

/// Send `request` and copy the store's JSON answer to `out`.
pub fn run(cfg: &Config, request: &Request, out: &mut impl Write) -> Result<()> {
    let url = url(&cfg.base_url, request)?;
    let http = build_http_client()?;
    let send = || -> Result<reqwest::blocking::Response> {
        let bearer = cfg
            .bearer()
            .context("reading this agent's credential for the swarm metrics store")?;
        http.get(url.clone())
            .bearer_auth(bearer)
            .send()
            .with_context(|| format!("querying the swarm metrics store at {}", url.path()))
    };

    let mut response = send()?;
    // A 401 can be a secret re-minted between the store read and the check,
    // so it is read again and the request sent once more.
    if response.status() == reqwest::StatusCode::UNAUTHORIZED {
        response = send()?;
    }

    let status = response.status();
    if !status.is_success() {
        // nginx answers HTML for an `auth_request` denial, the store answers
        // JSON naming the bad expression. Truncated to the first screen.
        let body = response.text().unwrap_or_default();
        let body = body.trim();
        let shown: String = body.chars().take(2000).collect();
        bail!(
            "the swarm metrics store refused the request ({status}): {shown}{}",
            if shown.len() < body.len() {
                " […]"
            } else {
                ""
            }
        );
    }
    copy_with_newline(&mut response, out)
}

/// Copy `from` to `out` unchanged, ending with exactly one newline when the
/// body did not end in one. Streams: a broad `series` answer can be large.
fn copy_with_newline(from: &mut impl std::io::Read, out: &mut impl Write) -> Result<()> {
    let mut buf = [0u8; 8192];
    let mut last = b'\n';
    loop {
        let n = from
            .read(&mut buf)
            .context("reading the answer from the swarm metrics store")?;
        if n == 0 {
            break;
        }
        out.write_all(&buf[..n]).context("writing to stdout")?;
        last = buf[n - 1];
    }
    if last != b'\n' {
        writeln!(out).context("writing to stdout")?;
    }
    Ok(())
}

/// The HTTP client, trusting the container's own roots.
///
/// No timeout: a range query over a long window legitimately takes long, and
/// a client-side cut would read as a failure the store never had.
fn build_http_client() -> Result<reqwest::blocking::Client> {
    reqwest::blocking::Client::builder()
        .build()
        .context("building the metrics-store HTTP client")
}

#[cfg(test)]
mod tests {
    use super::*;

    const BASE: &str = "https://metrics.example";

    fn pairs(u: &Url) -> Vec<(String, String)> {
        u.query_pairs()
            .map(|(k, v)| (k.into_owned(), v.into_owned()))
            .collect()
    }

    fn owned(p: &[(&str, &str)]) -> Vec<(String, String)> {
        p.iter()
            .map(|(k, v)| ((*k).to_owned(), (*v).to_owned()))
            .collect()
    }

    #[test]
    fn an_instant_query_carries_the_expression_and_time() {
        let u = url(
            BASE,
            &Request::Query {
                promql: r#"count by (__name__) ({__name__=~"jobq\\..*"})"#.into(),
                time: Some("-5m".into()),
            },
        )
        .expect("builds");
        assert_eq!(u.path(), "/api/v1/query");
        assert_eq!(
            pairs(&u),
            owned(&[
                ("query", r#"count by (__name__) ({__name__=~"jobq\\..*"})"#),
                ("time", "-5m"),
            ])
        );
    }

    /// An unset optional parameter is left to the store, not sent empty.
    #[test]
    fn an_unset_time_is_not_sent() {
        let u = url(
            BASE,
            &Request::Query {
                promql: "up".into(),
                time: None,
            },
        )
        .expect("builds");
        assert_eq!(pairs(&u), owned(&[("query", "up")]));
    }

    #[test]
    fn a_range_query_carries_its_window_and_step() {
        let u = url(
            BASE,
            &Request::QueryRange {
                promql: "rate(x[5m])".into(),
                start: "-1h".into(),
                end: Some("now".into()),
                step: Some("1m".into()),
            },
        )
        .expect("builds");
        assert_eq!(u.path(), "/api/v1/query_range");
        assert_eq!(
            pairs(&u),
            owned(&[
                ("query", "rate(x[5m])"),
                ("start", "-1h"),
                ("end", "now"),
                ("step", "1m"),
            ])
        );
    }

    /// Each selector is its own `match[]`, the Prometheus API's spelling.
    #[test]
    fn series_sends_every_selector_and_the_span() {
        let u = url(
            BASE,
            &Request::Series {
                matchers: vec![r#"{service.name="nats"}"#.into(), "up".into()],
                span: Span {
                    start: Some("-1d".into()),
                    end: None,
                    limit: Some(20),
                },
            },
        )
        .expect("builds");
        assert_eq!(u.path(), "/api/v1/series");
        assert_eq!(
            pairs(&u),
            owned(&[
                ("match[]", r#"{service.name="nats"}"#),
                ("match[]", "up"),
                ("start", "-1d"),
                ("limit", "20"),
            ])
        );
    }

    #[test]
    fn labels_has_its_own_path() {
        let u = url(
            BASE,
            &Request::Labels {
                matchers: vec![],
                span: Span::default(),
            },
        )
        .expect("builds");
        assert_eq!(u.path(), "/api/v1/labels");
        assert_eq!(u.query(), Some(""));
    }

    /// The label name is a path segment, so one with a dot (an OTLP
    /// attribute) stays one segment and a `/` cannot reach another route.
    #[test]
    fn a_label_name_is_one_encoded_path_segment() {
        let u = url(
            BASE,
            &Request::LabelValues {
                name: "service.name".into(),
                matchers: vec![],
                span: Span::default(),
            },
        )
        .expect("builds");
        assert_eq!(u.path(), "/api/v1/label/service.name/values");

        let u = url(
            BASE,
            &Request::LabelValues {
                name: "../../admin".into(),
                matchers: vec![],
                span: Span::default(),
            },
        )
        .expect("builds");
        assert_eq!(u.path(), "/api/v1/label/..%2F..%2Fadmin/values");
    }

    /// A base with a trailing slash or a path prefix still lands on
    /// `<base>/api/v1/...`, with no doubled slash.
    #[test]
    fn the_base_url_keeps_its_prefix_without_a_doubled_slash() {
        let q = Request::Query {
            promql: "up".into(),
            time: None,
        };
        assert_eq!(
            url("https://metrics.example/", &q).expect("builds").path(),
            "/api/v1/query"
        );
        for base in ["https://gw.example/metrics", "https://gw.example/metrics/"] {
            assert_eq!(
                url(base, &q).expect("builds").path(),
                "/metrics/api/v1/query",
                "{base}"
            );
        }
    }

    #[test]
    fn a_relative_base_is_refused() {
        let q = Request::Query {
            promql: "up".into(),
            time: None,
        };
        assert!(url("metrics.example", &q).is_err());
    }

    #[test]
    fn the_body_is_copied_and_newline_terminated_once() {
        let mut out = Vec::new();
        copy_with_newline(&mut &br#"{"status":"success"}"#[..], &mut out).expect("copies");
        assert_eq!(out, b"{\"status\":\"success\"}\n");

        let mut out = Vec::new();
        copy_with_newline(&mut &b"{}\n"[..], &mut out).expect("copies");
        assert_eq!(out, b"{}\n");
    }
}