refactor: InterceptorFactory return Result, not Option

sqlness/examples/interceptor-case/simple/input.result

@@ -5,7 +5,6 @@ Args: addr=127.0.0.1, bla=a, port=3306
 
 SELECT * FROM t;
 
--- SQLNESS 🤪ARG🤪 addr=127.0.0.1 port=3306
 -- SQLNESS   ARG port=3307
 SELECT *
 FROM t
@@ -110,12 +109,3 @@ INSERT INTO t (c) VALUES(1)  , (2)  , (3)  , (4) ;
 4
 2;
 
--- SQLNESS SORT_RESULT -1
-6
-2
-4;
-
-6
-2
-4;
-

sqlness/examples/interceptor-case/simple/input.sql

@@ -1,7 +1,6 @@
 -- SQLNESS ARG bla=a addr=127.0.0.1 port=3306
 SELECT * FROM t;
 
--- SQLNESS 🤪ARG🤪 addr=127.0.0.1 port=3306
 -- SQLNESS   ARG port=3307
 SELECT *
 FROM t
@@ -64,8 +63,3 @@ INSERT INTO t (c) VALUES
 4
 2
 2;
-
--- SQLNESS SORT_RESULT -1
-6
-2
-4;

sqlness/src/case.rs

@@ -11,7 +11,7 @@ use std::{
 use crate::{
     config::Config,
     error::Result,
-    interceptor::{InterceptorFactoryRef, InterceptorRef},
+    interceptor::{InterceptorRef, Registry},
     Database, SqlnessError,
 };
 
@@ -31,7 +31,7 @@ impl TestCase {
         })?;
 
         let mut queries = vec![];
-        let mut query = Query::with_interceptor_factories(cfg.interceptor_factories.clone());
+        let mut query = Query::with_interceptor_factories(cfg.interceptor_registry.clone());
 
         let reader = BufReader::new(file);
         for line in reader.lines() {
@@ -43,7 +43,7 @@ impl TestCase {
 
                 // intercept command start with INTERCEPTOR_PREFIX
                 if line.starts_with(&cfg.interceptor_prefix) {
-                    query.push_interceptor(&cfg.interceptor_prefix, line);
+                    query.push_interceptor(&cfg.interceptor_prefix, line)?;
                 }
                 continue;
             }
@@ -58,7 +58,7 @@ impl TestCase {
             // SQL statement ends with ';'
             if line.ends_with(QUERY_DELIMITER) {
                 queries.push(query);
-                query = Query::with_interceptor_factories(cfg.interceptor_factories.clone());
+                query = Query::with_interceptor_factories(cfg.interceptor_registry.clone());
             } else {
                 query.append_query_line("\n");
             }
@@ -101,26 +101,31 @@ struct Query {
     display_query: Vec<String>,
     /// Query to be executed
     execute_query: Vec<String>,
-    interceptor_factories: Vec<InterceptorFactoryRef>,
+    interceptor_registry: Registry,
     interceptors: Vec<InterceptorRef>,
 }
 
 impl Query {
-    pub fn with_interceptor_factories(interceptor_factories: Vec<InterceptorFactoryRef>) -> Self {
+    pub fn with_interceptor_factories(interceptor_registry: Registry) -> Self {
         Self {
-            interceptor_factories,
+            interceptor_registry,
             ..Default::default()
         }
     }
 
-    fn push_interceptor(&mut self, interceptor_prefix: &str, interceptor_line: String) {
-        let interceptor_text = interceptor_line
-            .trim_start_matches(interceptor_prefix)
-            .trim_start();
-        for factories in &self.interceptor_factories {
-            if let Some(interceptor) = factories.try_new(interceptor_text) {
-                self.interceptors.push(interceptor);
-            }
+    fn push_interceptor(
+        &mut self,
+        interceptor_prefix: &str,
+        interceptor_line: String,
+    ) -> Result<()> {
+        if let Some((_, remaining)) = interceptor_line.split_once(interceptor_prefix) {
+            let interceptor = self.interceptor_registry.create(remaining)?;
+            self.interceptors.push(interceptor);
+            Ok(())
+        } else {
+            Err(SqlnessError::MissingPrefix {
+                line: interceptor_line,
+            })
         }
     }
 

sqlness/src/config.rs

@@ -1,9 +1,8 @@
 // Copyright 2022 CeresDB Project Authors. Licensed under Apache-2.0.
 
+use crate::interceptor::Registry;
 use derive_builder::Builder;
 
-use crate::interceptor::{builtin_interceptors, InterceptorFactoryRef};
-
 /// Configurations of [`Runner`].
 ///
 /// [`Runner`]: crate::Runner
@@ -36,10 +35,9 @@ pub struct Config {
     /// Defaults to "true" (follow symbolic links).
     #[builder(default = "Config::default_follow_links()")]
     pub follow_links: bool,
-
     /// Interceptors used to pre-process input query and post-process query response
-    #[builder(default = "Config::default_interceptors()")]
-    pub interceptor_factories: Vec<InterceptorFactoryRef>,
+    #[builder(default = "Config::default_registry()")]
+    pub interceptor_registry: Registry,
 }
 
 impl Config {
@@ -75,8 +73,8 @@ impl Config {
         true
     }
 
-    fn default_interceptors() -> Vec<InterceptorFactoryRef> {
-        builtin_interceptors()
+    fn default_registry() -> Registry {
+        Registry::default()
     }
 }
 

sqlness/src/error.rs

@@ -29,6 +29,15 @@ pub enum SqlnessError {
 
     #[error("Invalid regexp, source error: {0}")]
     Regex(#[from] regex::Error),
+
+    #[error("Unknown interceptor prefix, value:{prefix}.")]
+    UnknownInterceptor { prefix: String },
+
+    #[error("Invalid interceptor context, prefix:{prefix}, msg:{msg}.")]
+    InvalidContext { prefix: String, msg: String },
+
+    #[error("Missing interceptor prefix, line:{line}.")]
+    MissingPrefix { line: String },
 }
 
 pub(crate) type Result<T> = std::result::Result<T, SqlnessError>;

sqlness/src/interceptor.rs

@@ -2,10 +2,12 @@
 
 //! Query interceptor implementations.
 
-use std::sync::Arc;
+use std::{collections::HashMap, sync::Arc};
 
 use crate::{
     case::QueryContext,
+    error::Result,
+    error::SqlnessError,
     interceptor::{
         arg::ArgInterceptorFactory, env::EnvInterceptorFactory, replace::ReplaceInterceptorFactory,
         sort_result::SortResultInterceptorFactory, template::TemplateInterceptorFactory,
@@ -31,16 +33,68 @@ pub trait Interceptor {
 pub type InterceptorFactoryRef = Arc<dyn InterceptorFactory>;
 
 pub trait InterceptorFactory {
-    fn try_new(&self, interceptor: &str) -> Option<InterceptorRef>;
+    fn try_new(&self, ctx: &str) -> Result<InterceptorRef>;
+}
+
+#[derive(Clone)]
+pub struct Registry {
+    factories: HashMap<String, InterceptorFactoryRef>,
+}
+
+impl Default for Registry {
+    fn default() -> Self {
+        Self {
+            factories: builtin_interceptors(),
+        }
+    }
+}
+
+impl Registry {
+    pub fn register(&mut self, prefix: &str, factory: InterceptorFactoryRef) {
+        self.factories.insert(prefix.to_string(), factory);
+    }
+
+    pub fn create(&self, ctx: &str) -> Result<InterceptorRef> {
+        let mut args = ctx.trim().splitn(2, ' ');
+        let prefix = args.next().ok_or_else(|| SqlnessError::MissingPrefix {
+            line: ctx.to_string(),
+        })?;
+        let context = args.next().unwrap_or_default();
+        if let Some(factory) = self.factories.get(prefix.trim()) {
+            factory.try_new(context.trim())
+        } else {
+            Err(SqlnessError::UnknownInterceptor {
+                prefix: prefix.to_string(),
+            })
+        }
+    }
 }
 
 /// Interceptors builtin sqlness
-pub fn builtin_interceptors() -> Vec<InterceptorFactoryRef> {
-    vec![
-        Arc::new(ArgInterceptorFactory {}),
-        Arc::new(ReplaceInterceptorFactory {}),
-        Arc::new(EnvInterceptorFactory {}),
-        Arc::new(SortResultInterceptorFactory {}),
-        Arc::new(TemplateInterceptorFactory {}),
+fn builtin_interceptors() -> HashMap<String, InterceptorFactoryRef> {
+    [
+        (
+            arg::PREFIX.to_string(),
+            Arc::new(ArgInterceptorFactory {}) as _,
+        ),
+        (
+            replace::PREFIX.to_string(),
+            Arc::new(ReplaceInterceptorFactory {}) as _,
+        ),
+        (
+            env::PREFIX.to_string(),
+            Arc::new(EnvInterceptorFactory {}) as _,
+        ),
+        (
+            sort_result::PREFIX.to_string(),
+            Arc::new(SortResultInterceptorFactory {}) as _,
+        ),
+        (
+            template::PREFIX.to_string(),
+            Arc::new(TemplateInterceptorFactory {}) as _,
+        ),
     ]
+    .into_iter()
+    .map(|(prefix, factory)| (prefix.to_string(), factory))
+    .collect()
 }

sqlness/src/interceptor/arg.rs

@@ -1,9 +1,10 @@
 // Copyright 2022 CeresDB Project Authors. Licensed under Apache-2.0.
 
 use crate::case::QueryContext;
+use crate::error::Result;
 use crate::interceptor::{Interceptor, InterceptorFactory, InterceptorRef};
 
-const PREFIX: &str = "ARG";
+pub const PREFIX: &str = "ARG";
 
 /// Pass arguments to the [QueryContext].
 ///
@@ -35,14 +36,9 @@ impl Interceptor for ArgInterceptor {
 pub struct ArgInterceptorFactory;
 
 impl InterceptorFactory for ArgInterceptorFactory {
-    fn try_new(&self, interceptor: &str) -> Option<InterceptorRef> {
-        if interceptor.starts_with(PREFIX) {
-            let args =
-                Self::separate_key_value_pairs(interceptor.trim_start_matches(PREFIX).trim_start());
-            Some(Box::new(ArgInterceptor { args }))
-        } else {
-            None
-        }
+    fn try_new(&self, ctx: &str) -> Result<InterceptorRef> {
+        let args = Self::separate_key_value_pairs(ctx);
+        Ok(Box::new(ArgInterceptor { args }))
     }
 }
 
@@ -64,7 +60,7 @@ mod test {
 
     #[test]
     fn cut_arg_string() {
-        let input = "ARG arg1=value1 arg2=value2 arg3=a=b=c arg4= arg5=,,,";
+        let input = "arg1=value1 arg2=value2 arg3=a=b=c arg4= arg5=,,,";
         let expected = vec![
             ("arg1".to_string(), "value1".to_string()),
             ("arg2".to_string(), "value2".to_string()),

sqlness/src/interceptor/env.rs

@@ -3,9 +3,10 @@
 use std::collections::HashMap;
 
 use crate::case::QueryContext;
+use crate::error::Result;
 use crate::interceptor::{Interceptor, InterceptorFactory, InterceptorRef};
 
-const PREFIX: &str = "ENV";
+pub const PREFIX: &str = "ENV";
 
 /// Read environment variables and fill them in query.
 ///
@@ -52,31 +53,24 @@ impl Interceptor for EnvInterceptor {
 pub struct EnvInterceptorFactory;
 
 impl InterceptorFactory for EnvInterceptorFactory {
-    fn try_new(&self, interceptor: &str) -> Option<InterceptorRef> {
-        Self::create(interceptor).map(|i| Box::new(i) as InterceptorRef)
+    fn try_new(&self, ctx: &str) -> Result<InterceptorRef> {
+        Self::create(ctx).map(|v| Box::new(v) as _)
     }
 }
 
 impl EnvInterceptorFactory {
-    fn create(interceptor: &str) -> Option<EnvInterceptor> {
-        if interceptor.starts_with(PREFIX) {
-            let input = interceptor
-                .trim_start_matches(PREFIX)
-                .trim_start()
-                .trim_end();
-            let envs = input.split(' ').collect::<Vec<_>>();
+    fn create(s: &str) -> Result<EnvInterceptor> {
+        let input = s.trim_start().trim_end();
+        let envs = input.split(' ').collect::<Vec<_>>();
 
-            let mut env_data = HashMap::new();
-            for env in envs {
-                if let Ok(value) = std::env::var(env) {
-                    env_data.insert(format!("${env}"), value);
-                }
+        let mut data = HashMap::new();
+        for env in envs {
+            if let Ok(value) = std::env::var(env) {
+                data.insert(format!("${env}"), value);
             }
-
-            Some(EnvInterceptor { data: env_data })
-        } else {
-            None
         }
+
+        Ok(EnvInterceptor { data })
     }
 }
 
@@ -86,7 +80,7 @@ mod test {
 
     #[test]
     fn cut_env_string() {
-        let input = "ENV SECRET NONEXISTENT";
+        let input = "SECRET NONEXISTENT";
         std::env::set_var("SECRET", "2333");
 
         let expected = [("$SECRET".to_string(), "2333".to_string())]