refactor: add *_async to trait

CeresDB · Jun 17, 2024 · 7bf1a2a · 7bf1a2a
1 parent f5ed4a3
commit 7bf1a2a
Show file tree

Hide file tree

Showing 8 changed files with 49 additions and 51 deletions.
diff --git a/sqlness/src/case.rs b/sqlness/src/case.rs
@@ -177,16 +177,15 @@ impl Query {
 
         for interceptor in &self.interceptors {
             interceptor
-                .before_execute(&mut self.execute_query, &mut context)
-                .await;
+                .before_execute_async(&mut self.execute_query, &mut context).await;
         }
 
         context
     }
 
     async fn after_execute_intercept(&mut self, result: &mut String) {
         for interceptor in &self.interceptors {
-            interceptor.after_execute(result).await;
+            interceptor.after_execute_async(result).await;
         }
     }
 

diff --git a/sqlness/src/interceptor.rs b/sqlness/src/interceptor.rs
@@ -26,10 +26,24 @@ pub type InterceptorRef = Box<dyn Interceptor + Send + Sync>;
 #[async_trait::async_trait]
 pub trait Interceptor {
     #[allow(unused_variables)]
-    async fn before_execute(&self, execute_query: &mut Vec<String>, context: &mut QueryContext) {}
+    fn before_execute(&self, execute_query: &mut Vec<String>, context: &mut QueryContext) {}
 
     #[allow(unused_variables)]
-    async fn after_execute(&self, result: &mut String) {}
+    async fn before_execute_async(
+        &self,
+        execute_query: &mut Vec<String>,
+        context: &mut QueryContext,
+    ) {
+        self.before_execute(execute_query, context)
+    }
+
+    #[allow(unused_variables)]
+    fn after_execute(&self, result: &mut String) {}
+
+    #[allow(unused_variables)]
+    async fn after_execute_async(&self, result: &mut String) {
+        self.after_execute(result)
+    }
 }
 
 pub type InterceptorFactoryRef = Arc<dyn InterceptorFactory>;

diff --git a/sqlness/src/interceptor/arg.rs b/sqlness/src/interceptor/arg.rs
@@ -25,9 +25,8 @@ pub struct ArgInterceptor {
     args: Vec<(String, String)>,
 }
 
-#[async_trait::async_trait]
 impl Interceptor for ArgInterceptor {
-    async fn before_execute(&self, _: &mut Vec<String>, context: &mut QueryContext) {
+    fn before_execute(&self, _: &mut Vec<String>, context: &mut QueryContext) {
         for (key, value) in &self.args {
             context.context.insert(key.to_string(), value.to_string());
         }

diff --git a/sqlness/src/interceptor/env.rs b/sqlness/src/interceptor/env.rs
@@ -39,9 +39,8 @@ pub struct EnvInterceptor {
     data: HashMap<String, String>,
 }
 
-#[async_trait::async_trait]
 impl Interceptor for EnvInterceptor {
-    async fn before_execute(&self, execute_query: &mut Vec<String>, _: &mut QueryContext) {
+    fn before_execute(&self, execute_query: &mut Vec<String>, _: &mut QueryContext) {
         for line in execute_query {
             for (key, value) in &self.data {
                 let rendered = line.replace(key, value);

diff --git a/sqlness/src/interceptor/replace.rs b/sqlness/src/interceptor/replace.rs
@@ -39,9 +39,8 @@ pub struct ReplaceInterceptor {
     replacement: String,
 }
 
-#[async_trait::async_trait]
 impl Interceptor for ReplaceInterceptor {
-    async fn after_execute(&self, result: &mut String) {
+    fn after_execute(&self, result: &mut String) {
         let re = Regex::new(&self.pattern).unwrap();
         let replaced = re.replace_all(result, &self.replacement);
         *result = replaced.to_string();
@@ -85,21 +84,21 @@ mod tests {
         assert!(interceptor.is_err());
     }
 
-    #[tokio::test]
-    async fn replace_without_replacement() {
+    #[test]
+    fn replace_without_replacement() {
         let interceptor = ReplaceInterceptorFactory {}.try_new("0").unwrap();
 
         let mut exec_result = "000010101".to_string();
-        interceptor.after_execute(&mut exec_result).await;
+        interceptor.after_execute(&mut exec_result);
         assert_eq!(exec_result, "111".to_string());
     }
 
-    #[tokio::test]
-    async fn simple_replace() {
+    #[test]
+    fn simple_replace() {
         let interceptor = ReplaceInterceptorFactory {}.try_new("00 2").unwrap();
 
         let mut exec_result = "0000010101".to_string();
-        interceptor.after_execute(&mut exec_result).await;
+        interceptor.after_execute(&mut exec_result);
         assert_eq!(exec_result, "22010101".to_string());
     }
 }
diff --git a/sqlness/src/interceptor/sleep.rs b/sqlness/src/interceptor/sleep.rs
@@ -41,7 +41,7 @@ pub struct SleepInterceptor {
 
 #[async_trait::async_trait]
 impl Interceptor for SleepInterceptor {
-    async fn before_execute(
+    async fn before_execute_async(
         &self,
         _execute_query: &mut Vec<String>,
         _context: &mut crate::case::QueryContext,
@@ -102,8 +102,8 @@ mod test {
         let input = "1500";
         let interceptor = SleepInterceptorFactory{}.try_new(input).unwrap();
         let now = Instant::now();
-        interceptor.before_execute(&mut vec![], &mut crate::QueryContext::default()).await;
+        interceptor.before_execute_async(&mut vec![], &mut crate::QueryContext::default()).await;
         let elasped = now.elapsed().as_millis() as u64;
         assert!(elasped >= 1500);
     }
-}
+}
diff --git a/sqlness/src/interceptor/sort_result.rs b/sqlness/src/interceptor/sort_result.rs
@@ -43,9 +43,8 @@ pub struct SortResultInterceptor {
     ignore_tail: usize,
 }
 
-#[async_trait::async_trait]
 impl Interceptor for SortResultInterceptor {
-    async fn after_execute(&self, result: &mut String) {
+    fn after_execute(&self, result: &mut String) {
         let mut lines = result.lines().collect::<VecDeque<_>>();
         let mut head = Vec::with_capacity(self.ignore_head);
         let mut tail = Vec::with_capacity(self.ignore_tail);
@@ -114,8 +113,8 @@ mod tests {
         assert!(interceptor.is_err());
     }
 
-    #[tokio::test]
-    async fn sort_result_full() {
+    #[test]
+    fn sort_result_full() {
         let interceptor = SortResultInterceptorFactory.try_new("").unwrap();
 
         let cases = [
@@ -146,13 +145,13 @@ mod tests {
         ];
 
         for (mut input, expected) in cases {
-            interceptor.after_execute(&mut input).await;
+            interceptor.after_execute(&mut input);
             assert_eq!(input, expected);
         }
     }
 
-    #[tokio::test]
-    async fn ignore_head_exceeds_length() {
+    #[test]
+    fn ignore_head_exceeds_length() {
         let interceptor = SortResultInterceptorFactory.try_new("10000").unwrap();
 
         let mut exec_result = String::from(
@@ -161,12 +160,12 @@ mod tests {
             \n1",
         );
         let expected = exec_result.clone();
-        interceptor.after_execute(&mut exec_result).await;
+        interceptor.after_execute(&mut exec_result);
         assert_eq!(exec_result, expected);
     }
 
-    #[tokio::test]
-    async fn ignore_tail_exceeds_length() {
+    #[test]
+    fn ignore_tail_exceeds_length() {
         let interceptor = SortResultInterceptorFactory.try_new("0 10000").unwrap();
 
         let mut exec_result = String::from(
@@ -175,7 +174,7 @@ mod tests {
             \n1",
         );
         let expected = exec_result.clone();
-        interceptor.after_execute(&mut exec_result).await;
+        interceptor.after_execute(&mut exec_result);
         assert_eq!(exec_result, expected);
     }
 }
diff --git a/sqlness/src/interceptor/template.rs b/sqlness/src/interceptor/template.rs
@@ -46,13 +46,8 @@ fn sql_delimiter() -> std::result::Result<String, minijinja::Error> {
     Ok(";".to_string())
 }
 
-#[async_trait::async_trait]
 impl Interceptor for TemplateInterceptor {
-    async fn before_execute(
-        &self,
-        execute_query: &mut Vec<String>,
-        _context: &mut crate::QueryContext,
-    ) {
+    fn before_execute(&self, execute_query: &mut Vec<String>, _context: &mut crate::QueryContext) {
         let input = execute_query.join("\n");
         let mut env = Environment::new();
         env.add_function("sql_delimiter", sql_delimiter);
@@ -86,22 +81,20 @@ impl InterceptorFactory for TemplateInterceptorFactory {
 mod tests {
     use super::*;
 
-    #[tokio::test]
-    async fn basic_template() {
+    #[test]
+    fn basic_template() {
         let interceptor = TemplateInterceptorFactory
             .try_new(r#"{"name": "test"}"#)
             .unwrap();
 
         let mut input = vec!["SELECT * FROM table where name = '{{name}}'".to_string()];
-        interceptor
-            .before_execute(&mut input, &mut crate::QueryContext::default())
-            .await;
+        interceptor.before_execute(&mut input, &mut crate::QueryContext::default());
 
         assert_eq!(input, vec!["SELECT * FROM table where name = 'test'"]);
     }
 
-    #[tokio::test]
-    async fn vector_template() {
+    #[test]
+    fn vector_template() {
         let interceptor = TemplateInterceptorFactory
             .try_new(r#"{"aggr": ["sum", "count", "avg"]}"#)
             .unwrap();
@@ -113,9 +106,7 @@ mod tests {
         ]
         .map(|v| v.to_string())
         .to_vec();
-        interceptor
-            .before_execute(&mut input, &mut crate::QueryContext::default())
-            .await;
+        interceptor.before_execute(&mut input, &mut crate::QueryContext::default());
 
         assert_eq!(
             input,
@@ -130,8 +121,7 @@ mod tests {
         );
     }
 
-    #[tokio::test]
-    async fn range_template() {
+    fn range_template() {
         let interceptor = TemplateInterceptorFactory.try_new(r#""#).unwrap();
 
         let mut input = [
@@ -144,8 +134,7 @@ mod tests {
         .map(|v| v.to_string())
         .to_vec();
         interceptor
-            .before_execute(&mut input, &mut crate::QueryContext::default())
-            .await;
+            .before_execute(&mut input, &mut crate::QueryContext::default());
 
         assert_eq!(
             input,