ottyl/src/promql.rs

//! Hand-rolled parser for a small PromQL subset: metric name + label matchers, an optional `[5m]` range, and a top-level `sum()`/`avg()`/`rate()`/`histogram()`/`gauge()` wrapper.
//! Only exact-match (`=`) label matchers are supported; `rate()` requires a duration range selector, while `histogram(metric[width])` reuses that `[...]` slot for a bare bucket width instead.

use std::time::Duration;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AggOp {
    Sum,
    Avg,
    Rate,
    /// Buckets counter values into fixed-width buckets; requires `bucket_width`.
    Histogram,
    /// The most-recent counter value per label combination.
    Gauge,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct LabelMatcher {
    pub name: String,
    pub value: String,
}

#[derive(Debug, Clone, PartialEq)]
pub struct PromqlQuery {
    pub agg: Option<AggOp>,
    pub metric_name: String,
    pub matchers: Vec<LabelMatcher>,
    /// Duration range selector, e.g. `[5m]`; only meaningful for `rate()`.
    pub range: Option<Duration>,
    /// Bare-number range selector, e.g. `[10]`; only meaningful for `histogram()`, where it's the bucket width.
    pub bucket_width: Option<f64>,
}

/// What a `[...]` range selector turned out to hold: a duration, or a bare width.
enum RangeSelector {
    Duration(Duration),
    BareWidth(f64),
}

/// Parses a PromQL query string, returning a human-readable error message on malformed input.
pub fn parse_query(input: &str) -> Result<PromqlQuery, String> {
    let mut parser = Parser::new(input);
    let query = parser.parse_query()?;
    parser.skip_whitespace();
    if !parser.is_at_end() {
        return Err(format!("unexpected trailing input: '{}'", parser.remaining()));
    }
    if query.agg == Some(AggOp::Rate) && query.range.is_none() {
        return Err("rate() requires a range selector, e.g. rate(my_metric[5m])".to_string());
    }
    if query.agg == Some(AggOp::Histogram) && query.bucket_width.is_none() {
        return Err(
            "histogram() requires a bucket width, e.g. histogram(my_metric[10])".to_string(),
        );
    }
    Ok(query)
}

struct Parser<'a> {
    input: &'a str,
    pos: usize,
}

impl<'a> Parser<'a> {
    fn new(input: &'a str) -> Self {
        Parser { input, pos: 0 }
    }

    fn remaining(&self) -> &'a str {
        &self.input[self.pos..]
    }

    fn is_at_end(&self) -> bool {
        self.pos >= self.input.len()
    }

    fn skip_whitespace(&mut self) {
        while self.remaining().starts_with(|c: char| c.is_whitespace()) {
            self.pos += 1;
        }
    }

    fn peek(&self) -> Option<char> {
        self.remaining().chars().next()
    }

    fn expect_char(&mut self, expected: char) -> Result<(), String> {
        self.skip_whitespace();
        match self.peek() {
            Some(c) if c == expected => {
                self.pos += c.len_utf8();
                Ok(())
            }
            Some(c) => Err(format!("expected '{expected}', found '{c}'")),
            None => Err(format!("expected '{expected}', found end of input")),
        }
    }

    fn parse_identifier(&mut self) -> Result<String, String> {
        self.skip_whitespace();
        let rest = self.remaining();
        let mut chars = rest.char_indices();
        match chars.next() {
            Some((_, c)) if c.is_ascii_alphabetic() || c == '_' => {}
            Some((_, c)) => return Err(format!("expected an identifier, found '{c}'")),
            None => return Err("expected an identifier, found end of input".to_string()),
        }
        let end = chars
            .find(|(_, c)| !(c.is_ascii_alphanumeric() || *c == '_'))
            .map(|(idx, _)| idx)
            .unwrap_or(rest.len());
        let ident = &rest[..end];
        self.pos += end;
        Ok(ident.to_string())
    }

    fn parse_quoted_string(&mut self) -> Result<String, String> {
        self.expect_char('"')?;
        let rest = self.remaining();
        let end = rest.find('"').ok_or("unterminated string literal (missing closing '\"')")?;
        let value = &rest[..end];
        if value.contains('\\') {
            return Err("escape sequences in label values are not supported".to_string());
        }
        self.pos += end;
        self.expect_char('"')?;
        Ok(value.to_string())
    }

    fn parse_label_matcher(&mut self) -> Result<LabelMatcher, String> {
        let name = self.parse_identifier()?;
        self.expect_char('=')?;
        self.skip_whitespace();
        if self.peek() == Some('~') {
            return Err("regex label matchers ('=~') are not supported, only '='".to_string());
        }
        let value = self.parse_quoted_string()?;
        Ok(LabelMatcher { name, value })
    }

    fn parse_label_matchers(&mut self) -> Result<Vec<LabelMatcher>, String> {
        self.skip_whitespace();
        if self.peek() != Some('{') {
            return Ok(Vec::new());
        }
        self.expect_char('{')?;
        let mut matchers = Vec::new();
        self.skip_whitespace();
        if self.peek() == Some('}') {
            self.expect_char('}')?;
            return Ok(matchers);
        }
        loop {
            matchers.push(self.parse_label_matcher()?);
            self.skip_whitespace();
            match self.peek() {
                Some(',') => {
                    self.pos += 1;
                }
                Some('}') => break,
                Some(c) => return Err(format!("expected ',' or '}}', found '{c}'")),
                None => return Err("unterminated label matcher list (missing '}')".to_string()),
            }
        }
        self.expect_char('}')?;
        Ok(matchers)
    }

    /// Parses `[<digits><unit>]` (a duration, e.g. `[5m]`) or a bare `[<digits>]` (a bucket width).
    fn parse_range_selector(&mut self) -> Result<Option<RangeSelector>, String> {
        self.skip_whitespace();
        if self.peek() != Some('[') {
            return Ok(None);
        }
        self.expect_char('[')?;
        let rest = self.remaining();
        let digits_end = rest
            .find(|c: char| !c.is_ascii_digit())
            .ok_or("unterminated range selector (missing ']')")?;
        if digits_end == 0 {
            return Err("range selector must start with a number, e.g. [5m] or [10]".to_string());
        }
        let digits: i64 = rest[..digits_end]
            .parse()
            .map_err(|_| "range selector number is not a valid number".to_string())?;
        self.pos += digits_end;

        if self.peek() == Some(']') {
            self.expect_char(']')?;
            return Ok(Some(RangeSelector::BareWidth(digits as f64)));
        }

        let unit = self.peek().ok_or("range selector is missing a duration unit or ']'")?;
        let seconds_per_unit: i64 = match unit {
            's' => 1,
            'm' => 60,
            'h' => 60 * 60,
            'd' => 24 * 60 * 60,
            other => {
                return Err(format!(
                    "unknown duration unit '{other}' (expected one of s, m, h, d)"
                ));
            }
        };
        self.pos += unit.len_utf8();
        self.expect_char(']')?;
        Ok(Some(RangeSelector::Duration(Duration::from_secs(
            (digits * seconds_per_unit).max(0) as u64,
        ))))
    }

    fn parse_vector_selector(&mut self) -> Result<PromqlQuery, String> {
        let metric_name = self.parse_identifier()?;
        let matchers = self.parse_label_matchers()?;
        let (range, bucket_width) = match self.parse_range_selector()? {
            None => (None, None),
            Some(RangeSelector::Duration(d)) => (Some(d), None),
            Some(RangeSelector::BareWidth(w)) => (None, Some(w)),
        };
        Ok(PromqlQuery { agg: None, metric_name, matchers, range, bucket_width })
    }

    fn parse_query(&mut self) -> Result<PromqlQuery, String> {
        self.skip_whitespace();
        let checkpoint = self.pos;
        let ident = self.parse_identifier()?;
        let agg = match ident.as_str() {
            "sum" => Some(AggOp::Sum),
            "avg" => Some(AggOp::Avg),
            "rate" => Some(AggOp::Rate),
            "histogram" => Some(AggOp::Histogram),
            "gauge" => Some(AggOp::Gauge),
            _ => None,
        };
        self.skip_whitespace();
        if agg.is_some() && self.peek() == Some('(') {
            self.expect_char('(')?;
            let mut inner = self.parse_vector_selector()?;
            self.expect_char(')')?;
            inner.agg = agg;
            Ok(inner)
        } else {
            // Rewind and parse as a plain vector selector, covering metric names that collide with an agg_op keyword.
            self.pos = checkpoint;
            self.parse_vector_selector()
        }
    }
}

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

    #[test]
    fn parses_bare_metric_name() {
        let query = parse_query("requests").unwrap();
        assert_eq!(query.agg, None);
        assert_eq!(query.metric_name, "requests");
        assert!(query.matchers.is_empty());
        assert_eq!(query.range, None);
    }

    #[test]
    fn parses_metric_with_single_label_matcher() {
        let query = parse_query(r#"requests{path="/a"}"#).unwrap();
        assert_eq!(query.metric_name, "requests");
        assert_eq!(
            query.matchers,
            vec![LabelMatcher { name: "path".to_string(), value: "/a".to_string() }]
        );
    }

    #[test]
    fn parses_metric_with_multiple_label_matchers() {
        let query = parse_query(r#"requests{path="/a", method="GET"}"#).unwrap();
        assert_eq!(query.matchers.len(), 2);
        assert_eq!(query.matchers[0].name, "path");
        assert_eq!(query.matchers[1].name, "method");
    }

    #[test]
    fn parses_range_selector() {
        let query = parse_query("requests[5m]").unwrap();
        assert_eq!(query.range, Some(Duration::from_secs(300)));
    }

    #[test]
    fn parses_label_matchers_and_range_together() {
        let query = parse_query(r#"requests{path="/a"}[1h]"#).unwrap();
        assert_eq!(query.matchers.len(), 1);
        assert_eq!(query.range, Some(Duration::from_secs(3600)));
    }

    #[test]
    fn parses_sum_aggregation() {
        let query = parse_query(r#"sum(requests{path="/a"})"#).unwrap();
        assert_eq!(query.agg, Some(AggOp::Sum));
        assert_eq!(query.metric_name, "requests");
    }

    #[test]
    fn parses_avg_aggregation() {
        let query = parse_query("avg(requests)").unwrap();
        assert_eq!(query.agg, Some(AggOp::Avg));
    }

    #[test]
    fn parses_rate_with_range() {
        let query = parse_query("rate(requests[5m])").unwrap();
        assert_eq!(query.agg, Some(AggOp::Rate));
        assert_eq!(query.range, Some(Duration::from_secs(300)));
    }

    #[test]
    fn rejects_rate_without_range() {
        let err = parse_query("rate(requests)").unwrap_err();
        assert!(err.contains("range selector"));
    }

    #[test]
    fn rejects_empty_input() {
        assert!(parse_query("").is_err());
    }

    #[test]
    fn rejects_unclosed_brace() {
        assert!(parse_query(r#"requests{path="/a""#).is_err());
    }

    #[test]
    fn rejects_unclosed_bracket() {
        assert!(parse_query("requests[5m").is_err());
    }

    #[test]
    fn rejects_unknown_aggregation() {
        // `count` isn't a supported agg_op, so this parses as a bare metric name and trips on the trailing `(...)`.
        let err = parse_query("count(requests)").unwrap_err();
        assert!(err.contains("unexpected trailing input"));
    }

    #[test]
    fn rejects_regex_label_matcher() {
        let err = parse_query(r#"requests{path=~"/a"}"#).unwrap_err();
        assert!(err.contains("regex label matchers"));
    }

    #[test]
    fn rejects_missing_metric_name() {
        assert!(parse_query(r#"{path="/a"}"#).is_err());
    }

    #[test]
    fn rejects_unknown_duration_unit() {
        let err = parse_query("requests[5x]").unwrap_err();
        assert!(err.contains("duration unit"));
    }

    #[test]
    fn rejects_trailing_garbage() {
        let err = parse_query("requests extra").unwrap_err();
        assert!(err.contains("unexpected trailing input"));
    }

    #[test]
    fn parses_histogram_with_bare_width() {
        let query = parse_query("histogram(latency_ms[10])").unwrap();
        assert_eq!(query.agg, Some(AggOp::Histogram));
        assert_eq!(query.metric_name, "latency_ms");
        assert_eq!(query.bucket_width, Some(10.0));
        assert_eq!(query.range, None);
    }

    #[test]
    fn rejects_histogram_without_width() {
        let err = parse_query("histogram(latency_ms)").unwrap_err();
        assert!(err.contains("bucket width"));
    }

    #[test]
    fn parses_gauge_with_no_range_selector() {
        let query = parse_query(r#"gauge(connections{host="a"})"#).unwrap();
        assert_eq!(query.agg, Some(AggOp::Gauge));
        assert_eq!(query.matchers, vec![LabelMatcher { name: "host".to_string(), value: "a".to_string() }]);
    }

    #[test]
    fn rate_range_selector_is_still_a_duration_not_a_bare_width() {
        let query = parse_query("rate(requests[5m])").unwrap();
        assert_eq!(query.range, Some(Duration::from_secs(300)));
        assert_eq!(query.bucket_width, None);
    }
}