prom_mock_rs/
query_engine.rs

1//! Simple query engine for basic metric selectors without full `PromQL`.
2//!
3//! This module provides a basic query parser and executor that supports
4//! simple metric selectors like `metric{label="value"}` without requiring
5//! a full `PromQL` implementation.
6
7use std::io;
8use std::sync::Arc;
9
10use regex::Regex;
11
12use crate::matchers::{EqualMatcher, LabelMatcher, NotEqualMatcher, NotRegexMatcher, RegexMatcher};
13use crate::storage::Storage;
14
15/// Simple query parser for basic selectors like: metric{a="b",c!="d",e=~"regex"}.
16#[derive(Clone)]
17pub struct SimpleQueryEngine {
18    storage: Arc<dyn Storage>,
19}
20
21impl SimpleQueryEngine {
22    /// Create a new query engine with the given storage backend.
23    ///
24    /// # Parameters
25    ///
26    /// - `storage` - Shared reference to any storage implementation
27    ///
28    /// # Returns
29    ///
30    /// Returns a new `SimpleQueryEngine` instance.
31    pub fn new(storage: Arc<dyn Storage>) -> Self {
32        Self { storage }
33    }
34
35    /// Parse and execute a simple metric selector query
36    pub fn query(&self, query: &str, start: i64, end: i64) -> io::Result<QueryResult> {
37        let selector = Self::parse_selector(query)?;
38        let series = self.storage.query_series(&selector.matchers);
39
40        let mut result_series = Vec::new();
41        for ts in series {
42            let samples = ts.samples_in_range(start, end);
43            if !samples.is_empty() {
44                result_series.push(QueryResultSeries {
45                    labels: ts.labels.clone(),
46                    samples: samples.into_iter().cloned().collect(),
47                });
48            }
49        }
50
51        Ok(QueryResult { series: result_series })
52    }
53
54    /// Parse a simple selector like: metric{a="b",c!="d",e=~"regex"}
55    fn parse_selector(query: &str) -> io::Result<MetricSelector> {
56        let query = query.trim();
57
58        // Split metric name and labels
59        let (metric_name, labels_part) = query.find('{').map_or((Some(query), None), |brace_pos| {
60            let name = query[..brace_pos].trim();
61            let labels = &query[brace_pos..];
62            (Some(name), Some(labels))
63        });
64
65        let mut matchers: Vec<Arc<dyn LabelMatcher>> = Vec::new();
66
67        // Add metric name matcher if present
68        if let Some(name) = metric_name {
69            if !name.is_empty() {
70                matchers.push(Arc::new(EqualMatcher::new("__name__", name)));
71            }
72        }
73
74        // Parse label matchers if present
75        if let Some(labels) = labels_part {
76            let label_matchers = Self::parse_label_matchers(labels)?;
77            matchers.extend(label_matchers);
78        }
79
80        Ok(MetricSelector { matchers })
81    }
82
83    /// Parse label matchers from string like: {a="b",c!="d",e=~"regex"}
84    fn parse_label_matchers(labels_str: &str) -> io::Result<Vec<Arc<dyn LabelMatcher>>> {
85        let labels_str = labels_str.trim();
86        if !labels_str.starts_with('{') || !labels_str.ends_with('}') {
87            return Err(io::Error::new(io::ErrorKind::InvalidInput, "invalid label syntax"));
88        }
89
90        let inner = &labels_str[1..labels_str.len() - 1];
91        if inner.trim().is_empty() {
92            return Ok(Vec::new());
93        }
94
95        let mut matchers = Vec::new();
96        let parts = Self::split_label_expressions(inner);
97
98        for part in parts {
99            let matcher = Self::parse_single_label_matcher(&part)?;
100            matchers.push(matcher);
101        }
102
103        Ok(matchers)
104    }
105
106    /// Split label expressions by comma, handling quoted strings
107    fn split_label_expressions(input: &str) -> Vec<String> {
108        let mut parts = Vec::new();
109        let mut current = String::new();
110        let mut in_quotes = false;
111        let mut escape_next = false;
112
113        for ch in input.chars() {
114            if escape_next {
115                current.push(ch);
116                escape_next = false;
117                continue;
118            }
119
120            match ch {
121                '\\' => {
122                    escape_next = true;
123                    current.push(ch);
124                }
125                '"' => {
126                    in_quotes = !in_quotes;
127                    current.push(ch);
128                }
129                ',' if !in_quotes => {
130                    parts.push(current.trim().to_string());
131                    current.clear();
132                }
133                _ => {
134                    current.push(ch);
135                }
136            }
137        }
138
139        if !current.trim().is_empty() {
140            parts.push(current.trim().to_string());
141        }
142
143        parts
144    }
145
146    /// Parse a single label matcher like: a="b" or c!="d" or e=~"regex"
147    fn parse_single_label_matcher(expr: &str) -> io::Result<Arc<dyn LabelMatcher>> {
148        let expr = expr.trim();
149
150        // Find operator
151        if let Some(pos) = expr.find("!=") {
152            let name = expr[..pos].trim().to_string();
153            let value = Self::parse_quoted_value(&expr[pos + 2..])?;
154            return Ok(Arc::new(NotEqualMatcher::new(name, value)));
155        }
156
157        if let Some(pos) = expr.find("!~") {
158            let name = expr[..pos].trim().to_string();
159            let pattern_str = Self::parse_quoted_value(&expr[pos + 2..])?;
160            let pattern = Regex::new(&pattern_str).map_err(|e| {
161                io::Error::new(io::ErrorKind::InvalidInput, format!("invalid regex: {e}"))
162            })?;
163            return Ok(Arc::new(NotRegexMatcher::new(name, pattern)));
164        }
165
166        if let Some(pos) = expr.find("=~") {
167            let name = expr[..pos].trim().to_string();
168            let pattern_str = Self::parse_quoted_value(&expr[pos + 2..])?;
169            let pattern = Regex::new(&pattern_str).map_err(|e| {
170                io::Error::new(io::ErrorKind::InvalidInput, format!("invalid regex: {e}"))
171            })?;
172            return Ok(Arc::new(RegexMatcher::new(name, pattern)));
173        }
174
175        if let Some(pos) = expr.find('=') {
176            let name = expr[..pos].trim().to_string();
177            let value = Self::parse_quoted_value(&expr[pos + 1..])?;
178            return Ok(Arc::new(EqualMatcher::new(name, value)));
179        }
180
181        Err(io::Error::new(io::ErrorKind::InvalidInput, format!("invalid label matcher: {expr}")))
182    }
183
184    /// Parse quoted value, removing quotes
185    fn parse_quoted_value(input: &str) -> io::Result<String> {
186        let input = input.trim();
187        if input.starts_with('"') && input.ends_with('"') && input.len() >= 2 {
188            Ok(input[1..input.len() - 1].to_string())
189        } else {
190            Err(io::Error::new(
191                io::ErrorKind::InvalidInput,
192                format!("expected quoted value: {input}"),
193            ))
194        }
195    }
196}
197
198#[derive(Debug)]
199struct MetricSelector {
200    matchers: Vec<Arc<dyn LabelMatcher>>,
201}
202
203/// Query result containing time series.
204#[derive(Debug)]
205pub struct QueryResult {
206    pub series: Vec<QueryResultSeries>,
207}
208
209#[derive(Debug)]
210pub struct QueryResultSeries {
211    pub labels: Vec<crate::storage::Label>,
212    pub samples: Vec<crate::storage::Sample>,
213}
214
215#[cfg(test)]
216mod tests {
217    use super::*;
218    use crate::storage::{Label, MemoryStorage, Sample, TimeSeries};
219
220    /// Test parsing of simple metric selectors with and without labels.
221    #[test]
222    fn test_parse_simple_selector() {
223        let _engine = SimpleQueryEngine::new(Arc::new(MemoryStorage::new()));
224
225        // Test simple metric name
226        let selector = SimpleQueryEngine::parse_selector("up").expect("valid syntax");
227        assert_eq!(selector.matchers.len(), 1);
228
229        // Test with labels
230        let selector =
231            SimpleQueryEngine::parse_selector(r#"http_requests{job="api",method!="POST"}"#)
232                .expect("valid syntax");
233        assert_eq!(selector.matchers.len(), 3); // __name__, job, method
234    }
235
236    /// Test splitting of label expressions with proper quote handling.
237    #[test]
238    fn test_split_label_expressions() {
239        let parts = SimpleQueryEngine::split_label_expressions(r#"a="b",c!="d",e=~"regex""#);
240        assert_eq!(parts.len(), 3);
241        assert_eq!(parts[0], r#"a="b""#);
242        assert_eq!(parts[1], r#"c!="d""#);
243        assert_eq!(parts[2], r#"e=~"regex""#);
244    }
245
246    /// Test parsing of individual label matchers for equality and inequality operations.
247    #[test]
248    fn test_parse_label_matcher() {
249        let matcher =
250            SimpleQueryEngine::parse_single_label_matcher(r#"job="api""#).expect("valid syntax");
251        assert_eq!(matcher.label_name(), "job");
252
253        let labels = vec![crate::storage::Label::new("job", "api")];
254        assert!(matcher.matches(&labels));
255
256        let matcher = SimpleQueryEngine::parse_single_label_matcher(r#"method!="POST""#)
257            .expect("valid syntax");
258        assert_eq!(matcher.label_name(), "method");
259
260        let labels = vec![crate::storage::Label::new("method", "GET")];
261        assert!(matcher.matches(&labels));
262    }
263
264    /// Test end-to-end query functionality with in-memory storage.
265    #[test]
266    fn test_query_with_storage() {
267        let storage = Arc::new(MemoryStorage::new());
268        let engine = SimpleQueryEngine::new(storage.clone());
269
270        // Add test data
271        let labels = vec![
272            Label::new("__name__", "http_requests"),
273            Label::new("job", "api"),
274            Label::new("method", "GET"),
275        ];
276        let mut ts = TimeSeries::new(labels);
277        ts.add_sample(Sample::new(1000, 10.0));
278        ts.add_sample(Sample::new(2000, 20.0));
279        storage.add_series(ts);
280
281        // Query
282        let result = engine.query(r#"http_requests{job="api"}"#, 0, 3000).expect("valid query");
283        assert_eq!(result.series.len(), 1);
284        assert_eq!(result.series[0].samples.len(), 2);
285    }
286
287    /// Test parsing edge cases and error conditions.
288    #[test]
289    fn test_parse_edge_cases() {
290        // Empty selector (no metric name, no labels)
291        let result = SimpleQueryEngine::parse_selector("{}");
292        assert!(result.is_ok());
293        let selector = result.unwrap();
294        assert_eq!(selector.matchers.len(), 0);
295
296        // Only metric name, no labels
297        let result = SimpleQueryEngine::parse_selector("cpu_usage");
298        assert!(result.is_ok());
299        let selector = result.unwrap();
300        assert_eq!(selector.matchers.len(), 1);
301        assert_eq!(selector.matchers[0].label_name(), "__name__");
302
303        // Empty metric name with labels
304        let result = SimpleQueryEngine::parse_selector(r#"{job="api"}"#);
305        assert!(result.is_ok());
306        let selector = result.unwrap();
307        assert_eq!(selector.matchers.len(), 1);
308        assert_eq!(selector.matchers[0].label_name(), "job");
309
310        // Whitespace handling
311        let result = SimpleQueryEngine::parse_selector(r#"  cpu_usage  { job = "api" }  "#);
312        assert!(result.is_ok());
313        let selector = result.unwrap();
314        assert_eq!(selector.matchers.len(), 2);
315    }
316
317    /// Test invalid label syntax error cases.
318    #[test]
319    fn test_parse_error_cases() {
320        // Invalid label syntax - missing closing brace
321        let result = SimpleQueryEngine::parse_selector(r#"metric{job="api""#);
322        assert!(result.is_err());
323
324        // Invalid label syntax - missing opening brace
325        let result = SimpleQueryEngine::parse_label_matchers(r#"job="api"}"#);
326        assert!(result.is_err());
327
328        // Invalid quoted value
329        let result = SimpleQueryEngine::parse_quoted_value("unquoted");
330        assert!(result.is_err());
331
332        // Invalid quoted value - missing closing quote
333        let result = SimpleQueryEngine::parse_quoted_value(r#""missing_close"#);
334        assert!(result.is_err());
335
336        // Invalid label matcher - no operator
337        let result = SimpleQueryEngine::parse_single_label_matcher("invalid_syntax");
338        assert!(result.is_err());
339
340        // Invalid regex pattern
341        let result = SimpleQueryEngine::parse_single_label_matcher(r#"test=~"[invalid""#);
342        assert!(result.is_err());
343    }
344
345    /// Test regex matchers (=~ and !~).
346    #[test]
347    fn test_regex_matchers() {
348        // Valid regex matcher
349        let result = SimpleQueryEngine::parse_single_label_matcher(r#"instance=~"server.*""#);
350        assert!(result.is_ok());
351        let matcher = result.unwrap();
352        assert_eq!(matcher.label_name(), "instance");
353
354        let labels = vec![crate::storage::Label::new("instance", "server1")];
355        assert!(matcher.matches(&labels));
356
357        let labels = vec![crate::storage::Label::new("instance", "client1")];
358        assert!(!matcher.matches(&labels));
359
360        // Valid not-regex matcher
361        let result = SimpleQueryEngine::parse_single_label_matcher(r#"method!~"POST|PUT""#);
362        assert!(result.is_ok());
363        let matcher = result.unwrap();
364        assert_eq!(matcher.label_name(), "method");
365
366        let labels = vec![crate::storage::Label::new("method", "GET")];
367        assert!(matcher.matches(&labels));
368
369        let labels = vec![crate::storage::Label::new("method", "POST")];
370        assert!(!matcher.matches(&labels));
371    }
372
373    /// Test complex label expression splitting with quotes and escapes.
374    #[test]
375    fn test_complex_label_splitting() {
376        // Test with escaped quotes
377        let parts =
378            SimpleQueryEngine::split_label_expressions(r#"a="value with \"quotes\"",b="normal""#);
379        assert_eq!(parts.len(), 2);
380        assert_eq!(parts[0], r#"a="value with \"quotes\"""#);
381        assert_eq!(parts[1], r#"b="normal""#);
382
383        // Test with commas inside quotes
384        let parts = SimpleQueryEngine::split_label_expressions(
385            r#"description="contains, comma",job="api""#,
386        );
387        assert_eq!(parts.len(), 2);
388        assert_eq!(parts[0], r#"description="contains, comma""#);
389        assert_eq!(parts[1], r#"job="api""#);
390
391        // Test with empty parts
392        let parts = SimpleQueryEngine::split_label_expressions("");
393        assert_eq!(parts.len(), 0);
394
395        // Test single expression
396        let parts = SimpleQueryEngine::split_label_expressions(r#"job="api""#);
397        assert_eq!(parts.len(), 1);
398        assert_eq!(parts[0], r#"job="api""#);
399    }
400
401    /// Test query with time filtering.
402    #[test]
403    fn test_query_with_time_filtering() {
404        let storage = Arc::new(MemoryStorage::new());
405        let engine = SimpleQueryEngine::new(storage.clone());
406
407        // Add test data with different timestamps
408        let labels = vec![Label::new("__name__", "cpu_usage"), Label::new("job", "api")];
409        let mut ts = TimeSeries::new(labels);
410        ts.add_sample(Sample::new(1000, 10.0));
411        ts.add_sample(Sample::new(2000, 20.0));
412        ts.add_sample(Sample::new(3000, 30.0));
413        ts.add_sample(Sample::new(4000, 40.0));
414        storage.add_series(ts);
415
416        // Query with time range that excludes some samples
417        let result = engine.query("cpu_usage", 1500, 3500).expect("valid query");
418        assert_eq!(result.series.len(), 1);
419        assert_eq!(result.series[0].samples.len(), 2); // Only samples at 2000 and 3000
420        assert_eq!(result.series[0].samples[0].timestamp, 2000);
421        assert_eq!(result.series[0].samples[1].timestamp, 3000);
422
423        // Query with no matches in time range
424        let result = engine.query("cpu_usage", 5000, 6000).expect("valid query");
425        assert_eq!(result.series.len(), 0); // No samples in range, so no series
426    }
427
428    /// Test query with complex selector and multiple series.
429    #[test]
430    fn test_complex_query() {
431        let storage = Arc::new(MemoryStorage::new());
432        let engine = SimpleQueryEngine::new(storage.clone());
433
434        // Add multiple series
435        storage.add_series({
436            let mut ts = TimeSeries::new(vec![
437                Label::new("__name__", "http_requests"),
438                Label::new("job", "api"),
439                Label::new("method", "GET"),
440            ]);
441            ts.add_sample(Sample::new(1000, 10.0));
442            ts
443        });
444
445        storage.add_series({
446            let mut ts = TimeSeries::new(vec![
447                Label::new("__name__", "http_requests"),
448                Label::new("job", "api"),
449                Label::new("method", "POST"),
450            ]);
451            ts.add_sample(Sample::new(1000, 5.0));
452            ts
453        });
454
455        storage.add_series({
456            let mut ts = TimeSeries::new(vec![
457                Label::new("__name__", "http_requests"),
458                Label::new("job", "web"),
459                Label::new("method", "GET"),
460            ]);
461            ts.add_sample(Sample::new(1000, 15.0));
462            ts
463        });
464
465        // Query with multiple matchers
466        let result = engine
467            .query(r#"http_requests{job="api",method!="POST"}"#, 0, 2000)
468            .expect("valid query");
469        assert_eq!(result.series.len(), 1); // Only the API GET requests
470        assert_eq!(result.series[0].samples[0].value, 10.0);
471
472        // Query with regex matcher
473        let result =
474            engine.query(r#"http_requests{job=~".*api.*"}"#, 0, 2000).expect("valid query");
475        assert_eq!(result.series.len(), 2); // Both API series (GET and POST)
476    }
477}
prom_mock_rs/query_engine.rs

prom_mock_rs/
query_engine.rs
