feat: implement ENV interceptor (#56)

Signed-off-by: Ruihang Xia <[email protected]>

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

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<dyn Display> {
+        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();
+}

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;
+

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;

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<y>\d{4})-(?P<m>\d{2})-(?P<d>\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;
+

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<y>\d{4})-(?P<m>\d{2})-(?P<d>\d{2}) $m/$d/$y
+2012-03-14, 2013-01-01 and 2014-07-05;

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;
+

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;

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<dyn Display> {
+        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();
+}

sqlness/src/case.rs

@@ -96,7 +96,10 @@ pub struct QueryContext {
 #[derive(Default)]
 struct Query {
     comment_lines: Vec<String>,
-    query_lines: Vec<String>,
+    /// Query to be displayed in the result file
+    display_query: Vec<String>,
+    /// Query to be executed
+    execute_query: Vec<String>,
     interceptor_factories: Vec<InterceptorFactoryRef>,
     interceptors: Vec<InterceptorRef>,
 }
@@ -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<W>(&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())?;

sqlness/src/interceptor.rs

@@ -7,20 +7,21 @@ 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;
 
 pub type InterceptorRef = Box<dyn Interceptor>;
 
 pub trait Interceptor {
     #[allow(unused_variables)]
-    fn before_execute(&self, query: &mut Vec<String>, context: &mut QueryContext) {}
+    fn before_execute(&self, execute_query: &mut Vec<String>, context: &mut QueryContext) {}
 
     #[allow(unused_variables)]
     fn after_execute(&self, result: &mut String) {}
@@ -37,6 +38,7 @@ pub fn builtin_interceptors() -> Vec<InterceptorFactoryRef> {
     vec![
         Arc::new(ArgInterceptorFactory {}),
         Arc::new(ReplaceInterceptorFactory {}),
+        Arc::new(EnvInterceptorFactory {}),
         Arc::new(SortResultInterceptorFactory {}),
     ]
 }