sqlness/interceptor/
arg.rs

1// Copyright 2022 CeresDB Project Authors. Licensed under Apache-2.0.
2
3use crate::case::QueryContext;
4use crate::error::Result;
5use crate::interceptor::{Interceptor, InterceptorFactory, InterceptorRef};
6
7pub const PREFIX: &str = "ARG";
8
9/// Pass arguments to the [QueryContext].
10///
11/// # Example
12/// ``` sql
13/// -- SQLNESS ARG arg1=value1 arg2=value2
14/// SELECT * FROM t;
15/// ```
16///
17/// # Format
18/// The arguments are in the format of `key=value`, without space. The key should not contains
19/// equal marker (`=`), and value can be any string without space. The arguments are separated
20/// by spaces.
21///
22/// It will overwrite existing key-value pair in context if the key name is same.
23#[derive(Debug)]
24pub struct ArgInterceptor {
25    args: Vec<(String, String)>,
26}
27
28impl Interceptor for ArgInterceptor {
29    fn before_execute(&self, _: &mut Vec<String>, context: &mut QueryContext) {
30        for (key, value) in &self.args {
31            context.context.insert(key.to_string(), value.to_string());
32        }
33    }
34}
35
36pub struct ArgInterceptorFactory;
37
38impl InterceptorFactory for ArgInterceptorFactory {
39    fn try_new(&self, ctx: &str) -> Result<InterceptorRef> {
40        let args = Self::separate_key_value_pairs(ctx);
41        Ok(Box::new(ArgInterceptor { args }))
42    }
43}
44
45impl ArgInterceptorFactory {
46    fn separate_key_value_pairs(input: &str) -> Vec<(String, String)> {
47        let mut result = Vec::new();
48        for pair in input.split(' ') {
49            if let Some((key, value)) = pair.split_once('=') {
50                result.push((key.to_string(), value.to_string()));
51            }
52        }
53        result
54    }
55}
56
57#[cfg(test)]
58mod test {
59    use super::*;
60
61    #[test]
62    fn cut_arg_string() {
63        let input = "arg1=value1 arg2=value2 arg3=a=b=c arg4= arg5=,,,";
64        let expected = vec![
65            ("arg1".to_string(), "value1".to_string()),
66            ("arg2".to_string(), "value2".to_string()),
67            ("arg3".to_string(), "a=b=c".to_string()),
68            ("arg4".to_string(), "".to_string()),
69            ("arg5".to_string(), ",,,".to_string()),
70        ];
71
72        let args = ArgInterceptorFactory::separate_key_value_pairs(input);
73        assert_eq!(args, expected);
74    }
75}
sqlness/interceptor/arg.rs

sqlness/interceptor/
arg.rs
