prometheus-parser 0.4.3-alpha9

// (C) Copyright 2019-2020 Hewlett Packard Enterprise Development LP

use std::collections::HashSet;
use std::fmt;
use std::iter::FromIterator;

use super::expression::Expression;
use super::misc::{PromDuration, Span, Subquery};
use super::return_value::{LabelSetOp, ReturnKind, ReturnValue};

/// Label equality operators
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum LabelOp {
    Equal,
    NotEqual,
    RegexEqual,
    RegexNotEqual,
}

impl fmt::Display for LabelOp {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "{}",
            match self {
                LabelOp::Equal => "=",
                LabelOp::NotEqual => "!=",
                LabelOp::RegexEqual => "=~",
                LabelOp::RegexNotEqual => "!~",
            }
        )
    }
}

/// A label matcher in a selector
#[derive(Debug, PartialEq, Eq, Clone)]
pub struct Label {
    pub key: String,
    pub op: LabelOp,
    pub value: String,

    pub span: Option<Span>,
}

impl Label {
    pub fn new<S: Into<String>>(op: LabelOp, key: S, value: S) -> Self {
        Label {
            op,
            key: key.into(),
            value: value.into(),
            span: None,
        }
    }

    pub fn equal<S: Into<String>>(key: S, value: S) -> Self {
        Label::new(LabelOp::Equal, key, value)
    }

    pub fn not_equal<S: Into<String>>(key: S, value: S) -> Self {
        Label::new(LabelOp::NotEqual, key, value)
    }

    pub fn regex_equal<S: Into<String>>(key: S, value: S) -> Self {
        Label::new(LabelOp::RegexEqual, key, value)
    }

    pub fn regex_notequal<S: Into<String>>(key: S, value: S) -> Self {
        Label::new(LabelOp::RegexNotEqual, key, value)
    }

    pub fn key<S: Into<String>>(mut self, key: S) -> Self {
        self.key = key.into();
        self
    }

    pub fn op(mut self, op: LabelOp) -> Self {
        self.op = op;
        self
    }

    pub fn value<S: Into<String>>(mut self, value: S) -> Self {
        self.value = value.into();
        self
    }

    pub fn span<S: Into<Span>>(mut self, span: S) -> Self {
        self.span = Some(span.into());
        self
    }
}

impl fmt::Display for Label {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}{}{:?}", self.key, self.op, self.value)
    }
}

/// A Selector that retrieves time series data from Prometheus
#[derive(Debug, PartialEq, Eq, Clone, Default)]
pub struct Selector {
    pub metric: Option<String>,
    pub labels: Vec<Label>,
    pub range: Option<PromDuration>,
    pub offset: Option<PromDuration>,
    pub subquery: Option<Subquery>,

    pub span: Option<Span>,
}

impl Selector {
    pub fn new() -> Self {
        Selector {
            metric: None,
            labels: vec![],
            range: None,
            offset: None,
            subquery: None,
            span: None,
        }
    }

    /// Sets or replaces this Selector's metric
    pub fn metric<S: Into<String>>(mut self, metric: S) -> Self {
        self.metric = Some(metric.into());
        self
    }

    /// Clears this Selector's metric
    pub fn clear_metric(mut self) -> Self {
        self.metric = None;
        self
    }

    /// Adds a label to this Selector
    pub fn label(mut self, label: Label) -> Self {
        self.labels.push(label);
        self
    }

    /// Replaces this Selector's labels with the given set
    pub fn labels(mut self, labels: Vec<Label>) -> Self {
        self.labels = labels;
        self
    }

    /// Clears this Selector's set of labels
    pub fn clear_labels(mut self) -> Self {
        self.labels.clear();
        self
    }

    /// Sets or replaces this Selector's range
    pub fn range(mut self, range: PromDuration) -> Self {
        self.range = Some(range);
        self
    }

    /// Clears this Selector's range
    pub fn clear_range(mut self) -> Self {
        self.range = None;
        self
    }

    /// Sets or replaces this Selector's offset
    pub fn offset(mut self, offset: PromDuration) -> Self {
        self.offset = Some(offset);
        self
    }

    pub fn clear_offset(mut self) -> Self {
        self.offset = None;
        self
    }

    pub fn subquery(mut self, subquery: Subquery) -> Self {
        self.subquery = Some(subquery);
        self
    }

    pub fn clear_subquery(mut self) -> Self {
        self.subquery = None;
        self
    }

    pub fn span<S: Into<Span>>(mut self, span: S) -> Self {
        self.span = Some(span.into());
        self
    }

    pub fn wrap(self) -> Expression {
        Expression::Selector(self)
    }

    pub fn return_value(&self) -> ReturnValue {
        let kind = match (self.range.is_some(), self.subquery.is_some()) {
            (false, false) => ReturnKind::InstantVector,
            (false, true) => ReturnKind::RangeVector,
            (true, false) => ReturnKind::RangeVector,

            // range + subquery is not allowed (however this is syntactically invalid)
            (true, true) => ReturnKind::unknown(
                "range and subquery are not allowed together",
                self.clone().wrap(),
            ),
        };

        // if a label is selected, we can infer that it must be present on the
        // output vector
        let mut label_ops = Vec::new();
        if !self.labels.is_empty() {
            label_ops.push(LabelSetOp::append(
                self.clone().wrap(),
                self.span,
                HashSet::from_iter(self.labels.iter().cloned().map(|l| l.key)),
            ));
        }

        ReturnValue { kind, label_ops }
    }
}

impl fmt::Display for Selector {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if let Some(metric) = &self.metric {
            write!(f, "{}", metric)?;
        }

        if !self.labels.is_empty() {
            write!(f, "{{")?;
            for (i, label) in self.labels.iter().enumerate() {
                if i > 0 {
                    write!(f, ",")?;
                }

                write!(f, "{}", label)?;
            }

            write!(f, "}}")?;
        };

        if let Some(range) = &self.range {
            write!(f, "[{}]", range)?;
        }

        if let Some(offset) = &self.offset {
            write!(f, " offset {}", offset)?;
        }

        if let Some(subquery) = &self.subquery {
            write!(f, "{}", subquery)?;
        }

        Ok(())
    }
}