From e64c65caf522263ce0668b84b14ad9d6d0575183 Mon Sep 17 00:00:00 2001 From: jiacai2050 Date: Fri, 22 Mar 2024 17:51:30 +0800 Subject: [PATCH 1/2] refactor factory --- sqlness/src/config.rs | 7 ++++--- sqlness/src/interceptor.rs | 23 +++++++++++++++-------- sqlness/src/interceptor/arg.rs | 2 +- 3 files changed, 20 insertions(+), 12 deletions(-) diff --git a/sqlness/src/config.rs b/sqlness/src/config.rs index a8538fa..8da9a8f 100644 --- a/sqlness/src/config.rs +++ b/sqlness/src/config.rs @@ -1,5 +1,7 @@ // Copyright 2022 CeresDB Project Authors. Licensed under Apache-2.0. +use std::collections::HashMap; + use derive_builder::Builder; use crate::interceptor::{builtin_interceptors, InterceptorFactoryRef}; @@ -36,10 +38,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, + pub interceptor_factories: HashMap, } impl Config { @@ -75,7 +76,7 @@ impl Config { true } - fn default_interceptors() -> Vec { + fn default_interceptors() -> HashMap { builtin_interceptors() } } diff --git a/sqlness/src/interceptor.rs b/sqlness/src/interceptor.rs index 2f30c25..70c86cf 100644 --- a/sqlness/src/interceptor.rs +++ b/sqlness/src/interceptor.rs @@ -2,7 +2,7 @@ //! Query interceptor implementations. -use std::sync::Arc; +use std::{collections::HashMap, sync::Arc}; use crate::{ case::QueryContext, @@ -35,12 +35,19 @@ pub trait InterceptorFactory { } /// Interceptors builtin sqlness -pub fn builtin_interceptors() -> Vec { - vec![ - Arc::new(ArgInterceptorFactory {}), - Arc::new(ReplaceInterceptorFactory {}), - Arc::new(EnvInterceptorFactory {}), - Arc::new(SortResultInterceptorFactory {}), - Arc::new(TemplateInterceptorFactory {}), +pub fn builtin_interceptors() -> HashMap { + [ + ( + arg::PREFIX.to_string(), + Arc::new(ArgInterceptorFactory {}) as InterceptorFactoryRef, + ), + // Arc::new(ArgInterceptorFactory {}), + // Arc::new(ReplaceInterceptorFactory {}), + // Arc::new(EnvInterceptorFactory {}), + // Arc::new(SortResultInterceptorFactory {}), + // Arc::new(TemplateInterceptorFactory {}), ] + .into_iter() + .map(|(prefix, factory)| (prefix.to_string(), factory)) + .collect() } diff --git a/sqlness/src/interceptor/arg.rs b/sqlness/src/interceptor/arg.rs index 6f313c7..2f784a6 100644 --- a/sqlness/src/interceptor/arg.rs +++ b/sqlness/src/interceptor/arg.rs @@ -3,7 +3,7 @@ use crate::case::QueryContext; use crate::interceptor::{Interceptor, InterceptorFactory, InterceptorRef}; -const PREFIX: &str = "ARG"; +pub const PREFIX: &str = "ARG"; /// Pass arguments to the [QueryContext]. /// From 4c6f8045cd80f969fead3eb16aa41fdee68a04c7 Mon Sep 17 00:00:00 2001 From: jiacai2050 Date: Sat, 23 Mar 2024 10:32:35 +0800 Subject: [PATCH 2/2] Factory return Result, not Option --- .../interceptor-case/simple/input.result | 10 --- .../interceptor-case/simple/input.sql | 6 -- sqlness/src/case.rs | 35 +++++---- sqlness/src/config.rs | 13 ++-- sqlness/src/error.rs | 9 +++ sqlness/src/interceptor.rs | 63 +++++++++++++-- sqlness/src/interceptor/arg.rs | 14 ++-- sqlness/src/interceptor/env.rs | 34 ++++----- sqlness/src/interceptor/replace.rs | 57 +++++++------- sqlness/src/interceptor/sort_result.rs | 76 +++++++++---------- sqlness/src/interceptor/template.rs | 44 +++++------ sqlness/src/lib.rs | 1 - 12 files changed, 190 insertions(+), 172 deletions(-) diff --git a/sqlness/examples/interceptor-case/simple/input.result b/sqlness/examples/interceptor-case/simple/input.result index e72d634..deca3f9 100644 --- a/sqlness/examples/interceptor-case/simple/input.result +++ b/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; - diff --git a/sqlness/examples/interceptor-case/simple/input.sql b/sqlness/examples/interceptor-case/simple/input.sql index 85bced9..113409a 100644 --- a/sqlness/examples/interceptor-case/simple/input.sql +++ b/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; diff --git a/sqlness/src/case.rs b/sqlness/src/case.rs index 328145f..a890f4e 100644 --- a/sqlness/src/case.rs +++ b/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, /// Query to be executed execute_query: Vec, - interceptor_factories: Vec, + interceptor_registry: Registry, interceptors: Vec, } impl Query { - pub fn with_interceptor_factories(interceptor_factories: Vec) -> 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, + }) } } diff --git a/sqlness/src/config.rs b/sqlness/src/config.rs index 8da9a8f..a891872 100644 --- a/sqlness/src/config.rs +++ b/sqlness/src/config.rs @@ -1,11 +1,8 @@ // Copyright 2022 CeresDB Project Authors. Licensed under Apache-2.0. -use std::collections::HashMap; - +use crate::interceptor::Registry; use derive_builder::Builder; -use crate::interceptor::{builtin_interceptors, InterceptorFactoryRef}; - /// Configurations of [`Runner`]. /// /// [`Runner`]: crate::Runner @@ -39,8 +36,8 @@ pub struct Config { #[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: HashMap, + #[builder(default = "Config::default_registry()")] + pub interceptor_registry: Registry, } impl Config { @@ -76,8 +73,8 @@ impl Config { true } - fn default_interceptors() -> HashMap { - builtin_interceptors() + fn default_registry() -> Registry { + Registry::default() } } diff --git a/sqlness/src/error.rs b/sqlness/src/error.rs index 5c5fb7e..e77a0af 100644 --- a/sqlness/src/error.rs +++ b/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 = std::result::Result; diff --git a/sqlness/src/interceptor.rs b/sqlness/src/interceptor.rs index 70c86cf..b3e8a14 100644 --- a/sqlness/src/interceptor.rs +++ b/sqlness/src/interceptor.rs @@ -6,6 +6,8 @@ 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,21 +33,66 @@ pub trait Interceptor { pub type InterceptorFactoryRef = Arc; pub trait InterceptorFactory { - fn try_new(&self, interceptor: &str) -> Option; + fn try_new(&self, ctx: &str) -> Result; +} + +#[derive(Clone)] +pub struct Registry { + factories: HashMap, +} + +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 { + 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() -> HashMap { +fn builtin_interceptors() -> HashMap { [ ( arg::PREFIX.to_string(), - Arc::new(ArgInterceptorFactory {}) as InterceptorFactoryRef, + 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 _, ), - // Arc::new(ArgInterceptorFactory {}), - // Arc::new(ReplaceInterceptorFactory {}), - // Arc::new(EnvInterceptorFactory {}), - // Arc::new(SortResultInterceptorFactory {}), - // Arc::new(TemplateInterceptorFactory {}), ] .into_iter() .map(|(prefix, factory)| (prefix.to_string(), factory)) diff --git a/sqlness/src/interceptor/arg.rs b/sqlness/src/interceptor/arg.rs index 2f784a6..81cc3c7 100644 --- a/sqlness/src/interceptor/arg.rs +++ b/sqlness/src/interceptor/arg.rs @@ -1,6 +1,7 @@ // Copyright 2022 CeresDB Project Authors. Licensed under Apache-2.0. use crate::case::QueryContext; +use crate::error::Result; use crate::interceptor::{Interceptor, InterceptorFactory, InterceptorRef}; pub const PREFIX: &str = "ARG"; @@ -35,14 +36,9 @@ impl Interceptor for ArgInterceptor { pub struct ArgInterceptorFactory; impl InterceptorFactory for ArgInterceptorFactory { - fn try_new(&self, interceptor: &str) -> Option { - 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 { + 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()), diff --git a/sqlness/src/interceptor/env.rs b/sqlness/src/interceptor/env.rs index e6c7259..2bb4477 100644 --- a/sqlness/src/interceptor/env.rs +++ b/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 { - Self::create(interceptor).map(|i| Box::new(i) as InterceptorRef) + fn try_new(&self, ctx: &str) -> Result { + Self::create(ctx).map(|v| Box::new(v) as _) } } impl EnvInterceptorFactory { - fn create(interceptor: &str) -> Option { - if interceptor.starts_with(PREFIX) { - let input = interceptor - .trim_start_matches(PREFIX) - .trim_start() - .trim_end(); - let envs = input.split(' ').collect::>(); + fn create(s: &str) -> Result { + let input = s.trim_start().trim_end(); + let envs = input.split(' ').collect::>(); - 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())] diff --git a/sqlness/src/interceptor/replace.rs b/sqlness/src/interceptor/replace.rs index c4b8c99..58c8785 100644 --- a/sqlness/src/interceptor/replace.rs +++ b/sqlness/src/interceptor/replace.rs @@ -1,10 +1,11 @@ // Copyright 2023 CeresDB Project Authors. Licensed under Apache-2.0. -use regex::Regex; - +use crate::error::Result; use crate::interceptor::{Interceptor, InterceptorFactory, InterceptorRef}; +use crate::SqlnessError; +use regex::Regex; -const PREFIX: &str = "REPLACE"; +pub const PREFIX: &str = "REPLACE"; /// Replace all matched occurrences in execution result to the given string. /// The pattern is treated as a regular expression. @@ -49,26 +50,27 @@ impl Interceptor for ReplaceInterceptor { pub struct ReplaceInterceptorFactory; impl InterceptorFactory for ReplaceInterceptorFactory { - fn try_new(&self, interceptor: &str) -> Option { - if interceptor.starts_with(PREFIX) { - let args = interceptor - .trim_start_matches(PREFIX) - .trim_start() - .trim_end(); - // TODO(ruihang): support pattern with blanks - let mut args = args.splitn(2, ' '); - let pattern = args.next()?.to_string(); - if pattern.is_empty() { - return None; - } - let replacement = args.next().unwrap_or("").to_string(); - Some(Box::new(ReplaceInterceptor { - pattern, - replacement, - })) - } else { - None + fn try_new(&self, ctx: &str) -> Result { + // TODO(ruihang): support pattern with blanks + let mut args = ctx.splitn(2, ' '); + let pattern = args + .next() + .ok_or_else(|| SqlnessError::InvalidContext { + prefix: PREFIX.to_string(), + msg: "Expect [replacement]".to_string(), + })? + .to_string(); + if pattern.is_empty() { + return Err(SqlnessError::InvalidContext { + prefix: PREFIX.to_string(), + msg: "Pattern shouldn't be empty".to_string(), + }); } + let replacement = args.next().unwrap_or("").to_string(); + Ok(Box::new(ReplaceInterceptor { + pattern, + replacement, + })) } } @@ -78,15 +80,13 @@ mod tests { #[test] fn construct_replace_with_empty_string() { - let input = "REPLACE "; - let interceptor = ReplaceInterceptorFactory {}.try_new(input); - assert!(interceptor.is_none()); + let interceptor = ReplaceInterceptorFactory {}.try_new(""); + assert!(interceptor.is_err()); } #[test] fn replace_without_replacement() { - let input = "REPLACE 0"; - let interceptor = ReplaceInterceptorFactory {}.try_new(input).unwrap(); + let interceptor = ReplaceInterceptorFactory {}.try_new("0").unwrap(); let mut exec_result = "000010101".to_string(); interceptor.after_execute(&mut exec_result); @@ -95,8 +95,7 @@ mod tests { #[test] fn simple_replace() { - let input = "REPLACE 00 2"; - let interceptor = ReplaceInterceptorFactory {}.try_new(input).unwrap(); + let interceptor = ReplaceInterceptorFactory {}.try_new("00 2").unwrap(); let mut exec_result = "0000010101".to_string(); interceptor.after_execute(&mut exec_result); diff --git a/sqlness/src/interceptor/sort_result.rs b/sqlness/src/interceptor/sort_result.rs index 1a0f66d..d3caec4 100644 --- a/sqlness/src/interceptor/sort_result.rs +++ b/sqlness/src/interceptor/sort_result.rs @@ -2,9 +2,13 @@ use std::collections::VecDeque; -use crate::interceptor::{Interceptor, InterceptorFactory, InterceptorRef}; +use crate::{ + error::Result, + interceptor::{Interceptor, InterceptorFactory, InterceptorRef}, + SqlnessError, +}; -const PREFIX: &str = "SORT_RESULT"; +pub const PREFIX: &str = "SORT_RESULT"; /// Sort the query result in lexicographical order. /// @@ -73,29 +77,29 @@ impl Interceptor for SortResultInterceptor { pub struct SortResultInterceptorFactory; impl InterceptorFactory for SortResultInterceptorFactory { - fn try_new(&self, interceptor: &str) -> Option { - Self::try_new_from_str(interceptor).map(|i| Box::new(i) as _) - } -} - -impl SortResultInterceptorFactory { - fn try_new_from_str(interceptor: &str) -> Option { - if interceptor.starts_with(PREFIX) { - let args = interceptor - .trim_start_matches(PREFIX) - .trim_start() - .trim_end(); - let mut args = args.splitn(2, ' ').filter(|s| !s.is_empty()); - let ignore_head = args.next().unwrap_or("0").parse().ok()?; - let ignore_tail = args.next().unwrap_or("0").parse().ok()?; - - Some(SortResultInterceptor { - ignore_head, - ignore_tail, - }) - } else { - None - } + fn try_new(&self, ctx: &str) -> Result { + let mut args = ctx.splitn(2, ' ').filter(|s| !s.is_empty()); + let ignore_head = + args.next() + .unwrap_or("0") + .parse() + .map_err(|e| SqlnessError::InvalidContext { + prefix: PREFIX.to_string(), + msg: format!("Expect number, err:{e}"), + })?; + let ignore_tail = + args.next() + .unwrap_or("0") + .parse() + .map_err(|e| SqlnessError::InvalidContext { + prefix: PREFIX.to_string(), + msg: format!("Expect number, err:{e}"), + })?; + + Ok(Box::new(SortResultInterceptor { + ignore_head, + ignore_tail, + })) } } @@ -103,25 +107,15 @@ impl SortResultInterceptorFactory { mod tests { use super::*; - #[test] - fn construct_with_empty_string() { - let input = "SORT_RESULT"; - let sort_result = SortResultInterceptorFactory::try_new_from_str(input).unwrap(); - assert_eq!(sort_result.ignore_head, 0); - assert_eq!(sort_result.ignore_tail, 0); - } - #[test] fn construct_with_negative() { - let input = "SORT_RESULT -1"; - let interceptor = SortResultInterceptorFactory.try_new(input); - assert!(interceptor.is_none()); + let interceptor = SortResultInterceptorFactory.try_new("-1"); + assert!(interceptor.is_err()); } #[test] fn sort_result_full() { - let input = "SORT_RESULT"; - let interceptor = SortResultInterceptorFactory.try_new(input).unwrap(); + let interceptor = SortResultInterceptorFactory.try_new("").unwrap(); let cases = [ ( @@ -158,8 +152,7 @@ mod tests { #[test] fn ignore_head_exceeds_length() { - let input = "SORT_RESULT 10000"; - let interceptor = SortResultInterceptorFactory.try_new(input).unwrap(); + let interceptor = SortResultInterceptorFactory.try_new("10000").unwrap(); let mut exec_result = String::from( "3\ @@ -173,8 +166,7 @@ mod tests { #[test] fn ignore_tail_exceeds_length() { - let input = "SORT_RESULT 0 10000"; - let interceptor = SortResultInterceptorFactory.try_new(input).unwrap(); + let interceptor = SortResultInterceptorFactory.try_new("0 10000").unwrap(); let mut exec_result = String::from( "3\ diff --git a/sqlness/src/interceptor/template.rs b/sqlness/src/interceptor/template.rs index c66b00e..07903b6 100644 --- a/sqlness/src/interceptor/template.rs +++ b/sqlness/src/interceptor/template.rs @@ -3,11 +3,13 @@ use minijinja::Environment; use serde_json::Value; -use super::{Interceptor, InterceptorFactory, InterceptorRef}; +use crate::error::Result; +use crate::interceptor::{Interceptor, InterceptorFactory, InterceptorRef}; +use crate::SqlnessError; pub struct TemplateInterceptorFactory; -const PREFIX: &str = "TEMPLATE"; +pub const PREFIX: &str = "TEMPLATE"; /// Templated query, powered by [minijinja](https://github.com/mitsuhiko/minijinja). /// The template syntax can be found [here](https://docs.rs/minijinja/latest/minijinja/syntax/index.html). @@ -37,10 +39,10 @@ const PREFIX: &str = "TEMPLATE"; /// #[derive(Debug)] pub struct TemplateInterceptor { - json_ctx: String, + data_bindings: Value, } -fn sql_delimiter() -> Result { +fn sql_delimiter() -> std::result::Result { Ok(";".to_string()) } @@ -51,12 +53,7 @@ impl Interceptor for TemplateInterceptor { env.add_function("sql_delimiter", sql_delimiter); env.add_template("sql", &input).unwrap(); let tmpl = env.get_template("sql").unwrap(); - let bindings: Value = if self.json_ctx.is_empty() { - serde_json::from_str("{}").unwrap() - } else { - serde_json::from_str(&self.json_ctx).unwrap() - }; - let rendered = tmpl.render(bindings).unwrap(); + let rendered = tmpl.render(&self.data_bindings).unwrap(); *execute_query = rendered .split('\n') .map(|v| v.to_string()) @@ -67,19 +64,18 @@ impl Interceptor for TemplateInterceptor { } impl InterceptorFactory for TemplateInterceptorFactory { - fn try_new(&self, interceptor: &str) -> Option { - Self::try_new_from_str(interceptor).map(|i| Box::new(i) as _) - } -} - -impl TemplateInterceptorFactory { - fn try_new_from_str(interceptor: &str) -> Option { - if interceptor.starts_with(PREFIX) { - let json_ctx = interceptor.trim_start_matches(PREFIX).to_string(); - Some(TemplateInterceptor { json_ctx }) + fn try_new(&self, ctx: &str) -> Result { + let data_bindings = if ctx.is_empty() { + serde_json::from_str("{}") } else { - None + serde_json::from_str(ctx) } + .map_err(|e| SqlnessError::InvalidContext { + prefix: PREFIX.to_string(), + msg: format!("Expect json, err:{e}"), + })?; + + Ok(Box::new(TemplateInterceptor { data_bindings })) } } @@ -90,7 +86,7 @@ mod tests { #[test] fn basic_template() { let interceptor = TemplateInterceptorFactory - .try_new(r#"TEMPLATE {"name": "test"}"#) + .try_new(r#"{"name": "test"}"#) .unwrap(); let mut input = vec!["SELECT * FROM table where name = '{{name}}'".to_string()]; @@ -102,7 +98,7 @@ mod tests { #[test] fn vector_template() { let interceptor = TemplateInterceptorFactory - .try_new(r#"TEMPLATE {"aggr": ["sum", "count", "avg"]}"#) + .try_new(r#"{"aggr": ["sum", "count", "avg"]}"#) .unwrap(); let mut input = [ @@ -129,7 +125,7 @@ mod tests { #[test] fn range_template() { - let interceptor = TemplateInterceptorFactory.try_new(r#"TEMPLATE"#).unwrap(); + let interceptor = TemplateInterceptorFactory.try_new(r#""#).unwrap(); let mut input = [ "INSERT INTO t (c) VALUES", diff --git a/sqlness/src/lib.rs b/sqlness/src/lib.rs index ea3f4cf..b14f5fc 100644 --- a/sqlness/src/lib.rs +++ b/sqlness/src/lib.rs @@ -68,5 +68,4 @@ pub use config::{Config, ConfigBuilder, DatabaseConfig, DatabaseConfigBuilder}; pub use database::Database; pub use environment::EnvController; pub use error::SqlnessError; -pub use interceptor::builtin_interceptors; pub use runner::Runner;