Make InterpreterFactor.get to async

Author: BohuTANG

src/query/service/src/interpreters/access/accessor.rs

@@ -26,7 +26,7 @@ use crate::sql::plans::Plan;
 pub trait AccessChecker: Sync + Send {
     // Check the access permission for the old plan.
     // TODO(bohu): Remove after new plan done.
-    fn check(&self, plan: &PlanNode) -> Result<()>;
+    async fn check(&self, plan: &PlanNode) -> Result<()>;
 
     // Check the access permission for the old plan.
     async fn check_new(&self, _plan: &Plan) -> Result<()>;
@@ -43,9 +43,9 @@ impl Accessor {
         Accessor { accessors }
     }
 
-    pub fn check(&self, plan: &PlanNode) -> Result<()> {
+    pub async fn check(&self, plan: &PlanNode) -> Result<()> {
         for accessor in self.accessors.values() {
-            accessor.check(plan)?;
+            accessor.check(plan).await?;
         }
         Ok(())
     }

src/query/service/src/interpreters/access/management_mode_access.rs

@@ -36,7 +36,7 @@ impl ManagementModeAccess {
 #[async_trait::async_trait]
 impl AccessChecker for ManagementModeAccess {
     // Check what we can do if in management mode.
-    fn check(&self, plan: &PlanNode) -> Result<()> {
+    async fn check(&self, plan: &PlanNode) -> Result<()> {
         // Allows for management-mode.
         if self.ctx.get_config().query.management_mode {
             return match plan {

src/query/service/src/interpreters/interpreter_factory.rs

@@ -36,11 +36,12 @@ pub struct InterpreterFactory;
 /// InterpreterFactory provides `get` method which transforms the PlanNode into the corresponding interpreter.
 /// Such as: SelectPlan -> SelectInterpreter, ExplainPlan -> ExplainInterpreter, ...
 impl InterpreterFactory {
-    pub fn get(ctx: Arc<QueryContext>, plan: PlanNode) -> Result<Arc<dyn Interpreter>> {
+    pub async fn get(ctx: Arc<QueryContext>, plan: PlanNode) -> Result<Arc<dyn Interpreter>> {
         // Check the access permission.
         let access_checker = Accessor::create(ctx.clone());
         access_checker
             .check(&plan)
+            .await
             .map(|e| error!("Access.denied(v1): {:?}", e))?;
 
         let inner = Self::create_interpreter(ctx.clone(), &plan)?;

src/query/service/src/servers/http/clickhouse_handler.rs

@@ -235,7 +235,9 @@ pub async fn clickhouse_handler_get(
         context.attach_query_str(&sql);
         let format = get_format_with_default(format, default_format)?;
         let schema = plan.schema();
-        let interpreter = InterpreterFactory::get(context.clone(), plan).map_err(BadRequest)?;
+        let interpreter = InterpreterFactory::get(context.clone(), plan)
+            .await
+            .map_err(BadRequest)?;
         execute(context, interpreter, schema, format, None, params)
             .await
             .map_err(InternalServerError)
@@ -312,7 +314,9 @@ pub async fn clickhouse_handler_post(
         let format = get_format_with_default(format, default_format)?;
 
         let schema = plan.schema();
-        let interpreter = InterpreterFactory::get(ctx.clone(), plan).map_err(BadRequest)?;
+        let interpreter = InterpreterFactory::get(ctx.clone(), plan)
+            .await
+            .map_err(BadRequest)?;
         execute(ctx, interpreter, schema, format, None, params)
             .await
             .map_err(InternalServerError)

src/query/service/src/servers/http/v1/load.rs

@@ -72,7 +72,7 @@ fn execute_query(
     source_builder: SourcePipeBuilder,
 ) -> impl Future<Output = Result<()>> {
     async move {
-        let interpreter = InterpreterFactory::get(context.clone(), node)?;
+        let interpreter = InterpreterFactory::get(context.clone(), node).await?;
 
         if let Err(cause) = interpreter.start().await {
             error!("interpreter.start error: {:?}", cause);
@@ -202,8 +202,9 @@ pub async fn streaming_load(
             StatusCode::BAD_REQUEST,
         )),
     }?;
-    let interpreter =
-        InterpreterFactory::get(context.clone(), plan.clone()).map_err(InternalServerError)?;
+    let interpreter = InterpreterFactory::get(context.clone(), plan.clone())
+        .await
+        .map_err(InternalServerError)?;
     let _ = interpreter
         .set_source_pipe_builder(Some(source_pipe_builder))
         .map_err(|e| error!("interpreter.set_source_pipe_builder.error: {:?}", e));

src/query/service/src/servers/http/v1/query/execute_state.rs

@@ -256,7 +256,7 @@ impl ExecuteState {
             };
 
             is_select = matches!(&plan, PlanNode::Select(_));
-            InterpreterFactory::get(ctx.clone(), plan)
+            InterpreterFactory::get(ctx.clone(), plan).await
         }?;
 
         if is_v2 && is_select {

src/query/service/src/servers/mysql/mysql_interactive_worker.rs

@@ -344,18 +344,23 @@ impl<W: AsyncWrite + Send + Unpin> InteractiveWorkerBase<W> {
                     Ok((_, h)) => h.clone(),
                     Err(_) => vec![],
                 };
-                let mut has_result_set = true;
-                let interpreter = if settings.get_enable_planner_v2()? != 0 {
+
+                let (interpreter, has_result_set) = if settings.get_enable_planner_v2()? != 0 {
                     let mut planner = Planner::new(context.clone());
-                    let plan_res = planner.plan_sql(query).await?;
-                    has_result_set = has_result_set_by_plan(&plan_res.0);
-                    InterpreterFactoryV2::get(context.clone(), &plan_res.0).await
+                    let plan = planner.plan_sql(query).await?;
+                    let has_result_set = has_result_set_by_plan(&plan.0);
+                    (
+                        InterpreterFactoryV2::get(context.clone(), &plan.0).await,
+                        has_result_set,
+                    )
                 } else {
-                    let (plan, _) = PlanParser::parse_with_hint(query, context.clone()).await;
-                    plan.and_then(|v| {
-                        has_result_set = has_result_set_by_plan_node(&v);
-                        InterpreterFactory::get(context.clone(), v)
-                    })
+                    let (plan_res, _) = PlanParser::parse_with_hint(query, context.clone()).await;
+                    let plan = plan_res?;
+                    let has_result_set = has_result_set_by_plan_node(&plan);
+                    (
+                        InterpreterFactory::get(context.clone(), plan).await,
+                        has_result_set,
+                    )
                 };
 
                 let hint = hints

src/query/service/tests/it/api/http/status.rs

@@ -62,7 +62,7 @@ async fn run_query(query_ctx: &Arc<QueryContext>) -> Result<Arc<dyn Interpreter>
         .await?;
     query_ctx.set_current_user(user);
     let plan = PlanParser::parse(query_ctx.clone(), sql).await?;
-    InterpreterFactory::get(query_ctx.clone(), plan)
+    InterpreterFactory::get(query_ctx.clone(), plan).await
 }
 
 #[tokio::test]

src/query/service/tests/it/interpreters/async_insert_queue.rs

@@ -134,7 +134,7 @@ async fn test_async_insert_queue() -> Result<()> {
     {
         let query = "select * from default.test";
         let plan = PlanParser::parse(ctx.clone(), query).await?;
-        let executor = InterpreterFactory::get(ctx.clone(), plan.clone())?;
+        let executor = InterpreterFactory::get(ctx.clone(), plan.clone()).await?;
         let stream = executor.execute(ctx.clone()).await?;
         let result = stream.try_collect::<Vec<_>>().await?;
         let expected = vec![
@@ -390,7 +390,7 @@ async fn test_async_insert_queue_no_wait() -> Result<()> {
     {
         let query = "select * from default.test";
         let plan = PlanParser::parse(ctx.clone(), query).await?;
-        let executor = InterpreterFactory::get(ctx.clone(), plan.clone())?;
+        let executor = InterpreterFactory::get(ctx.clone(), plan.clone()).await?;
         let stream = executor.execute(ctx.clone()).await?;
         let result = stream.try_collect::<Vec<_>>().await?;
         let expected = vec![

src/query/service/tests/it/interpreters/interpreter_explain.rs

@@ -29,7 +29,7 @@ async fn test_explain_interpreter() -> Result<()> {
     ";
 
     let plan = PlanParser::parse(ctx.clone(), query).await?;
-    let executor = InterpreterFactory::get(ctx.clone(), plan)?;
+    let executor = InterpreterFactory::get(ctx.clone(), plan).await?;
     assert_eq!(executor.name(), "ExplainInterpreter");
 
     let stream = executor.execute(ctx).await?;