bk_promql_parser/parser/
ast.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use crate::label::{Labels, Matchers, METRIC_NAME};
16use crate::parser::token::{
17    self, token_display, T_BOTTOMK, T_COUNT_VALUES, T_END, T_QUANTILE, T_START, T_TOPK,
18};
19use crate::parser::token::{Token, TokenId, TokenType};
20use crate::parser::value::ValueType;
21use crate::parser::{indent, Function, FunctionArgs, Prettier, MAX_CHARACTERS_PER_LINE};
22use crate::util::display_duration;
23use std::fmt::{self, Write};
24use std::ops::Neg;
25use std::sync::Arc;
26use std::time::{Duration, SystemTime};
27
28/// LabelModifier acts as
29///
30/// # Aggregation Modifier
31///
32/// - Exclude means `ignoring`
33/// - Include means `on`
34///
35/// # Vector Match Modifier
36///
37/// - Exclude means `without` removes the listed labels from the result vector,
38/// while all other labels are preserved in the output.
39/// - Include means `by` does the opposite and drops labels that are not listed in the by clause,
40/// even if their label values are identical between all elements of the vector.
41///
42/// if empty listed labels, meaning no grouping
43#[derive(Debug, Clone, PartialEq, Eq)]
44pub enum LabelModifier {
45    Include(Labels),
46    Exclude(Labels),
47}
48
49impl LabelModifier {
50    pub fn labels(&self) -> &Labels {
51        match self {
52            LabelModifier::Include(l) => l,
53            LabelModifier::Exclude(l) => l,
54        }
55    }
56
57    pub fn is_include(&self) -> bool {
58        matches!(*self, LabelModifier::Include(_))
59    }
60
61    pub fn include(ls: Vec<&str>) -> Self {
62        Self::Include(Labels::new(ls))
63    }
64
65    pub fn exclude(ls: Vec<&str>) -> Self {
66        Self::Exclude(Labels::new(ls))
67    }
68}
69
70/// The label list provided with the group_left or group_right modifier contains
71/// additional labels from the "one"-side to be included in the result metrics.
72#[derive(Debug, Clone, PartialEq, Eq)]
73pub enum VectorMatchCardinality {
74    OneToOne,
75    ManyToOne(Labels),
76    OneToMany(Labels),
77    ManyToMany, // logical/set binary operators
78}
79
80impl VectorMatchCardinality {
81    pub fn labels(&self) -> Option<&Labels> {
82        match self {
83            VectorMatchCardinality::ManyToOne(l) => Some(l),
84            VectorMatchCardinality::OneToMany(l) => Some(l),
85            VectorMatchCardinality::ManyToMany => None,
86            VectorMatchCardinality::OneToOne => None,
87        }
88    }
89
90    pub fn many_to_one(ls: Vec<&str>) -> Self {
91        Self::ManyToOne(Labels::new(ls))
92    }
93
94    pub fn one_to_many(ls: Vec<&str>) -> Self {
95        Self::OneToMany(Labels::new(ls))
96    }
97}
98
99/// Binary Expr Modifier
100#[derive(Debug, Clone, PartialEq, Eq)]
101pub struct BinModifier {
102    /// The matching behavior for the operation if both operands are Vectors.
103    /// If they are not this field is None.
104    pub card: VectorMatchCardinality,
105
106    /// on/ignoring on labels.
107    /// like a + b, no match modifier is needed.
108    pub matching: Option<LabelModifier>,
109    /// If a comparison operator, return 0/1 rather than filtering.
110    pub return_bool: bool,
111}
112
113impl fmt::Display for BinModifier {
114    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
115        let mut s = String::from(self.bool_str());
116
117        if let Some(matching) = &self.matching {
118            match matching {
119                LabelModifier::Include(ls) => write!(s, "on ({ls}) ")?,
120                LabelModifier::Exclude(ls) if !ls.is_empty() => write!(s, "ignoring ({ls}) ")?,
121                _ => (),
122            }
123        }
124
125        match &self.card {
126            VectorMatchCardinality::ManyToOne(ls) => write!(s, "group_left ({ls}) ")?,
127            VectorMatchCardinality::OneToMany(ls) => write!(s, "group_right ({ls}) ")?,
128            _ => (),
129        }
130
131        if s.trim().is_empty() {
132            write!(f, "")
133        } else {
134            write!(f, " {}", s.trim_end()) // there is a leading space here
135        }
136    }
137}
138
139impl Default for BinModifier {
140    fn default() -> Self {
141        Self {
142            card: VectorMatchCardinality::OneToOne,
143            matching: None,
144            return_bool: false,
145        }
146    }
147}
148
149impl BinModifier {
150    pub fn with_card(mut self, card: VectorMatchCardinality) -> Self {
151        self.card = card;
152        self
153    }
154
155    pub fn with_matching(mut self, matching: Option<LabelModifier>) -> Self {
156        self.matching = matching;
157        self
158    }
159
160    pub fn with_return_bool(mut self, return_bool: bool) -> Self {
161        self.return_bool = return_bool;
162        self
163    }
164
165    pub fn is_labels_joint(&self) -> bool {
166        matches!(
167            (self.card.labels(), &self.matching),
168            (Some(labels), Some(matching)) if labels.is_joint(matching.labels())
169        )
170    }
171
172    pub fn intersect_labels(&self) -> Option<Vec<String>> {
173        if let Some(labels) = self.card.labels() {
174            if let Some(matching) = &self.matching {
175                return Some(labels.intersect(matching.labels()).labels);
176            }
177        };
178        None
179    }
180
181    pub fn is_matching_on(&self) -> bool {
182        matches!(&self.matching, Some(matching) if matching.is_include())
183    }
184
185    pub fn is_matching_labels_not_empty(&self) -> bool {
186        matches!(&self.matching, Some(matching) if !matching.labels().is_empty())
187    }
188
189    pub fn bool_str(&self) -> &str {
190        if self.return_bool {
191            "bool "
192        } else {
193            ""
194        }
195    }
196}
197
198#[derive(Debug, Clone, PartialEq, Eq)]
199pub enum Offset {
200    Pos(Duration),
201    Neg(Duration),
202}
203
204impl fmt::Display for Offset {
205    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
206        match self {
207            Offset::Pos(dur) => write!(f, "{}", display_duration(dur)),
208            Offset::Neg(dur) => write!(f, "-{}", display_duration(dur)),
209        }
210    }
211}
212#[derive(Debug, Clone, PartialEq, Eq)]
213pub enum AtModifier {
214    Start,
215    End,
216    /// at can be earlier than UNIX_EPOCH
217    At(SystemTime),
218}
219
220impl fmt::Display for AtModifier {
221    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
222        match self {
223            AtModifier::Start => write!(f, "@ {}()", token_display(T_START)),
224            AtModifier::End => write!(f, "@ {}()", token_display(T_END)),
225            AtModifier::At(time) => {
226                let d = time
227                    .duration_since(SystemTime::UNIX_EPOCH)
228                    .unwrap_or(Duration::ZERO); // This should not happen
229                write!(f, "@ {:.3}", d.as_secs() as f64)
230            }
231        }
232    }
233}
234impl TryFrom<TokenId> for AtModifier {
235    type Error = String;
236
237    fn try_from(id: TokenId) -> Result<Self, Self::Error> {
238        match id {
239            T_START => Ok(AtModifier::Start),
240            T_END => Ok(AtModifier::End),
241            _ => Err(format!(
242                "invalid @ modifier preprocessor '{}', START or END is valid.",
243                token::token_display(id)
244            )),
245        }
246    }
247}
248
249impl TryFrom<Token> for AtModifier {
250    type Error = String;
251
252    fn try_from(token: Token) -> Result<Self, Self::Error> {
253        AtModifier::try_from(token.id())
254    }
255}
256
257impl TryFrom<NumberLiteral> for AtModifier {
258    type Error = String;
259
260    fn try_from(num: NumberLiteral) -> Result<Self, Self::Error> {
261        AtModifier::try_from(num.val)
262    }
263}
264
265impl TryFrom<Expr> for AtModifier {
266    type Error = String;
267
268    fn try_from(ex: Expr) -> Result<Self, Self::Error> {
269        match ex {
270            Expr::NumberLiteral(nl) => AtModifier::try_from(nl),
271            _ => Err("invalid float value after @ modifier".into()),
272        }
273    }
274}
275
276impl TryFrom<f64> for AtModifier {
277    type Error = String;
278
279    fn try_from(secs: f64) -> Result<Self, Self::Error> {
280        let err_info = format!("timestamp out of bounds for @ modifier: {secs}");
281
282        if secs.is_nan() || secs.is_infinite() || secs >= f64::MAX || secs <= f64::MIN {
283            return Err(err_info);
284        }
285        let milli = (secs * 1000f64).round().abs() as u64;
286
287        let duration = Duration::from_millis(milli);
288        let mut st = Some(SystemTime::UNIX_EPOCH);
289        if secs.is_sign_positive() {
290            st = SystemTime::UNIX_EPOCH.checked_add(duration);
291        }
292        if secs.is_sign_negative() {
293            st = SystemTime::UNIX_EPOCH.checked_sub(duration);
294        }
295
296        st.map(Self::At).ok_or(err_info)
297    }
298}
299
300/// EvalStmt holds an expression and information on the range it should
301/// be evaluated on.
302#[allow(rustdoc::broken_intra_doc_links)]
303#[derive(Debug, Clone)]
304pub struct EvalStmt {
305    /// Expression to be evaluated.
306    pub expr: Expr,
307
308    /// The time boundaries for the evaluation. If start equals end an instant
309    /// is evaluated.
310    pub start: SystemTime,
311    pub end: SystemTime,
312    /// Time between two evaluated instants for the range [start:end].
313    pub interval: Duration,
314    /// Lookback delta to use for this evaluation.
315    pub lookback_delta: Duration,
316}
317
318/// Grammar:
319/// ``` norust
320/// <aggr-op> [without|by (<label list>)] ([parameter,] <vector expression>)
321/// <aggr-op>([parameter,] <vector expression>) [without|by (<label list>)]
322/// ```
323///
324/// parameter is only required for `count_values`, `quantile`, `topk` and `bottomk`.
325#[derive(Debug, Clone, PartialEq, Eq)]
326pub struct AggregateExpr {
327    /// The used aggregation operation.
328    pub op: TokenType,
329    /// The Vector expression over which is aggregated.
330    pub expr: Box<Expr>,
331    /// Parameter used by some aggregators.
332    pub param: Option<Box<Expr>>,
333    /// modifier is optional for some aggregation operators, like sum.
334    pub modifier: Option<LabelModifier>,
335}
336
337impl AggregateExpr {
338    fn get_op_string(&self) -> String {
339        let mut s = self.op.to_string();
340
341        if let Some(modifier) = &self.modifier {
342            match modifier {
343                LabelModifier::Exclude(ls) => write!(s, " without ({ls}) ").unwrap(),
344                LabelModifier::Include(ls) if !ls.is_empty() => write!(s, " by ({ls}) ").unwrap(),
345                _ => (),
346            }
347        }
348        s
349    }
350}
351
352impl fmt::Display for AggregateExpr {
353    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
354        write!(f, "{}", self.get_op_string())?;
355
356        write!(f, "(")?;
357        if let Some(param) = &self.param {
358            write!(f, "{param}, ")?;
359        }
360        write!(f, "{})", self.expr)?;
361
362        Ok(())
363    }
364}
365
366impl Prettier for AggregateExpr {
367    fn format(&self, level: usize, max: usize) -> String {
368        let mut s = format!("{}{}(\n", indent(level), self.get_op_string());
369        if let Some(param) = &self.param {
370            writeln!(s, "{},", param.pretty(level + 1, max)).unwrap();
371        }
372        writeln!(s, "{}", self.expr.pretty(level + 1, max)).unwrap();
373        write!(s, "{})", indent(level)).unwrap();
374        s
375    }
376}
377
378/// UnaryExpr will negate the expr
379#[derive(Debug, Clone, PartialEq, Eq)]
380pub struct UnaryExpr {
381    pub expr: Box<Expr>,
382}
383
384impl fmt::Display for UnaryExpr {
385    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
386        write!(f, "-{}", self.expr)
387    }
388}
389
390impl Prettier for UnaryExpr {
391    fn pretty(&self, level: usize, max: usize) -> String {
392        format!(
393            "{}-{}",
394            indent(level),
395            self.expr.pretty(level, max).trim_start()
396        )
397    }
398}
399
400/// Grammar:
401/// ``` norust
402/// <vector expr> <bin-op> ignoring(<label list>) group_left(<label list>) <vector expr>
403/// <vector expr> <bin-op> ignoring(<label list>) group_right(<label list>) <vector expr>
404/// <vector expr> <bin-op> on(<label list>) group_left(<label list>) <vector expr>
405/// <vector expr> <bin-op> on(<label list>) group_right(<label list>) <vector expr>
406/// ```
407#[derive(Debug, Clone, PartialEq, Eq)]
408pub struct BinaryExpr {
409    /// The operation of the expression.
410    pub op: TokenType,
411    /// The operands on the left sides of the operator.
412    pub lhs: Box<Expr>,
413    /// The operands on the right sides of the operator.
414    pub rhs: Box<Expr>,
415
416    pub modifier: Option<BinModifier>,
417}
418
419impl BinaryExpr {
420    pub fn is_matching_on(&self) -> bool {
421        matches!(&self.modifier, Some(modifier) if modifier.is_matching_on())
422    }
423
424    pub fn is_matching_labels_not_empty(&self) -> bool {
425        matches!(&self.modifier, Some(modifier) if modifier.is_matching_labels_not_empty())
426    }
427
428    pub fn return_bool(&self) -> bool {
429        matches!(&self.modifier, Some(modifier) if modifier.return_bool)
430    }
431
432    /// check if labels of card and matching are joint
433    pub fn is_labels_joint(&self) -> bool {
434        matches!(&self.modifier, Some(modifier) if modifier.is_labels_joint())
435    }
436
437    /// intersect labels of card and matching
438    pub fn intersect_labels(&self) -> Option<Vec<String>> {
439        self.modifier
440            .as_ref()
441            .and_then(|modifier| modifier.intersect_labels())
442    }
443
444    fn get_op_matching_string(&self) -> String {
445        match &self.modifier {
446            Some(modifier) => format!("{}{modifier}", self.op),
447            None => self.op.to_string(),
448        }
449    }
450}
451
452impl fmt::Display for BinaryExpr {
453    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
454        write!(
455            f,
456            "{} {} {}",
457            self.lhs,
458            self.get_op_matching_string(),
459            self.rhs
460        )
461    }
462}
463
464impl Prettier for BinaryExpr {
465    fn format(&self, level: usize, max: usize) -> String {
466        format!(
467            "{}\n{}{}\n{}",
468            self.lhs.pretty(level + 1, max),
469            indent(level),
470            self.get_op_matching_string(),
471            self.rhs.pretty(level + 1, max)
472        )
473    }
474}
475
476#[derive(Debug, Clone, PartialEq, Eq)]
477pub struct ParenExpr {
478    pub expr: Box<Expr>,
479}
480
481impl fmt::Display for ParenExpr {
482    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
483        write!(f, "({})", self.expr)
484    }
485}
486
487impl Prettier for ParenExpr {
488    fn format(&self, level: usize, max: usize) -> String {
489        format!(
490            "{}(\n{}\n{})",
491            indent(level),
492            self.expr.pretty(level + 1, max),
493            indent(level)
494        )
495    }
496}
497
498/// Grammar:
499/// ```norust
500/// <instant_query> '[' <range> ':' [<resolution>] ']' [ @ <float_literal> ] [ offset <duration> ]
501/// ```
502#[derive(Debug, Clone, PartialEq, Eq)]
503pub struct SubqueryExpr {
504    pub expr: Box<Expr>,
505    pub offset: Option<Offset>,
506    pub at: Option<AtModifier>,
507    pub range: Duration,
508    /// Default is the global evaluation interval.
509    pub step: Option<Duration>,
510}
511
512impl SubqueryExpr {
513    fn get_time_suffix_string(&self) -> String {
514        let step = match &self.step {
515            Some(step) => display_duration(step),
516            None => String::from(""),
517        };
518        let range = display_duration(&self.range);
519
520        let mut s = format!("[{range}:{step}]");
521
522        if let Some(at) = &self.at {
523            write!(s, " {at}").unwrap();
524        }
525
526        if let Some(offset) = &self.offset {
527            write!(s, " offset {offset}").unwrap();
528        }
529        s
530    }
531}
532
533impl fmt::Display for SubqueryExpr {
534    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
535        write!(f, "{}{}", self.expr, self.get_time_suffix_string())
536    }
537}
538
539impl Prettier for SubqueryExpr {
540    fn pretty(&self, level: usize, max: usize) -> String {
541        format!(
542            "{}{}",
543            self.expr.pretty(level, max),
544            self.get_time_suffix_string()
545        )
546    }
547}
548
549#[derive(Debug, Clone)]
550pub struct NumberLiteral {
551    pub val: f64,
552}
553
554impl NumberLiteral {
555    pub fn new(val: f64) -> Self {
556        Self { val }
557    }
558}
559
560impl PartialEq for NumberLiteral {
561    fn eq(&self, other: &Self) -> bool {
562        self.val == other.val || self.val.is_nan() && other.val.is_nan()
563    }
564}
565
566impl Eq for NumberLiteral {}
567
568impl Neg for NumberLiteral {
569    type Output = Self;
570
571    fn neg(self) -> Self::Output {
572        NumberLiteral { val: -self.val }
573    }
574}
575
576impl fmt::Display for NumberLiteral {
577    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
578        if self.val == f64::INFINITY {
579            write!(f, "Inf")
580        } else if self.val == f64::NEG_INFINITY {
581            write!(f, "-Inf")
582        } else if f64::is_nan(self.val) {
583            write!(f, "NaN")
584        } else {
585            write!(f, "{}", self.val)
586        }
587    }
588}
589
590impl Prettier for NumberLiteral {
591    fn needs_split(&self, _max: usize) -> bool {
592        false
593    }
594}
595
596#[derive(Debug, Clone, PartialEq, Eq)]
597pub struct StringLiteral {
598    pub val: String,
599}
600
601impl fmt::Display for StringLiteral {
602    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
603        write!(f, "\"{}\"", self.val)
604    }
605}
606
607impl Prettier for StringLiteral {
608    fn needs_split(&self, _max: usize) -> bool {
609        false
610    }
611}
612
613#[derive(Debug, Clone, PartialEq, Eq)]
614pub struct VectorSelector {
615    pub name: Option<String>,
616    pub matchers: Matchers,
617    pub offset: Option<Offset>,
618    pub at: Option<AtModifier>,
619}
620
621impl VectorSelector {
622    pub fn new(name: Option<String>, matchers: Matchers) -> Self {
623        VectorSelector {
624            name,
625            matchers,
626            offset: None,
627            at: None,
628        }
629    }
630}
631
632impl Default for VectorSelector {
633    fn default() -> Self {
634        Self {
635            name: None,
636            matchers: Matchers::empty(),
637            offset: None,
638            at: None,
639        }
640    }
641}
642
643impl From<String> for VectorSelector {
644    fn from(name: String) -> Self {
645        VectorSelector {
646            name: Some(name),
647            offset: None,
648            at: None,
649            matchers: Matchers::empty(),
650        }
651    }
652}
653
654/// directly create an instant vector with only METRIC_NAME matcher.
655///
656/// # Examples
657///
658/// Basic usage:
659///
660/// ``` rust
661/// use bk_promql_parser::label::Matchers;
662/// use bk_promql_parser::parser::VectorSelector;
663///
664/// let vs = VectorSelector {
665///     name: Some(String::from("foo")),
666///     offset: None,
667///     at: None,
668///     matchers: Matchers::empty(),
669/// };
670///
671/// assert_eq!(VectorSelector::from("foo"), vs);
672/// ```
673impl From<&str> for VectorSelector {
674    fn from(name: &str) -> Self {
675        VectorSelector::from(name.to_string())
676    }
677}
678
679impl Neg for VectorSelector {
680    type Output = UnaryExpr;
681
682    fn neg(self) -> Self::Output {
683        let ex = Expr::VectorSelector(self);
684        UnaryExpr { expr: Box::new(ex) }
685    }
686}
687
688impl fmt::Display for VectorSelector {
689    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
690        if let Some(name) = &self.name {
691            write!(f, "{name}")?;
692        }
693        let matchers = &self.matchers.to_string();
694        if !matchers.is_empty() {
695            write!(f, "{{{matchers}}}")?;
696        }
697        if let Some(at) = &self.at {
698            write!(f, " {at}")?;
699        }
700        if let Some(offset) = &self.offset {
701            write!(f, " offset {offset}")?;
702        }
703        Ok(())
704    }
705}
706
707impl Prettier for VectorSelector {
708    fn needs_split(&self, _max: usize) -> bool {
709        false
710    }
711}
712
713#[derive(Debug, Clone, PartialEq, Eq)]
714pub struct MatrixSelector {
715    pub vs: VectorSelector,
716    pub range: Duration,
717}
718
719impl fmt::Display for MatrixSelector {
720    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
721        if let Some(name) = &self.vs.name {
722            write!(f, "{name}")?;
723        }
724
725        let matchers = &self.vs.matchers.to_string();
726        if !matchers.is_empty() {
727            write!(f, "{{{matchers}}}")?;
728        }
729
730        write!(f, "[{}]", display_duration(&self.range))?;
731
732        if let Some(at) = &self.vs.at {
733            write!(f, " {at}")?;
734        }
735
736        if let Some(offset) = &self.vs.offset {
737            write!(f, " offset {offset}")?;
738        }
739
740        Ok(())
741    }
742}
743
744impl Prettier for MatrixSelector {
745    fn needs_split(&self, _max: usize) -> bool {
746        false
747    }
748}
749
750/// Call represents Prometheus Function.
751/// Some functions have special cases:
752///
753/// ## exp
754///
755/// exp(v instant-vector) calculates the exponential function for all elements in v.
756/// Special cases are:
757///
758/// ```promql
759/// Exp(+Inf) = +Inf
760/// Exp(NaN) = NaN
761/// ```
762///
763/// ## ln
764///
765/// ln(v instant-vector) calculates the natural logarithm for all elements in v.
766/// Special cases are:
767///
768/// ```promql
769/// ln(+Inf) = +Inf
770/// ln(0) = -Inf
771/// ln(x < 0) = NaN
772/// ln(NaN) = NaN
773/// ```
774///
775/// TODO: support more special cases of function call
776///
777///  - acos()
778///  - acosh()
779///  - asin()
780///  - asinh()
781///  - atan()
782///  - atanh()
783///  - cos()
784///  - cosh()
785///  - sin()
786///  - sinh()
787///  - tan()
788///  - tanh()
789#[derive(Debug, Clone, PartialEq, Eq)]
790pub struct Call {
791    pub func: Function,
792    pub args: FunctionArgs,
793}
794
795impl fmt::Display for Call {
796    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
797        write!(f, "{}({})", self.func.name, self.args)
798    }
799}
800
801impl Prettier for Call {
802    fn format(&self, level: usize, max: usize) -> String {
803        format!(
804            "{}{}(\n{}\n{})",
805            indent(level),
806            self.func.name,
807            self.args.pretty(level + 1, max),
808            indent(level)
809        )
810    }
811}
812
813/// Node for extending the AST. [Extension] won't be generate by this parser itself.
814#[derive(Debug, Clone)]
815pub struct Extension {
816    pub expr: Arc<dyn ExtensionExpr>,
817}
818
819/// The interface for extending the AST with custom expression node.
820pub trait ExtensionExpr: std::fmt::Debug + Send + Sync {
821    fn as_any(&self) -> &dyn std::any::Any;
822
823    fn name(&self) -> &str;
824
825    fn value_type(&self) -> ValueType;
826
827    #[allow(clippy::mut_from_ref)]
828    fn children(&self) -> &mut [Expr];
829}
830
831impl PartialEq for Extension {
832    fn eq(&self, other: &Self) -> bool {
833        format!("{:?}", self) == format!("{:?}", other)
834    }
835}
836
837impl Eq for Extension {}
838
839#[derive(Debug, Clone, PartialEq, Eq)]
840pub enum Expr {
841    /// Aggregate represents an aggregation operation on a Vector.
842    Aggregate(AggregateExpr),
843
844    /// Unary represents a unary operation on another expression.
845    /// Currently unary operations are only supported for Scalars.
846    Unary(UnaryExpr),
847
848    /// Binary represents a binary expression between two child expressions.
849    Binary(BinaryExpr),
850
851    /// Paren wraps an expression so it cannot be disassembled as a consequence
852    /// of operator precedence.
853    Paren(ParenExpr),
854
855    /// SubqueryExpr represents a subquery.
856    Subquery(SubqueryExpr),
857
858    /// NumberLiteral represents a number.
859    NumberLiteral(NumberLiteral),
860
861    /// StringLiteral represents a string.
862    StringLiteral(StringLiteral),
863
864    /// VectorSelector represents a Vector selection.
865    VectorSelector(VectorSelector),
866
867    /// MatrixSelector represents a Matrix selection.
868    MatrixSelector(MatrixSelector),
869
870    /// Call represents a function call.
871    Call(Call),
872
873    /// Extension represents an extension expression. It is for user to attach additional
874    /// information to the AST. This parser won't generate Extension node.
875    Extension(Extension),
876}
877
878impl Expr {
879    pub(crate) fn new_vector_selector(
880        name: Option<String>,
881        matchers: Matchers,
882    ) -> Result<Self, String> {
883        let vs = VectorSelector::new(name, matchers);
884        Ok(Self::VectorSelector(vs))
885    }
886
887    pub(crate) fn new_unary_expr(expr: Expr) -> Result<Self, String> {
888        match expr {
889            Expr::StringLiteral(_) => Err("unary expression only allowed on expressions of type scalar or vector, got: string".into()),
890            Expr::MatrixSelector(_) => Err("unary expression only allowed on expressions of type scalar or vector, got: matrix".into()),
891            _ => Ok(-expr),
892        }
893    }
894
895    pub(crate) fn new_subquery_expr(
896        expr: Expr,
897        range: Duration,
898        step: Option<Duration>,
899    ) -> Result<Self, String> {
900        let se = Expr::Subquery(SubqueryExpr {
901            expr: Box::new(expr),
902            offset: None,
903            at: None,
904            range,
905            step,
906        });
907        Ok(se)
908    }
909
910    pub(crate) fn new_paren_expr(expr: Expr) -> Result<Self, String> {
911        let ex = Expr::Paren(ParenExpr {
912            expr: Box::new(expr),
913        });
914        Ok(ex)
915    }
916
917    /// NOTE: @ and offset is not set here.
918    pub(crate) fn new_matrix_selector(expr: Expr, range: Duration) -> Result<Self, String> {
919        match expr {
920            Expr::VectorSelector(VectorSelector {
921                offset: Some(_), ..
922            }) => Err("no offset modifiers allowed before range".into()),
923            Expr::VectorSelector(VectorSelector { at: Some(_), .. }) => {
924                Err("no @ modifiers allowed before range".into())
925            }
926            Expr::VectorSelector(vs) => {
927                let ms = Expr::MatrixSelector(MatrixSelector { vs, range });
928                Ok(ms)
929            }
930            _ => Err("ranges only allowed for vector selectors".into()),
931        }
932    }
933
934    pub(crate) fn at_expr(self, at: AtModifier) -> Result<Self, String> {
935        let already_set_err = Err("@ <timestamp> may not be set multiple times".into());
936        match self {
937            Expr::VectorSelector(mut vs) => match vs.at {
938                None => {
939                    vs.at = Some(at);
940                    Ok(Expr::VectorSelector(vs))
941                }
942                Some(_) => already_set_err,
943            },
944            Expr::MatrixSelector(mut ms) => match ms.vs.at {
945                None => {
946                    ms.vs.at = Some(at);
947                    Ok(Expr::MatrixSelector(ms))
948                }
949                Some(_) => already_set_err,
950            },
951            Expr::Subquery(mut s) => match s.at {
952                None => {
953                    s.at = Some(at);
954                    Ok(Expr::Subquery(s))
955                }
956                Some(_) => already_set_err,
957            },
958            _ => {
959                Err("@ modifier must be preceded by an vector selector or matrix selector or a subquery".into())
960            }
961        }
962    }
963
964    /// set offset field for specified Expr, but CAN ONLY be set once.
965    pub(crate) fn offset_expr(self, offset: Offset) -> Result<Self, String> {
966        let already_set_err = Err("offset may not be set multiple times".into());
967        match self {
968            Expr::VectorSelector(mut vs) => match vs.offset {
969                None => {
970                    vs.offset = Some(offset);
971                    Ok(Expr::VectorSelector(vs))
972                }
973                Some(_) => already_set_err,
974            },
975            Expr::MatrixSelector(mut ms) => match ms.vs.offset {
976                None => {
977                    ms.vs.offset = Some(offset);
978                    Ok(Expr::MatrixSelector(ms))
979                }
980                Some(_) => already_set_err,
981            },
982            Expr::Subquery(mut s) => match s.offset {
983                None => {
984                    s.offset = Some(offset);
985                    Ok(Expr::Subquery(s))
986                }
987                Some(_) => already_set_err,
988            },
989            _ => {
990                Err("offset modifier must be preceded by an vector selector or matrix selector or a subquery".into())
991            }
992        }
993    }
994
995    pub(crate) fn new_call(func: Function, args: FunctionArgs) -> Result<Expr, String> {
996        Ok(Expr::Call(Call { func, args }))
997    }
998
999    pub(crate) fn new_binary_expr(
1000        lhs: Expr,
1001        op: TokenId,
1002        modifier: Option<BinModifier>,
1003        rhs: Expr,
1004    ) -> Result<Expr, String> {
1005        let ex = BinaryExpr {
1006            op: TokenType::new(op),
1007            lhs: Box::new(lhs),
1008            rhs: Box::new(rhs),
1009            modifier,
1010        };
1011        Ok(Expr::Binary(ex))
1012    }
1013
1014    pub(crate) fn new_aggregate_expr(
1015        op: TokenId,
1016        modifier: Option<LabelModifier>,
1017        args: FunctionArgs,
1018    ) -> Result<Expr, String> {
1019        let op = TokenType::new(op);
1020        if args.is_empty() {
1021            return Err(format!(
1022                "no arguments for aggregate expression '{op}' provided"
1023            ));
1024        }
1025        let mut desired_args_count = 1;
1026        let mut param = None;
1027        if op.is_aggregator_with_param() {
1028            desired_args_count = 2;
1029            param = args.first();
1030        }
1031        if args.len() != desired_args_count {
1032            return Err(format!(
1033                "wrong number of arguments for aggregate expression provided, expected {}, got {}",
1034                desired_args_count,
1035                args.len()
1036            ));
1037        }
1038
1039        match args.last() {
1040            Some(expr) => Ok(Expr::Aggregate(AggregateExpr {
1041                op,
1042                expr,
1043                param,
1044                modifier,
1045            })),
1046            None => Err(
1047                "aggregate operation needs a single instant vector parameter, but found none"
1048                    .into(),
1049            ),
1050        }
1051    }
1052
1053    pub fn value_type(&self) -> ValueType {
1054        match self {
1055            Expr::Aggregate(_) => ValueType::Vector,
1056            Expr::Unary(ex) => ex.expr.value_type(),
1057            Expr::Binary(ex) => {
1058                if ex.lhs.value_type() == ValueType::Scalar
1059                    && ex.rhs.value_type() == ValueType::Scalar
1060                {
1061                    ValueType::Scalar
1062                } else {
1063                    ValueType::Vector
1064                }
1065            }
1066            Expr::Paren(ex) => ex.expr.value_type(),
1067            Expr::Subquery(_) => ValueType::Matrix,
1068            Expr::NumberLiteral(_) => ValueType::Scalar,
1069            Expr::StringLiteral(_) => ValueType::String,
1070            Expr::VectorSelector(_) => ValueType::Vector,
1071            Expr::MatrixSelector(_) => ValueType::Matrix,
1072            Expr::Call(ex) => ex.func.return_type,
1073            Expr::Extension(ex) => ex.expr.value_type(),
1074        }
1075    }
1076
1077    /// only Some if expr is [Expr::NumberLiteral]
1078    pub(crate) fn scalar_value(&self) -> Option<f64> {
1079        match self {
1080            Expr::NumberLiteral(nl) => Some(nl.val),
1081            _ => None,
1082        }
1083    }
1084
1085    pub fn prettify(&self) -> String {
1086        self.pretty(0, MAX_CHARACTERS_PER_LINE)
1087    }
1088}
1089
1090impl From<String> for Expr {
1091    fn from(val: String) -> Self {
1092        Expr::StringLiteral(StringLiteral { val })
1093    }
1094}
1095
1096impl From<&str> for Expr {
1097    fn from(s: &str) -> Self {
1098        Expr::StringLiteral(StringLiteral { val: s.into() })
1099    }
1100}
1101
1102impl From<f64> for Expr {
1103    fn from(val: f64) -> Self {
1104        Expr::NumberLiteral(NumberLiteral { val })
1105    }
1106}
1107
1108/// directly create an Expr::VectorSelector from instant vector
1109///
1110/// # Examples
1111///
1112/// Basic usage:
1113///
1114/// ``` rust
1115/// use bk_promql_parser::label::Matchers;
1116/// use bk_promql_parser::parser::{Expr, VectorSelector};
1117///
1118/// let name = String::from("foo");
1119/// let vs = VectorSelector::new(Some(name), Matchers::empty());
1120///
1121/// assert_eq!(Expr::VectorSelector(vs), Expr::from(VectorSelector::from("foo")));
1122/// ```
1123impl From<VectorSelector> for Expr {
1124    fn from(vs: VectorSelector) -> Self {
1125        Expr::VectorSelector(vs)
1126    }
1127}
1128
1129impl Neg for Expr {
1130    type Output = Self;
1131
1132    fn neg(self) -> Self::Output {
1133        match self {
1134            Expr::NumberLiteral(nl) => Expr::NumberLiteral(-nl),
1135            _ => Expr::Unary(UnaryExpr {
1136                expr: Box::new(self),
1137            }),
1138        }
1139    }
1140}
1141
1142impl fmt::Display for Expr {
1143    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
1144        match self {
1145            Expr::Aggregate(ex) => write!(f, "{ex}"),
1146            Expr::Unary(ex) => write!(f, "{ex}"),
1147            Expr::Binary(ex) => write!(f, "{ex}"),
1148            Expr::Paren(ex) => write!(f, "{ex}"),
1149            Expr::Subquery(ex) => write!(f, "{ex}"),
1150            Expr::NumberLiteral(ex) => write!(f, "{ex}"),
1151            Expr::StringLiteral(ex) => write!(f, "{ex}"),
1152            Expr::VectorSelector(ex) => write!(f, "{ex}"),
1153            Expr::MatrixSelector(ex) => write!(f, "{ex}"),
1154            Expr::Call(ex) => write!(f, "{ex}"),
1155            Expr::Extension(ext) => write!(f, "{ext:?}"),
1156        }
1157    }
1158}
1159
1160impl Prettier for Expr {
1161    fn pretty(&self, level: usize, max: usize) -> String {
1162        match self {
1163            Expr::Aggregate(ex) => ex.pretty(level, max),
1164            Expr::Unary(ex) => ex.pretty(level, max),
1165            Expr::Binary(ex) => ex.pretty(level, max),
1166            Expr::Paren(ex) => ex.pretty(level, max),
1167            Expr::Subquery(ex) => ex.pretty(level, max),
1168            Expr::NumberLiteral(ex) => ex.pretty(level, max),
1169            Expr::StringLiteral(ex) => ex.pretty(level, max),
1170            Expr::VectorSelector(ex) => ex.pretty(level, max),
1171            Expr::MatrixSelector(ex) => ex.pretty(level, max),
1172            Expr::Call(ex) => ex.pretty(level, max),
1173            Expr::Extension(ext) => format!("{ext:?}"),
1174        }
1175    }
1176}
1177
1178/// check_ast checks the validity of the provided AST. This includes type checking.
1179/// Recursively check correct typing for child nodes and raise errors in case of bad typing.
1180pub(crate) fn check_ast(expr: Expr) -> Result<Expr, String> {
1181    match expr {
1182        Expr::Binary(ex) => check_ast_for_binary_expr(ex),
1183        Expr::Aggregate(ex) => check_ast_for_aggregate_expr(ex),
1184        Expr::Call(ex) => check_ast_for_call(ex),
1185        Expr::Unary(ex) => check_ast_for_unary(ex),
1186        Expr::Subquery(ex) => check_ast_for_subquery(ex),
1187        Expr::VectorSelector(ex) => check_ast_for_vector_selector(ex),
1188        Expr::Paren(_) => Ok(expr),
1189        Expr::NumberLiteral(_) => Ok(expr),
1190        Expr::StringLiteral(_) => Ok(expr),
1191        Expr::MatrixSelector(_) => Ok(expr),
1192        Expr::Extension(_) => Ok(expr),
1193    }
1194}
1195
1196fn expect_type(
1197    expected: ValueType,
1198    actual: Option<ValueType>,
1199    context: &str,
1200) -> Result<bool, String> {
1201    match actual {
1202        Some(actual) => {
1203            if actual == expected {
1204                Ok(true)
1205            } else {
1206                Err(format!(
1207                    "expected type {expected} in {context}, got {actual}"
1208                ))
1209            }
1210        }
1211        None => Err(format!("expected type {expected} in {context}, got None")),
1212    }
1213}
1214
1215/// the original logic is redundant in prometheus, and the following coding blocks
1216/// have been optimized for readability, but all logic SHOULD be covered.
1217fn check_ast_for_binary_expr(mut ex: BinaryExpr) -> Result<Expr, String> {
1218    if !ex.op.is_operator() {
1219        return Err(format!(
1220            "binary expression does not support operator '{}'",
1221            ex.op
1222        ));
1223    }
1224
1225    if ex.return_bool() && !ex.op.is_comparison_operator() {
1226        return Err("bool modifier can only be used on comparison operators".into());
1227    }
1228
1229    if ex.op.is_comparison_operator()
1230        && ex.lhs.value_type() == ValueType::Scalar
1231        && ex.rhs.value_type() == ValueType::Scalar
1232        && !ex.return_bool()
1233    {
1234        return Err("comparisons between scalars must use BOOL modifier".into());
1235    }
1236
1237    // For `on` matching, a label can only appear in one of the lists.
1238    // Every time series of the result vector must be uniquely identifiable.
1239    if ex.is_matching_on() && ex.is_labels_joint() {
1240        if let Some(labels) = ex.intersect_labels() {
1241            if let Some(label) = labels.first() {
1242                return Err(format!(
1243                    "label '{label}' must not occur in ON and GROUP clause at once"
1244                ));
1245            }
1246        };
1247    }
1248
1249    if ex.op.is_set_operator() {
1250        if ex.lhs.value_type() == ValueType::Scalar || ex.rhs.value_type() == ValueType::Scalar {
1251            return Err(format!(
1252                "set operator '{}' not allowed in binary scalar expression",
1253                ex.op
1254            ));
1255        }
1256
1257        if ex.lhs.value_type() == ValueType::Vector && ex.rhs.value_type() == ValueType::Vector {
1258            if let Some(ref modifier) = ex.modifier {
1259                if matches!(modifier.card, VectorMatchCardinality::OneToMany(_))
1260                    || matches!(modifier.card, VectorMatchCardinality::ManyToOne(_))
1261                {
1262                    return Err(format!("no grouping allowed for '{}' operation", ex.op));
1263                }
1264            };
1265        }
1266
1267        match &mut ex.modifier {
1268            Some(modifier) => {
1269                if modifier.card == VectorMatchCardinality::OneToOne {
1270                    modifier.card = VectorMatchCardinality::ManyToMany;
1271                }
1272            }
1273            None => {
1274                ex.modifier =
1275                    Some(BinModifier::default().with_card(VectorMatchCardinality::ManyToMany));
1276            }
1277        }
1278    }
1279
1280    if ex.lhs.value_type() != ValueType::Scalar && ex.lhs.value_type() != ValueType::Vector {
1281        return Err("binary expression must contain only scalar and instant vector types".into());
1282    }
1283    if ex.rhs.value_type() != ValueType::Scalar && ex.rhs.value_type() != ValueType::Vector {
1284        return Err("binary expression must contain only scalar and instant vector types".into());
1285    }
1286
1287    if (ex.lhs.value_type() != ValueType::Vector || ex.rhs.value_type() != ValueType::Vector)
1288        && ex.is_matching_labels_not_empty()
1289    {
1290        return Err("vector matching only allowed between vectors".into());
1291    }
1292
1293    Ok(Expr::Binary(ex))
1294}
1295
1296fn check_ast_for_aggregate_expr(ex: AggregateExpr) -> Result<Expr, String> {
1297    if !ex.op.is_aggregator() {
1298        return Err(format!(
1299            "aggregation operator expected in aggregation expression but got '{}'",
1300            ex.op
1301        ));
1302    }
1303
1304    expect_type(
1305        ValueType::Vector,
1306        Some(ex.expr.value_type()),
1307        "aggregation expression",
1308    )?;
1309
1310    if matches!(ex.op.id(), T_TOPK | T_BOTTOMK | T_QUANTILE) {
1311        expect_type(
1312            ValueType::Scalar,
1313            ex.param.as_ref().map(|ex| ex.value_type()),
1314            "aggregation expression",
1315        )?;
1316    }
1317
1318    if ex.op.id() == T_COUNT_VALUES {
1319        expect_type(
1320            ValueType::String,
1321            ex.param.as_ref().map(|ex| ex.value_type()),
1322            "aggregation expression",
1323        )?;
1324    }
1325
1326    Ok(Expr::Aggregate(ex))
1327}
1328
1329fn check_ast_for_call(ex: Call) -> Result<Expr, String> {
1330    let expected_args_len = ex.func.arg_types.len();
1331    let name = ex.func.name;
1332    let actual_args_len = ex.args.len();
1333
1334    if ex.func.variadic {
1335        let expected_args_len_without_default = expected_args_len - 1;
1336        if expected_args_len_without_default > actual_args_len {
1337            return Err(format!(
1338                "expected at least {expected_args_len_without_default} argument(s) in call to '{name}', got {actual_args_len}"
1339            ));
1340        }
1341
1342        // `label_join` do not have a maximum arguments threshold.
1343        // this hard code SHOULD be careful if new functions are supported by Prometheus.
1344        if actual_args_len > expected_args_len && name.ne("label_join") {
1345            return Err(format!(
1346                "expected at most {expected_args_len} argument(s) in call to '{name}', got {actual_args_len}"
1347            ));
1348        }
1349    }
1350
1351    if !ex.func.variadic && expected_args_len != actual_args_len {
1352        return Err(format!(
1353            "expected {expected_args_len} argument(s) in call to '{name}', got {actual_args_len}"
1354        ));
1355    }
1356
1357    // special cases from https://prometheus.io/docs/prometheus/latest/querying/functions
1358    if name.eq("exp") {
1359        if let Some(val) = ex.args.first().and_then(|ex| ex.scalar_value()) {
1360            if val.is_nan() || val.is_infinite() {
1361                return Ok(Expr::Call(ex));
1362            }
1363        }
1364    } else if name.eq("ln") || name.eq("log2") || name.eq("log10") {
1365        if let Some(val) = ex.args.first().and_then(|ex| ex.scalar_value()) {
1366            if val.is_nan() || val.is_infinite() || val <= 0.0 {
1367                return Ok(Expr::Call(ex));
1368            }
1369        }
1370    }
1371
1372    for (mut idx, actual_arg) in ex.args.args.iter().enumerate() {
1373        // this only happens when function args are variadic
1374        if idx >= ex.func.arg_types.len() {
1375            idx = ex.func.arg_types.len() - 1;
1376        }
1377
1378        expect_type(
1379            ex.func.arg_types[idx],
1380            Some(actual_arg.value_type()),
1381            &format!("call to function '{name}'"),
1382        )?;
1383    }
1384
1385    Ok(Expr::Call(ex))
1386}
1387
1388fn check_ast_for_unary(ex: UnaryExpr) -> Result<Expr, String> {
1389    let value_type = ex.expr.value_type();
1390    if value_type != ValueType::Scalar && value_type != ValueType::Vector {
1391        return Err(format!(
1392            "unary expression only allowed on expressions of type scalar or vector, got {value_type}"
1393        ));
1394    }
1395
1396    Ok(Expr::Unary(ex))
1397}
1398
1399fn check_ast_for_subquery(ex: SubqueryExpr) -> Result<Expr, String> {
1400    let value_type = ex.expr.value_type();
1401    if value_type != ValueType::Vector {
1402        return Err(format!(
1403            "subquery is only allowed on vector, got {value_type} instead"
1404        ));
1405    }
1406
1407    Ok(Expr::Subquery(ex))
1408}
1409
1410fn check_ast_for_vector_selector(ex: VectorSelector) -> Result<Expr, String> {
1411    match ex.name {
1412        Some(ref name) => match ex.matchers.find_matcher_value(METRIC_NAME) {
1413            Some(val) => Err(format!(
1414                "metric name must not be set twice: '{}' or '{}'",
1415                name, val
1416            )),
1417            None => Ok(Expr::VectorSelector(ex)),
1418        },
1419        _ => Ok(Expr::VectorSelector(ex)),
1420    }
1421}
1422
1423#[cfg(test)]
1424mod tests {
1425    use super::*;
1426    use crate::label::{MatchOp, Matcher};
1427
1428    #[test]
1429    fn test_valid_at_modifier() {
1430        let cases = vec![
1431            // tuple: (seconds, elapsed milliseconds before or after UNIX_EPOCH)
1432            (0.0, 0),
1433            (1000.3, 1000300),    // after UNIX_EPOCH
1434            (1000.9, 1000900),    // after UNIX_EPOCH
1435            (1000.9991, 1000999), // after UNIX_EPOCH
1436            (1000.9999, 1001000), // after UNIX_EPOCH
1437            (-1000.3, 1000300),   // before UNIX_EPOCH
1438            (-1000.9, 1000900),   // before UNIX_EPOCH
1439        ];
1440
1441        for (secs, elapsed) in cases {
1442            match AtModifier::try_from(secs).unwrap() {
1443                AtModifier::At(st) => {
1444                    if secs.is_sign_positive() || secs == 0.0 {
1445                        assert_eq!(
1446                            elapsed,
1447                            st.duration_since(SystemTime::UNIX_EPOCH)
1448                                .unwrap()
1449                                .as_millis()
1450                        )
1451                    } else if secs.is_sign_negative() {
1452                        assert_eq!(
1453                            elapsed,
1454                            SystemTime::UNIX_EPOCH
1455                                .duration_since(st)
1456                                .unwrap()
1457                                .as_millis()
1458                        )
1459                    }
1460                }
1461                _ => panic!(),
1462            }
1463        }
1464
1465        assert_eq!(
1466            AtModifier::try_from(Expr::from(1.0)),
1467            AtModifier::try_from(1.0),
1468        );
1469    }
1470
1471    #[test]
1472    fn test_invalid_at_modifier() {
1473        let cases = vec![
1474            f64::MAX,
1475            f64::MIN,
1476            f64::NAN,
1477            f64::INFINITY,
1478            f64::NEG_INFINITY,
1479        ];
1480
1481        for secs in cases {
1482            assert!(AtModifier::try_from(secs).is_err())
1483        }
1484
1485        assert_eq!(
1486            AtModifier::try_from(token::T_ADD),
1487            Err("invalid @ modifier preprocessor '+', START or END is valid.".into())
1488        );
1489
1490        assert_eq!(
1491            AtModifier::try_from(Expr::from("string literal")),
1492            Err("invalid float value after @ modifier".into())
1493        );
1494    }
1495
1496    #[test]
1497    fn test_binary_labels() {
1498        assert_eq!(
1499            &Labels::new(vec!["foo", "bar"]),
1500            LabelModifier::Include(Labels::new(vec!["foo", "bar"])).labels()
1501        );
1502
1503        assert_eq!(
1504            &Labels::new(vec!["foo", "bar"]),
1505            LabelModifier::Exclude(Labels::new(vec!["foo", "bar"])).labels()
1506        );
1507
1508        assert_eq!(
1509            &Labels::new(vec!["foo", "bar"]),
1510            VectorMatchCardinality::OneToMany(Labels::new(vec!["foo", "bar"]))
1511                .labels()
1512                .unwrap()
1513        );
1514
1515        assert_eq!(
1516            &Labels::new(vec!["foo", "bar"]),
1517            VectorMatchCardinality::ManyToOne(Labels::new(vec!["foo", "bar"]))
1518                .labels()
1519                .unwrap()
1520        );
1521
1522        assert_eq!(VectorMatchCardinality::OneToOne.labels(), None);
1523        assert_eq!(VectorMatchCardinality::ManyToMany.labels(), None);
1524    }
1525
1526    #[test]
1527    fn test_neg() {
1528        assert_eq!(
1529            -VectorSelector::from("foo"),
1530            UnaryExpr {
1531                expr: Box::new(Expr::from(VectorSelector::from("foo")))
1532            }
1533        )
1534    }
1535
1536    #[test]
1537    fn test_scalar_value() {
1538        assert_eq!(Some(1.0), Expr::from(1.0).scalar_value());
1539        assert_eq!(None, Expr::from("1.0").scalar_value());
1540    }
1541
1542    #[test]
1543    fn test_at_expr() {
1544        assert_eq!(
1545            "@ <timestamp> may not be set multiple times",
1546            Expr::from(VectorSelector::from("foo"))
1547                .at_expr(AtModifier::try_from(1.0).unwrap())
1548                .and_then(|ex| ex.at_expr(AtModifier::try_from(1.0).unwrap()))
1549                .unwrap_err()
1550        );
1551
1552        assert_eq!(
1553            "@ <timestamp> may not be set multiple times",
1554            Expr::new_matrix_selector(
1555                Expr::from(VectorSelector::from("foo")),
1556                Duration::from_secs(1),
1557            )
1558            .and_then(|ex| ex.at_expr(AtModifier::try_from(1.0).unwrap()))
1559            .and_then(|ex| ex.at_expr(AtModifier::try_from(1.0).unwrap()))
1560            .unwrap_err()
1561        );
1562
1563        assert_eq!(
1564            "@ <timestamp> may not be set multiple times",
1565            Expr::new_subquery_expr(
1566                Expr::from(VectorSelector::from("foo")),
1567                Duration::from_secs(1),
1568                None,
1569            )
1570            .and_then(|ex| ex.at_expr(AtModifier::try_from(1.0).unwrap()))
1571            .and_then(|ex| ex.at_expr(AtModifier::try_from(1.0).unwrap()))
1572            .unwrap_err()
1573        )
1574    }
1575
1576    #[test]
1577    fn test_offset_expr() {
1578        assert_eq!(
1579            "offset may not be set multiple times",
1580            Expr::from(VectorSelector::from("foo"))
1581                .offset_expr(Offset::Pos(Duration::from_secs(1000)))
1582                .and_then(|ex| ex.offset_expr(Offset::Pos(Duration::from_secs(1000))))
1583                .unwrap_err()
1584        );
1585
1586        assert_eq!(
1587            "offset may not be set multiple times",
1588            Expr::new_matrix_selector(
1589                Expr::from(VectorSelector::from("foo")),
1590                Duration::from_secs(1),
1591            )
1592            .and_then(|ex| ex.offset_expr(Offset::Pos(Duration::from_secs(1000))))
1593            .and_then(|ex| ex.offset_expr(Offset::Pos(Duration::from_secs(1000))))
1594            .unwrap_err()
1595        );
1596
1597        assert_eq!(
1598            "offset may not be set multiple times",
1599            Expr::new_subquery_expr(
1600                Expr::from(VectorSelector::from("foo")),
1601                Duration::from_secs(1),
1602                None,
1603            )
1604            .and_then(|ex| ex.offset_expr(Offset::Pos(Duration::from_secs(1000))))
1605            .and_then(|ex| ex.offset_expr(Offset::Pos(Duration::from_secs(1000))))
1606            .unwrap_err()
1607        );
1608    }
1609
1610    #[test]
1611    fn test_expr_to_string() {
1612        let mut cases = vec![
1613            ("1", "1"),
1614            ("- 1", "-1"),
1615            ("+ 1", "1"),
1616            ("Inf", "Inf"),
1617            ("inf", "Inf"),
1618            ("+Inf", "Inf"),
1619            ("- Inf", "-Inf"),
1620            (".5", "0.5"),
1621            ("5.", "5"),
1622            ("123.4567", "123.4567"),
1623            ("5e-3", "0.005"),
1624            ("5e3", "5000"),
1625            ("0xc", "12"),
1626            ("0755", "493"),
1627            ("08", "8"),
1628            ("+5.5e-3", "0.0055"),
1629            ("-0755", "-493"),
1630            ("NaN", "NaN"),
1631            ("NAN", "NaN"),
1632            ("- 1^2", "-1 ^ 2"),
1633            ("+1 + -2 * 1", "1 + -2 * 1"),
1634            ("1 + 2/(3*1)", "1 + 2 / (3 * 1)"),
1635            ("foo*sum", "foo * sum"),
1636            ("foo * on(test,blub) bar", "foo * on (test, blub) bar"),
1637            (
1638                r#"up{job="hi", instance="in"} offset 5m @ 100"#,
1639                r#"up{instance="in",job="hi"} @ 100.000 offset 5m"#,
1640            ),
1641            (
1642                r#"up{job="hi", instance="in"}"#,
1643                r#"up{instance="in",job="hi"}"#,
1644            ),
1645            ("sum (up) by (job,instance)", "sum by (job, instance) (up)"),
1646            (
1647                "foo / on(test,blub) group_left(bar) bar",
1648                "foo / on (test, blub) group_left (bar) bar",
1649            ),
1650            (
1651                "foo / on(test,blub) group_right(bar) bar",
1652                "foo / on (test, blub) group_right (bar) bar",
1653            ),
1654            (
1655                r#"foo{a="b",foo!="bar",test=~"test",bar!~"baz"}"#,
1656                r#"foo{a="b",bar!~"baz",foo!="bar",test=~"test"}"#,
1657            ),
1658            (
1659                r#"{__name__=~"foo.+",__name__=~".*bar"}"#,
1660                r#"{__name__=~".*bar",__name__=~"foo.+"}"#,
1661            ),
1662            (
1663                r#"test{a="b"}[5y] OFFSET 3d"#,
1664                r#"test{a="b"}[5y] offset 3d"#,
1665            ),
1666            (
1667                r#"{a="b"}[5y] OFFSET 3d"#,
1668                r#"{a="b"}[5y] offset 3d"#,
1669            ),
1670            (
1671                "sum(some_metric) without(and, by, avg, count, alert, annotations)",
1672                "sum without (and, by, avg, count, alert, annotations) (some_metric)",
1673            ),
1674            (
1675                r#"floor(some_metric{foo!="bar"})"#,
1676                r#"floor(some_metric{foo!="bar"})"#,
1677            ),
1678            (
1679                "sum(rate(http_request_duration_seconds[10m])) / count(rate(http_request_duration_seconds[10m]))",
1680                "sum(rate(http_request_duration_seconds[10m])) / count(rate(http_request_duration_seconds[10m]))",
1681            ),
1682            ("rate(some_metric[5m])", "rate(some_metric[5m])"),
1683            ("round(some_metric,5)", "round(some_metric, 5)"),
1684            (
1685                r#"absent(sum(nonexistent{job="myjob"}))"#,
1686                r#"absent(sum(nonexistent{job="myjob"}))"#,
1687            ),
1688            (
1689                "histogram_quantile(0.9,rate(http_request_duration_seconds_bucket[10m]))",
1690                "histogram_quantile(0.9, rate(http_request_duration_seconds_bucket[10m]))",
1691            ),
1692            (
1693                "histogram_quantile(0.9,sum(rate(http_request_duration_seconds_bucket[10m])) by(job,le))",
1694                "histogram_quantile(0.9, sum by (job, le) (rate(http_request_duration_seconds_bucket[10m])))",
1695            ),
1696            (
1697                r#"label_join(up{job="api-server",src1="a",src2="b",src3="c"}, "foo", ",", "src1", "src2", "src3")"#,
1698                r#"label_join(up{job="api-server",src1="a",src2="b",src3="c"}, "foo", ",", "src1", "src2", "src3")"#,
1699            ),
1700            (
1701                r#"min_over_time(rate(foo{bar="baz"}[2s])[5m:])[4m:3s] @ 100"#,
1702                r#"min_over_time(rate(foo{bar="baz"}[2s])[5m:])[4m:3s] @ 100.000"#,
1703            ),
1704            (
1705                r#"min_over_time(rate(foo{bar="baz"}[2s])[5m:])[4m:3s]"#,
1706                r#"min_over_time(rate(foo{bar="baz"}[2s])[5m:])[4m:3s]"#,
1707            ),
1708            (
1709                r#"min_over_time(rate(foo{bar="baz"}[2s])[5m:] offset 4m)[4m:3s]"#,
1710                r#"min_over_time(rate(foo{bar="baz"}[2s])[5m:] offset 4m)[4m:3s]"#,
1711            ),
1712            ("some_metric OFFSET 1m [10m:5s]", "some_metric offset 1m[10m:5s]"),
1713            ("some_metric @123 [10m:5s]", "some_metric @ 123.000[10m:5s]")
1714        ];
1715
1716        // the following cases are from https://github.com/prometheus/prometheus/blob/main/promql/parser/printer_test.go
1717        let mut cases1 = vec![
1718            (
1719                r#"sum by() (task:errors:rate10s{job="s"})"#,
1720                r#"sum(task:errors:rate10s{job="s"})"#,
1721            ),
1722            (
1723                r#"sum by(code) (task:errors:rate10s{job="s"})"#,
1724                r#"sum by (code) (task:errors:rate10s{job="s"})"#,
1725            ),
1726            (
1727                r#"sum without() (task:errors:rate10s{job="s"})"#,
1728                r#"sum without () (task:errors:rate10s{job="s"})"#,
1729            ),
1730            (
1731                r#"sum without(instance) (task:errors:rate10s{job="s"})"#,
1732                r#"sum without (instance) (task:errors:rate10s{job="s"})"#,
1733            ),
1734            (
1735                r#"topk(5, task:errors:rate10s{job="s"})"#,
1736                r#"topk(5, task:errors:rate10s{job="s"})"#,
1737            ),
1738            (
1739                r#"count_values("value", task:errors:rate10s{job="s"})"#,
1740                r#"count_values("value", task:errors:rate10s{job="s"})"#,
1741            ),
1742            ("a - on() c", "a - on () c"),
1743            ("a - on(b) c", "a - on (b) c"),
1744            ("a - on(b) group_left(x) c", "a - on (b) group_left (x) c"),
1745            (
1746                "a - on(b) group_left(x, y) c",
1747                "a - on (b) group_left (x, y) c",
1748            ),
1749            ("a - on(b) group_left c", "a - on (b) group_left () c"),
1750            ("a - ignoring(b) c", "a - ignoring (b) c"),
1751            ("a - ignoring() c", "a - c"),
1752            ("up > bool 0", "up > bool 0"),
1753            ("a offset 1m", "a offset 1m"),
1754            ("a offset -7m", "a offset -7m"),
1755            (r#"a{c="d"}[5m] offset 1m"#, r#"a{c="d"}[5m] offset 1m"#),
1756            ("a[5m] offset 1m", "a[5m] offset 1m"),
1757            ("a[12m] offset -3m", "a[12m] offset -3m"),
1758            ("a[1h:5m] offset 1m", "a[1h:5m] offset 1m"),
1759            (r#"{__name__="a"}"#, r#"{__name__="a"}"#),
1760            (r#"a{b!="c"}[1m]"#, r#"a{b!="c"}[1m]"#),
1761            (r#"a{b=~"c"}[1m]"#, r#"a{b=~"c"}[1m]"#),
1762            (r#"a{b!~"c"}[1m]"#, r#"a{b!~"c"}[1m]"#),
1763            ("a @ 10", "a @ 10.000"),
1764            ("a[1m] @ 10", "a[1m] @ 10.000"),
1765            ("a @ start()", "a @ start()"),
1766            ("a @ end()", "a @ end()"),
1767            ("a[1m] @ start()", "a[1m] @ start()"),
1768            ("a[1m] @ end()", "a[1m] @ end()"),
1769        ];
1770
1771        cases.append(&mut cases1);
1772        for (input, expected) in cases {
1773            let expr = crate::parser::parse(input).unwrap();
1774            assert_eq!(expected, expr.to_string())
1775        }
1776    }
1777
1778    #[test]
1779    fn test_vector_selector_to_string() {
1780        let cases = vec![
1781            (VectorSelector::default(), ""),
1782            (VectorSelector::from("foobar"), "foobar"),
1783            (
1784                {
1785                    let name = Some(String::from("foobar"));
1786                    let matchers = Matchers::one(Matcher::new(MatchOp::Equal, "a", "x"));
1787                    VectorSelector::new(name, matchers)
1788                },
1789                r#"foobar{a="x"}"#,
1790            ),
1791            (
1792                {
1793                    let matchers = Matchers::new(vec![
1794                        Matcher::new(MatchOp::Equal, "a", "x"),
1795                        Matcher::new(MatchOp::Equal, "b", "y"),
1796                    ]);
1797                    VectorSelector::new(None, matchers)
1798                },
1799                r#"{a="x",b="y"}"#,
1800            ),
1801            (
1802                {
1803                    let matchers =
1804                        Matchers::one(Matcher::new(MatchOp::Equal, METRIC_NAME, "foobar"));
1805                    VectorSelector::new(None, matchers)
1806                },
1807                r#"{__name__="foobar"}"#,
1808            ),
1809        ];
1810
1811        for (vs, expect) in cases {
1812            assert_eq!(expect, vs.to_string())
1813        }
1814    }
1815
1816    #[test]
1817    fn test_aggregate_expr_pretty() {
1818        let cases = vec![
1819            ("sum(foo)", "sum(foo)"),
1820            (
1821                r#"sum by() (task:errors:rate10s{job="s"})"#,
1822                r#"sum(
1823  task:errors:rate10s{job="s"}
1824)"#,
1825            ),
1826            (
1827                r#"sum without(job,foo) (task:errors:rate10s{job="s"})"#,
1828                r#"sum without (job, foo) (
1829  task:errors:rate10s{job="s"}
1830)"#,
1831            ),
1832            (
1833                r#"sum(task:errors:rate10s{job="s"}) without(job,foo)"#,
1834                r#"sum without (job, foo) (
1835  task:errors:rate10s{job="s"}
1836)"#,
1837            ),
1838            (
1839                r#"sum by(job,foo) (task:errors:rate10s{job="s"})"#,
1840                r#"sum by (job, foo) (
1841  task:errors:rate10s{job="s"}
1842)"#,
1843            ),
1844            (
1845                r#"sum (task:errors:rate10s{job="s"}) by(job,foo)"#,
1846                r#"sum by (job, foo) (
1847  task:errors:rate10s{job="s"}
1848)"#,
1849            ),
1850            (
1851                r#"topk(10, ask:errors:rate10s{job="s"})"#,
1852                r#"topk(
1853  10,
1854  ask:errors:rate10s{job="s"}
1855)"#,
1856            ),
1857            (
1858                r#"sum by(job,foo) (sum by(job,foo) (task:errors:rate10s{job="s"}))"#,
1859                r#"sum by (job, foo) (
1860  sum by (job, foo) (
1861    task:errors:rate10s{job="s"}
1862  )
1863)"#,
1864            ),
1865            (
1866                r#"sum by(job,foo) (sum by(job,foo) (sum by(job,foo) (task:errors:rate10s{job="s"})))"#,
1867                r#"sum by (job, foo) (
1868  sum by (job, foo) (
1869    sum by (job, foo) (
1870      task:errors:rate10s{job="s"}
1871    )
1872  )
1873)"#,
1874            ),
1875            (
1876                r#"sum by(job,foo)
1877(sum by(job,foo) (task:errors:rate10s{job="s"}))"#,
1878                r#"sum by (job, foo) (
1879  sum by (job, foo) (
1880    task:errors:rate10s{job="s"}
1881  )
1882)"#,
1883            ),
1884            (
1885                r#"sum by(job,foo)
1886(sum(task:errors:rate10s{job="s"}) without(job,foo))"#,
1887                r#"sum by (job, foo) (
1888  sum without (job, foo) (
1889    task:errors:rate10s{job="s"}
1890  )
1891)"#,
1892            ),
1893            (
1894                r#"sum by(job,foo) # Comment 1.
1895(sum by(job,foo) ( # Comment 2.
1896task:errors:rate10s{job="s"}))"#,
1897                r#"sum by (job, foo) (
1898  sum by (job, foo) (
1899    task:errors:rate10s{job="s"}
1900  )
1901)"#,
1902            ),
1903        ];
1904
1905        for (input, expect) in cases {
1906            let expr = crate::parser::parse(input);
1907            assert_eq!(expect, expr.unwrap().pretty(0, 10));
1908        }
1909    }
1910
1911    #[test]
1912    fn test_binary_expr_pretty() {
1913        let cases = vec![
1914            ("a+b", "a + b"),
1915            (
1916                "a == bool 1",
1917                "  a
1918== bool
1919  1",
1920            ),
1921            (
1922                "a == 1024000",
1923                "  a
1924==
1925  1024000",
1926            ),
1927            (
1928                "a + ignoring(job) b",
1929                "  a
1930+ ignoring (job)
1931  b",
1932            ),
1933            (
1934                "foo_1 + foo_2",
1935                "  foo_1
1936+
1937  foo_2",
1938            ),
1939            (
1940                "foo_1 + foo_2 + foo_3",
1941                "    foo_1
1942  +
1943    foo_2
1944+
1945  foo_3",
1946            ),
1947            (
1948                "foo + baar + foo_3",
1949                "  foo + baar
1950+
1951  foo_3",
1952            ),
1953            (
1954                "foo_1 + foo_2 + foo_3 + foo_4",
1955                "      foo_1
1956    +
1957      foo_2
1958  +
1959    foo_3
1960+
1961  foo_4",
1962            ),
1963            (
1964                "foo_1 + ignoring(foo) foo_2 + ignoring(job) group_left foo_3 + on(instance) group_right foo_4",
1965
1966                 "      foo_1
1967    + ignoring (foo)
1968      foo_2
1969  + ignoring (job) group_left ()
1970    foo_3
1971+ on (instance) group_right ()
1972  foo_4",
1973            ),
1974        ];
1975
1976        for (input, expect) in cases {
1977            let expr = crate::parser::parse(input);
1978            assert_eq!(expect, expr.unwrap().pretty(0, 10));
1979        }
1980    }
1981
1982    #[test]
1983    fn test_call_expr_pretty() {
1984        let cases = vec![
1985            (
1986                "rate(foo[1m])",
1987                "rate(
1988  foo[1m]
1989)",
1990            ),
1991            (
1992                "sum_over_time(foo[1m])",
1993                "sum_over_time(
1994  foo[1m]
1995)",
1996            ),
1997            (
1998                "rate(long_vector_selector[10m:1m] @ start() offset 1m)",
1999                "rate(
2000  long_vector_selector[10m:1m] @ start() offset 1m
2001)",
2002            ),
2003            (
2004                "histogram_quantile(0.9, rate(foo[1m]))",
2005                "histogram_quantile(
2006  0.9,
2007  rate(
2008    foo[1m]
2009  )
2010)",
2011            ),
2012            (
2013                "histogram_quantile(0.9, rate(foo[1m] @ start()))",
2014                "histogram_quantile(
2015  0.9,
2016  rate(
2017    foo[1m] @ start()
2018  )
2019)",
2020            ),
2021            (
2022                "max_over_time(rate(demo_api_request_duration_seconds_count[1m])[1m:] @ start() offset 1m)",
2023                "max_over_time(
2024  rate(
2025    demo_api_request_duration_seconds_count[1m]
2026  )[1m:] @ start() offset 1m
2027)",
2028            ),
2029            (
2030                r#"label_replace(up{job="api-server",service="a:c"}, "foo", "$1", "service", "(.*):.*")"#,
2031                r#"label_replace(
2032  up{job="api-server",service="a:c"},
2033  "foo",
2034  "$1",
2035  "service",
2036  "(.*):.*"
2037)"#,
2038            ),
2039            (
2040                r#"label_replace(label_replace(up{job="api-server",service="a:c"}, "foo", "$1", "service", "(.*):.*"), "foo", "$1", "service", "(.*):.*")"#,
2041                r#"label_replace(
2042  label_replace(
2043    up{job="api-server",service="a:c"},
2044    "foo",
2045    "$1",
2046    "service",
2047    "(.*):.*"
2048  ),
2049  "foo",
2050  "$1",
2051  "service",
2052  "(.*):.*"
2053)"#,
2054            ),
2055        ];
2056
2057        for (input, expect) in cases {
2058            let expr = crate::parser::parse(input);
2059            assert_eq!(expect, expr.unwrap().pretty(0, 10));
2060        }
2061    }
2062
2063    #[test]
2064    fn test_paren_expr_pretty() {
2065        let cases = vec![
2066            ("(foo)", "(foo)"),
2067            (
2068                "(_foo_long_)",
2069                "(
2070  _foo_long_
2071)",
2072            ),
2073            (
2074                "((foo_long))",
2075                "(
2076  (foo_long)
2077)",
2078            ),
2079            (
2080                "((_foo_long_))",
2081                "(
2082  (
2083    _foo_long_
2084  )
2085)",
2086            ),
2087            (
2088                "(((foo_long)))",
2089                "(
2090  (
2091    (foo_long)
2092  )
2093)",
2094            ),
2095            ("(1 + 2)", "(1 + 2)"),
2096            (
2097                "(foo + bar)",
2098                "(
2099  foo + bar
2100)",
2101            ),
2102            (
2103                "(foo_long + bar_long)",
2104                "(
2105    foo_long
2106  +
2107    bar_long
2108)",
2109            ),
2110            (
2111                "(foo_long + bar_long + bar_2_long)",
2112                "(
2113      foo_long
2114    +
2115      bar_long
2116  +
2117    bar_2_long
2118)",
2119            ),
2120            (
2121                "((foo_long + bar_long) + bar_2_long)",
2122                "(
2123    (
2124        foo_long
2125      +
2126        bar_long
2127    )
2128  +
2129    bar_2_long
2130)",
2131            ),
2132            (
2133                "(1111 + 2222)",
2134                "(
2135    1111
2136  +
2137    2222
2138)",
2139            ),
2140            (
2141                "(sum_over_time(foo[1m]))",
2142                "(
2143  sum_over_time(
2144    foo[1m]
2145  )
2146)",
2147            ),
2148            (
2149                r#"(label_replace(up{job="api-server",service="a:c"}, "foo", "$1", "service", "(.*):.*"))"#,
2150                r#"(
2151  label_replace(
2152    up{job="api-server",service="a:c"},
2153    "foo",
2154    "$1",
2155    "service",
2156    "(.*):.*"
2157  )
2158)"#,
2159            ),
2160            (
2161                r#"(label_replace(label_replace(up{job="api-server",service="a:c"}, "foo", "$1", "service", "(.*):.*"), "foo", "$1", "service", "(.*):.*"))"#,
2162                r#"(
2163  label_replace(
2164    label_replace(
2165      up{job="api-server",service="a:c"},
2166      "foo",
2167      "$1",
2168      "service",
2169      "(.*):.*"
2170    ),
2171    "foo",
2172    "$1",
2173    "service",
2174    "(.*):.*"
2175  )
2176)"#,
2177            ),
2178            (
2179                r#"(label_replace(label_replace((up{job="api-server",service="a:c"}), "foo", "$1", "service", "(.*):.*"), "foo", "$1", "service", "(.*):.*"))"#,
2180                r#"(
2181  label_replace(
2182    label_replace(
2183      (
2184        up{job="api-server",service="a:c"}
2185      ),
2186      "foo",
2187      "$1",
2188      "service",
2189      "(.*):.*"
2190    ),
2191    "foo",
2192    "$1",
2193    "service",
2194    "(.*):.*"
2195  )
2196)"#,
2197            ),
2198        ];
2199
2200        for (input, expect) in cases {
2201            let expr = crate::parser::parse(input);
2202            assert_eq!(expect, expr.unwrap().pretty(0, 10));
2203        }
2204    }
2205
2206    #[test]
2207    fn test_unary_expr_pretty() {
2208        let cases = vec![
2209            ("-1", "-1"),
2210            ("-vector_selector", "-vector_selector"),
2211            (
2212                "(-vector_selector)",
2213                "(
2214  -vector_selector
2215)",
2216            ),
2217            (
2218                "-histogram_quantile(0.9,rate(foo[1m]))",
2219                "-histogram_quantile(
2220  0.9,
2221  rate(
2222    foo[1m]
2223  )
2224)",
2225            ),
2226            (
2227                "-histogram_quantile(0.99, sum by (le) (rate(foo[1m])))",
2228                "-histogram_quantile(
2229  0.99,
2230  sum by (le) (
2231    rate(
2232      foo[1m]
2233    )
2234  )
2235)",
2236            ),
2237            (
2238                "-histogram_quantile(0.9, -rate(foo[1m] @ start()))",
2239                "-histogram_quantile(
2240  0.9,
2241  -rate(
2242    foo[1m] @ start()
2243  )
2244)",
2245            ),
2246            (
2247                "(-histogram_quantile(0.9, -rate(foo[1m] @ start())))",
2248                "(
2249  -histogram_quantile(
2250    0.9,
2251    -rate(
2252      foo[1m] @ start()
2253    )
2254  )
2255)",
2256            ),
2257        ];
2258
2259        for (input, expect) in cases {
2260            let expr = crate::parser::parse(input);
2261            assert_eq!(expect, expr.unwrap().pretty(0, 10));
2262        }
2263    }
2264
2265    #[test]
2266    fn test_expr_pretty() {
2267        // Following queries have been taken from https://monitoring.mixins.dev/
2268        let cases = vec![
2269            (
2270                r#"(node_filesystem_avail_bytes{job="node",fstype!=""} / node_filesystem_size_bytes{job="node",fstype!=""} * 100 < 40 and predict_linear(node_filesystem_avail_bytes{job="node",fstype!=""}[6h], 24*60*60) < 0 and node_filesystem_readonly{job="node",fstype!=""} == 0)"#,
2271                r#"(
2272            node_filesystem_avail_bytes{fstype!="",job="node"}
2273          /
2274            node_filesystem_size_bytes{fstype!="",job="node"}
2275        *
2276          100
2277      <
2278        40
2279    and
2280        predict_linear(
2281          node_filesystem_avail_bytes{fstype!="",job="node"}[6h],
2282            24 * 60
2283          *
2284            60
2285        )
2286      <
2287        0
2288  and
2289      node_filesystem_readonly{fstype!="",job="node"}
2290    ==
2291      0
2292)"#,
2293            ),
2294            (
2295                r#"(node_filesystem_avail_bytes{job="node",fstype!=""} / node_filesystem_size_bytes{job="node",fstype!=""} * 100 < 20 and predict_linear(node_filesystem_avail_bytes{job="node",fstype!=""}[6h], 4*60*60) < 0 and node_filesystem_readonly{job="node",fstype!=""} == 0)"#,
2296                r#"(
2297            node_filesystem_avail_bytes{fstype!="",job="node"}
2298          /
2299            node_filesystem_size_bytes{fstype!="",job="node"}
2300        *
2301          100
2302      <
2303        20
2304    and
2305        predict_linear(
2306          node_filesystem_avail_bytes{fstype!="",job="node"}[6h],
2307            4 * 60
2308          *
2309            60
2310        )
2311      <
2312        0
2313  and
2314      node_filesystem_readonly{fstype!="",job="node"}
2315    ==
2316      0
2317)"#,
2318            ),
2319            (
2320                r#"(node_timex_offset_seconds > 0.05 and deriv(node_timex_offset_seconds[5m]) >= 0) or (node_timex_offset_seconds < -0.05 and deriv(node_timex_offset_seconds[5m]) <= 0)"#,
2321                r#"  (
2322        node_timex_offset_seconds
2323      >
2324        0.05
2325    and
2326        deriv(
2327          node_timex_offset_seconds[5m]
2328        )
2329      >=
2330        0
2331  )
2332or
2333  (
2334        node_timex_offset_seconds
2335      <
2336        -0.05
2337    and
2338        deriv(
2339          node_timex_offset_seconds[5m]
2340        )
2341      <=
2342        0
2343  )"#,
2344            ),
2345            (
2346                r#"1 - ((node_memory_MemAvailable_bytes{job="node"} or (node_memory_Buffers_bytes{job="node"} + node_memory_Cached_bytes{job="node"} + node_memory_MemFree_bytes{job="node"} + node_memory_Slab_bytes{job="node"}) ) / node_memory_MemTotal_bytes{job="node"})"#,
2347                r#"  1
2348-
2349  (
2350      (
2351          node_memory_MemAvailable_bytes{job="node"}
2352        or
2353          (
2354                  node_memory_Buffers_bytes{job="node"}
2355                +
2356                  node_memory_Cached_bytes{job="node"}
2357              +
2358                node_memory_MemFree_bytes{job="node"}
2359            +
2360              node_memory_Slab_bytes{job="node"}
2361          )
2362      )
2363    /
2364      node_memory_MemTotal_bytes{job="node"}
2365  )"#,
2366            ),
2367            (
2368                r#"min by (job, integration) (rate(alertmanager_notifications_failed_total{job="alertmanager", integration=~".*"}[5m]) / rate(alertmanager_notifications_total{job="alertmanager", integration="~.*"}[5m])) > 0.01"#,
2369                r#"  min by (job, integration) (
2370      rate(
2371        alertmanager_notifications_failed_total{integration=~".*",job="alertmanager"}[5m]
2372      )
2373    /
2374      rate(
2375        alertmanager_notifications_total{integration="~.*",job="alertmanager"}[5m]
2376      )
2377  )
2378>
2379  0.01"#,
2380            ),
2381            (
2382                r#"(count by (job) (changes(process_start_time_seconds{job="alertmanager"}[10m]) > 4) / count by (job) (up{job="alertmanager"})) >= 0.5"#,
2383                r#"  (
2384      count by (job) (
2385          changes(
2386            process_start_time_seconds{job="alertmanager"}[10m]
2387          )
2388        >
2389          4
2390      )
2391    /
2392      count by (job) (
2393        up{job="alertmanager"}
2394      )
2395  )
2396>=
2397  0.5"#,
2398            ),
2399        ];
2400
2401        for (input, expect) in cases {
2402            let expr = crate::parser::parse(input);
2403            assert_eq!(expect, expr.unwrap().pretty(0, 10));
2404        }
2405    }
2406
2407    #[test]
2408    fn test_step_invariant_pretty() {
2409        let cases = vec![
2410            ("a @ 1", "a @ 1.000"),
2411            ("a @ start()", "a @ start()"),
2412            ("vector_selector @ start()", "vector_selector @ start()"),
2413        ];
2414
2415        for (input, expect) in cases {
2416            let expr = crate::parser::parse(input);
2417            assert_eq!(expect, expr.unwrap().pretty(0, 10));
2418        }
2419    }
2420
2421    #[test]
2422    fn test_prettify() {
2423        let cases = vec![
2424            ("vector_selector", "vector_selector"),
2425            (
2426                r#"vector_selector{fooooooooooooooooo="barrrrrrrrrrrrrrrrrrr",barrrrrrrrrrrrrrrrrrr="fooooooooooooooooo",process_name="alertmanager"}"#,
2427                r#"vector_selector{barrrrrrrrrrrrrrrrrrr="fooooooooooooooooo",fooooooooooooooooo="barrrrrrrrrrrrrrrrrrr",process_name="alertmanager"}"#,
2428            ),
2429            (
2430                r#"matrix_selector{fooooooooooooooooo="barrrrrrrrrrrrrrrrrrr",barrrrrrrrrrrrrrrrrrr="fooooooooooooooooo",process_name="alertmanager"}[1y2w3d]"#,
2431                r#"matrix_selector{barrrrrrrrrrrrrrrrrrr="fooooooooooooooooo",fooooooooooooooooo="barrrrrrrrrrrrrrrrrrr",process_name="alertmanager"}[382d]"#,
2432            ),
2433        ];
2434
2435        for (input, expect) in cases {
2436            assert_eq!(expect, crate::parser::parse(input).unwrap().prettify());
2437        }
2438    }
2439}