diff --git a/Makefile b/Makefile index a6db970..5e45779 100644 --- a/Makefile +++ b/Makefile @@ -27,7 +27,7 @@ cli-test: example: good-example bad-example -good-example: basic-example interceptor-arg-example interceptor-replace-example interceptor-sort-result-example +good-example: basic-example interceptor-arg-example interceptor-replace-example interceptor-sort-result-example interceptor-env-example basic-example: cd $(DIR)/sqlness; cargo run --example basic @@ -43,3 +43,6 @@ interceptor-replace-example: interceptor-sort-result-example: cd $(DIR)/sqlness; cargo run --example interceptor_sort_result + +interceptor-env-example: + cd $(DIR)/sqlness; cargo run --example interceptor_env diff --git a/sqlness/examples/echo.rs b/sqlness/examples/echo.rs new file mode 100644 index 00000000..10a3a49 --- /dev/null +++ b/sqlness/examples/echo.rs @@ -0,0 +1,53 @@ +// Copyright 2022 CeresDB Project Authors. Licensed under Apache-2.0. + +//! Shows how an REPLACE interceptor works. + +use std::{fmt::Display, path::Path}; + +use async_trait::async_trait; +use sqlness::{ConfigBuilder, Database, EnvController, QueryContext, Runner}; + +struct MyController; +struct MyDB; + +#[async_trait] +impl Database for MyDB { + async fn query(&self, _: QueryContext, query: String) -> Box { + return Box::new(query); + } +} + +impl MyDB { + fn new(_env: &str, _config: Option<&Path>) -> Self { + MyDB + } + + fn stop(self) {} +} + +#[async_trait] +impl EnvController for MyController { + type DB = MyDB; + + async fn start(&self, env: &str, config: Option<&Path>) -> Self::DB { + MyDB::new(env, config) + } + + async fn stop(&self, _env: &str, database: Self::DB) { + database.stop(); + } +} + +#[tokio::main] +async fn main() { + let env = MyController; + let config = ConfigBuilder::default() + .case_dir("examples/echo".to_string()) + .build() + .unwrap(); + let runner = Runner::new(config, env); + + println!("Run testcase..."); + + runner.run().await.unwrap(); +} diff --git a/sqlness/examples/echo/interceptor-env/env.result b/sqlness/examples/echo/interceptor-env/env.result new file mode 100644 index 00000000..596cb7f --- /dev/null +++ b/sqlness/examples/echo/interceptor-env/env.result @@ -0,0 +1,10 @@ +-- SQLNESS ENV SECRET +select 23333 from data; + +select 23333 from data; + +-- SQLNESS ENV SECRET NONEXISTENT +select 23333 from data; + +select 23333 from data; + diff --git a/sqlness/examples/echo/interceptor-env/env.sql b/sqlness/examples/echo/interceptor-env/env.sql new file mode 100644 index 00000000..5160163 --- /dev/null +++ b/sqlness/examples/echo/interceptor-env/env.sql @@ -0,0 +1,5 @@ +-- SQLNESS ENV SECRET +select {{ SECRET }} from data; + +-- SQLNESS ENV SECRET NONEXISTENT +select {{ SECRET }} from data; \ No newline at end of file diff --git a/sqlness/examples/echo/interceptor-replace/replace.result b/sqlness/examples/echo/interceptor-replace/replace.result new file mode 100644 index 00000000..60d17d6 --- /dev/null +++ b/sqlness/examples/echo/interceptor-replace/replace.result @@ -0,0 +1,21 @@ +-- SQLNESS REPLACE 00 +SELECT 0; + +SELECT 0; + +-- SQLNESS REPLACE 00 +SELECT 00; + +SELECT ; + +-- SQLNESS REPLACE 0 1 +SELECT 0; + +SELECT 1; + +-- example of capture group replacement +-- SQLNESS REPLACE (?P\d{4})-(?P\d{2})-(?P\d{2}) $m/$d/$y +2012-03-14, 2013-01-01 and 2014-07-05; + +03/14/2012, 01/01/2013 and 07/05/2014; + diff --git a/sqlness/examples/echo/interceptor-replace/replace.sql b/sqlness/examples/echo/interceptor-replace/replace.sql new file mode 100644 index 00000000..f47f540 --- /dev/null +++ b/sqlness/examples/echo/interceptor-replace/replace.sql @@ -0,0 +1,12 @@ +-- SQLNESS REPLACE 00 +SELECT 0; + +-- SQLNESS REPLACE 00 +SELECT 00; + +-- SQLNESS REPLACE 0 1 +SELECT 0; + +-- example of capture group replacement +-- SQLNESS REPLACE (?P\d{4})-(?P\d{2})-(?P\d{2}) $m/$d/$y +2012-03-14, 2013-01-01 and 2014-07-05; diff --git a/sqlness/examples/interceptor-env/simple/env.result b/sqlness/examples/interceptor-env/simple/env.result new file mode 100644 index 00000000..6a72536 --- /dev/null +++ b/sqlness/examples/interceptor-env/simple/env.result @@ -0,0 +1,22 @@ +-- multiple env in one line +-- SQLNESS ENV ENV1 ENV2 NONEXISTENT1 NONEXISTENT2 NONEXISTENT3 +SELECT $ENV1, $ENV2, $NONEXISTENT1 FROM t; + +SELECT value1, value2, $NONEXISTENT1 FROM t; + +-- multiple env in multiple lines +-- SQLNESS ENV ENV1 +-- SQLNESS ENV ENV2 +-- SQLNESS ENV NONEXISTENT1 +-- SQLNESS ENV NONEXISTENT2 +-- SQLNESS ENV NONEXISTENT3 +SELECT $ENV1, $ENV2, $NONEXISTENT1, FROM t; + +SELECT value1, value2, $NONEXISTENT1, FROM t; + +-- Undeclared env won't be rendered +-- SQLNESS ENV ENV2 +SELECT $ENV1, $ENV2, $NONEXISTENT1 FROM t; + +SELECT $ENV1, value2, $NONEXISTENT1 FROM t; + diff --git a/sqlness/examples/interceptor-env/simple/env.sql b/sqlness/examples/interceptor-env/simple/env.sql new file mode 100644 index 00000000..e237480 --- /dev/null +++ b/sqlness/examples/interceptor-env/simple/env.sql @@ -0,0 +1,15 @@ +-- multiple env in one line +-- SQLNESS ENV ENV1 ENV2 NONEXISTENT1 NONEXISTENT2 NONEXISTENT3 +SELECT $ENV1, $ENV2, $NONEXISTENT1 FROM t; + +-- multiple env in multiple lines +-- SQLNESS ENV ENV1 +-- SQLNESS ENV ENV2 +-- SQLNESS ENV NONEXISTENT1 +-- SQLNESS ENV NONEXISTENT2 +-- SQLNESS ENV NONEXISTENT3 +SELECT $ENV1, $ENV2, $NONEXISTENT1, FROM t; + +-- Undeclared env won't be rendered +-- SQLNESS ENV ENV2 +SELECT $ENV1, $ENV2, $NONEXISTENT1 FROM t; diff --git a/sqlness/examples/interceptor_env.rs b/sqlness/examples/interceptor_env.rs new file mode 100644 index 00000000..1ae0ab5 --- /dev/null +++ b/sqlness/examples/interceptor_env.rs @@ -0,0 +1,55 @@ +// Copyright 2023 CeresDB Project Authors. Licensed under Apache-2.0. + +//! Shows how an SORT_RESULT interceptor works. + +use std::{fmt::Display, path::Path}; + +use async_trait::async_trait; +use sqlness::{ConfigBuilder, Database, EnvController, QueryContext, Runner}; + +struct MyController; +struct MyDB; + +#[async_trait] +impl Database for MyDB { + async fn query(&self, _: QueryContext, query: String) -> Box { + return Box::new(query); + } +} + +impl MyDB { + fn new(_env: &str, _config: Option<&Path>) -> Self { + MyDB + } + + fn stop(self) {} +} + +#[async_trait] +impl EnvController for MyController { + type DB = MyDB; + + async fn start(&self, env: &str, config: Option<&Path>) -> Self::DB { + MyDB::new(env, config) + } + + async fn stop(&self, _env: &str, database: Self::DB) { + database.stop(); + } +} + +#[tokio::main] +async fn main() { + std::env::set_var("ENV1", "value1"); + std::env::set_var("ENV2", "value2"); + let env = MyController; + let config = ConfigBuilder::default() + .case_dir("examples/interceptor-env".to_string()) + .build() + .unwrap(); + let runner = Runner::new(config, env); + + println!("Run testcase..."); + + runner.run().await.unwrap(); +} diff --git a/sqlness/src/case.rs b/sqlness/src/case.rs index 36b484a..ed581eb 100644 --- a/sqlness/src/case.rs +++ b/sqlness/src/case.rs @@ -96,7 +96,10 @@ pub struct QueryContext { #[derive(Default)] struct Query { comment_lines: Vec, - query_lines: Vec, + /// Query to be displayed in the result file + display_query: Vec, + /// Query to be executed + execute_query: Vec, interceptor_factories: Vec, interceptors: Vec, } @@ -125,7 +128,8 @@ impl Query { } fn append_query_line(&mut self, line: &str) { - self.query_lines.push(line.to_string()); + self.display_query.push(line.to_string()); + self.execute_query.push(line.to_string()); } async fn execute(&mut self, db: &dyn Database, writer: &mut W) -> Result<()> @@ -145,11 +149,15 @@ impl Query { Ok(()) } + /// Run pre-execution interceptors. + /// + /// Interceptors may change either the query to be displayed or the query to be executed, + /// so we need to return the query to caller. fn before_execute_intercept(&mut self) -> QueryContext { let mut context = QueryContext::default(); for interceptor in &self.interceptors { - interceptor.before_execute(&mut self.query_lines, &mut context); + interceptor.before_execute(&mut self.execute_query, &mut context); } context @@ -161,8 +169,9 @@ impl Query { } } + /// Concat the query to be executed to a single string. fn concat_query_lines(&self) -> String { - self.query_lines + self.execute_query .iter() .fold(String::new(), |query, str| query + str) .trim_start() @@ -178,7 +187,7 @@ impl Query { writer.write_all(comment.as_bytes())?; writer.write("\n".as_bytes())?; } - for line in &self.query_lines { + for line in &self.display_query { writer.write_all(line.as_bytes())?; } writer.write("\n\n".as_bytes())?; diff --git a/sqlness/src/interceptor.rs b/sqlness/src/interceptor.rs index 2aea588..094e998 100644 --- a/sqlness/src/interceptor.rs +++ b/sqlness/src/interceptor.rs @@ -7,12 +7,13 @@ use std::sync::Arc; use crate::{ case::QueryContext, interceptor::{ - arg::ArgInterceptorFactory, replace::ReplaceInterceptorFactory, + arg::ArgInterceptorFactory, env::EnvInterceptorFactory, replace::ReplaceInterceptorFactory, sort_result::SortResultInterceptorFactory, }, }; pub mod arg; +pub mod env; pub mod replace; pub mod sort_result; @@ -20,7 +21,7 @@ pub type InterceptorRef = Box; pub trait Interceptor { #[allow(unused_variables)] - fn before_execute(&self, query: &mut Vec, context: &mut QueryContext) {} + fn before_execute(&self, execute_query: &mut Vec, context: &mut QueryContext) {} #[allow(unused_variables)] fn after_execute(&self, result: &mut String) {} @@ -37,6 +38,7 @@ pub fn builtin_interceptors() -> Vec { vec![ Arc::new(ArgInterceptorFactory {}), Arc::new(ReplaceInterceptorFactory {}), + Arc::new(EnvInterceptorFactory {}), Arc::new(SortResultInterceptorFactory {}), ] } diff --git a/sqlness/src/interceptor/env.rs b/sqlness/src/interceptor/env.rs new file mode 100644 index 00000000..e6c7259 --- /dev/null +++ b/sqlness/src/interceptor/env.rs @@ -0,0 +1,99 @@ +// Copyright 2023 CeresDB Project Authors. Licensed under Apache-2.0. + +use std::collections::HashMap; + +use crate::case::QueryContext; +use crate::interceptor::{Interceptor, InterceptorFactory, InterceptorRef}; + +const PREFIX: &str = "ENV"; + +/// Read environment variables and fill them in query. +/// +/// # Example +/// ``` sql +/// -- SQLNESS ENV SECRET +/// SELECT $SECRET; +/// ``` +/// +/// Environment variables declared in `ENV` interceptor will be replaced in the +/// going to be executed. It won't be rendered in the result file so you can +/// safely put secret things in your query. +/// +/// Note that only decalred and present environment variables will be replaced. +/// +/// You can either declare multiple env in one intercetor or separate them into +/// different interceptors. The following two examples are equivalent: +/// +/// ``` sql +/// -- SQLNESS ENV SECRET1 SECRET2 +/// SELECT $SECRET1, $SECRET2; +/// +/// -- SQLNESS ENV SECRET1 +/// -- SQLNESS ENV SECRET2 +/// SELECT $SECRET1, $SECRET2; +/// ```` +#[derive(Debug)] +pub struct EnvInterceptor { + /// Environment variables to be replaced. + data: HashMap, +} + +impl Interceptor for EnvInterceptor { + fn before_execute(&self, execute_query: &mut Vec, _: &mut QueryContext) { + for line in execute_query { + for (key, value) in &self.data { + let rendered = line.replace(key, value); + *line = rendered; + } + } + } +} + +pub struct EnvInterceptorFactory; + +impl InterceptorFactory for EnvInterceptorFactory { + fn try_new(&self, interceptor: &str) -> Option { + Self::create(interceptor).map(|i| Box::new(i) as InterceptorRef) + } +} + +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::>(); + + let mut env_data = HashMap::new(); + for env in envs { + if let Ok(value) = std::env::var(env) { + env_data.insert(format!("${env}"), value); + } + } + + Some(EnvInterceptor { data: env_data }) + } else { + None + } + } +} + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn cut_env_string() { + let input = "ENV SECRET NONEXISTENT"; + std::env::set_var("SECRET", "2333"); + + let expected = [("$SECRET".to_string(), "2333".to_string())] + .into_iter() + .collect(); + + let interceptor = EnvInterceptorFactory::create(input).unwrap(); + assert_eq!(interceptor.data, expected); + } +} diff --git a/sqlness/src/interceptor/replace.rs b/sqlness/src/interceptor/replace.rs index 3922537..c4b8c99 100644 --- a/sqlness/src/interceptor/replace.rs +++ b/sqlness/src/interceptor/replace.rs @@ -1,4 +1,4 @@ -// Copyright 2022 CeresDB Project Authors. Licensed under Apache-2.0. +// Copyright 2023 CeresDB Project Authors. Licensed under Apache-2.0. use regex::Regex;