sqlness/interceptor/
replace.rs

1// Copyright 2023 CeresDB Project Authors. Licensed under Apache-2.0.
2
3use crate::error::Result;
4use crate::interceptor::{Interceptor, InterceptorFactory, InterceptorRef};
5use crate::SqlnessError;
6use regex::Regex;
7
8pub const PREFIX: &str = "REPLACE";
9
10/// Replace all matched occurrences in execution result to the given string.
11/// The pattern is treated as a regular expression.
12///
13/// Grammar:
14/// ``` text
15/// -- SQLNESS REPLACE <pattern> <replacement>
16/// ```
17///
18/// `replacement` is optional. If not specified, it will be replaced with an empty string.
19///
20/// # Example
21/// `.sql` file:
22/// ``` sql
23/// -- SQLNESS REPLACE 0 1
24/// SELECT 0;
25/// ```
26///
27/// `.result` file:
28/// ``` sql
29/// -- SQLNESS REPLACE 0 1
30/// SELECT 0;
31///
32/// 1
33/// ```
34///
35/// Multiple `REPLACE` statements are allowed to one query. They will be evaluated in order.
36#[derive(Debug)]
37pub struct ReplaceInterceptor {
38    pattern: String,
39    replacement: String,
40}
41
42impl Interceptor for ReplaceInterceptor {
43    fn after_execute(&self, result: &mut String) {
44        let re = Regex::new(&self.pattern).unwrap();
45        let replaced = re.replace_all(result, &self.replacement);
46        *result = replaced.to_string();
47    }
48}
49
50pub struct ReplaceInterceptorFactory;
51
52impl InterceptorFactory for ReplaceInterceptorFactory {
53    fn try_new(&self, ctx: &str) -> Result<InterceptorRef> {
54        // TODO(ruihang): support pattern with blanks
55        let mut args = ctx.splitn(2, ' ');
56        let pattern = args
57            .next()
58            .ok_or_else(|| SqlnessError::InvalidContext {
59                prefix: PREFIX.to_string(),
60                msg: "Expect <pattern> [replacement]".to_string(),
61            })?
62            .to_string();
63        if pattern.is_empty() {
64            return Err(SqlnessError::InvalidContext {
65                prefix: PREFIX.to_string(),
66                msg: "Pattern shouldn't be empty".to_string(),
67            });
68        }
69        let replacement = args.next().unwrap_or("").to_string();
70        Ok(Box::new(ReplaceInterceptor {
71            pattern,
72            replacement,
73        }))
74    }
75}
76
77#[cfg(test)]
78mod tests {
79    use super::*;
80
81    #[test]
82    fn construct_replace_with_empty_string() {
83        let interceptor = ReplaceInterceptorFactory {}.try_new("");
84        assert!(interceptor.is_err());
85    }
86
87    #[test]
88    fn replace_without_replacement() {
89        let interceptor = ReplaceInterceptorFactory {}.try_new("0").unwrap();
90
91        let mut exec_result = "000010101".to_string();
92        interceptor.after_execute(&mut exec_result);
93        assert_eq!(exec_result, "111".to_string());
94    }
95
96    #[test]
97    fn simple_replace() {
98        let interceptor = ReplaceInterceptorFactory {}.try_new("00 2").unwrap();
99
100        let mut exec_result = "0000010101".to_string();
101        interceptor.after_execute(&mut exec_result);
102        assert_eq!(exec_result, "22010101".to_string());
103    }
104}
sqlness/interceptor/replace.rs

sqlness/interceptor/
replace.rs
