chore(cubesql): Do not call async Node functions while planning

Author: MazterQyou

rust/cubesql/cubesql/src/compile/query_engine.rs

@@ -43,6 +43,7 @@ use datafusion::{
     sql::{parser::Statement as DFStatement, planner::SqlToRel},
     variable::VarType,
 };
+use uuid::Uuid;
 
 #[async_trait::async_trait]
 pub trait QueryEngine {
@@ -81,6 +82,7 @@ pub trait QueryEngine {
         span_id: Option<Arc<SpanId>>,
         meta: Arc<MetaContext>,
         state: Arc<SessionState>,
+        compiler_id: Uuid,
     ) -> CompilationResult<QueryPlan> {
         let ctx = self.create_session_ctx(state.clone())?;
         let cube_ctx = self.create_cube_ctx(state.clone(), meta.clone(), ctx.clone())?;
@@ -140,7 +142,7 @@ pub trait QueryEngine {
         let mut finalized_graph = self
             .compiler_cache_ref()
             .rewrite(
-                state.auth_context().unwrap(),
+                compiler_id,
                 cube_ctx.clone(),
                 converter.take_egraph(),
                 &query_params.unwrap(),
@@ -186,7 +188,13 @@ pub trait QueryEngine {
         let mut rewriter = Rewriter::new(finalized_graph, cube_ctx.clone());
 
         let result = rewriter
-            .find_best_plan(root, state.auth_context().unwrap(), qtrace, span_id.clone())
+            .find_best_plan(
+                root,
+                compiler_id,
+                state.auth_context().unwrap(),
+                qtrace,
+                span_id.clone(),
+            )
             .await
             .map_err(|e| match e.cause {
                 CubeErrorCauseType::Internal(_) => CompilationError::Internal(

rust/cubesql/cubesql/src/compile/rewrite/rewriter.rs

@@ -31,6 +31,7 @@ use std::{
     sync::Arc,
     time::Duration,
 };
+use uuid::Uuid;
 
 pub struct Rewriter {
     graph: EGraph<LogicalPlanLanguage, LogicalPlanAnalysis>,
@@ -229,7 +230,7 @@ impl Rewriter {
 
     pub async fn run_rewrite_to_completion(
         &mut self,
-        auth_context: AuthContextRef,
+        compiler_id: Uuid,
         qtrace: &mut Option<Qtrace>,
     ) -> Result<EGraph<LogicalPlanLanguage, LogicalPlanAnalysis>, CubeError> {
         let cube_context = self.cube_context.clone();
@@ -243,7 +244,7 @@ impl Rewriter {
             .server
             .compiler_cache
             .rewrite_rules(
-                auth_context.clone(),
+                compiler_id,
                 cube_context.session_state.protocol.clone(),
                 false,
             )
@@ -311,6 +312,7 @@ impl Rewriter {
     pub async fn find_best_plan(
         &mut self,
         root: Id,
+        compiler_id: Uuid,
         auth_context: AuthContextRef,
         qtrace: &mut Option<Qtrace>,
         span_id: Option<Arc<SpanId>>,
@@ -326,7 +328,7 @@ impl Rewriter {
             .server
             .compiler_cache
             .rewrite_rules(
-                auth_context.clone(),
+                compiler_id,
                 cube_context.session_state.protocol.clone(),
                 true,
             )

rust/cubesql/cubesql/src/compile/router.rs

@@ -31,6 +31,7 @@ use datafusion::{
 };
 use itertools::Itertools;
 use sqlparser::ast::{escape_single_quote_string, ObjectName};
+use uuid::Uuid;
 
 #[derive(Clone)]
 pub struct QueryRouter {
@@ -52,6 +53,20 @@ impl QueryRouter {
         }
     }
 
+    async fn get_compiler_id_and_refresh_cache_if_needed(&self) -> CompilationResult<Uuid> {
+        self.session_manager
+            .server
+            .compiler_cache
+            .get_compiler_id_and_refresh_if_needed(
+                self.state.auth_context().ok_or_else(|| {
+                    CompilationError::internal("Unable to get auth context".to_string())
+                })?,
+                self.state.protocol.clone(),
+            )
+            .await
+            .map_err(|e| CompilationError::internal(e.to_string()))
+    }
+
     /// Common case for both planners: meta & olap
     /// This method tries to detect what planner to use as earlier as possible
     /// and forward context to correct planner
@@ -61,6 +76,9 @@ impl QueryRouter {
         qtrace: &mut Option<Qtrace>,
         span_id: Option<Arc<SpanId>>,
     ) -> CompilationResult<QueryPlan> {
+        self.reauthenticate_if_needed().await?;
+        let compiler_id = self.get_compiler_id_and_refresh_cache_if_needed().await?;
+
         let planning_start = SystemTime::now();
         if let Some(span_id) = span_id.as_ref() {
             if let Some(auth_context) = self.state.auth_context() {
@@ -81,7 +99,7 @@ impl QueryRouter {
             }
         }
         let result = self
-            .create_df_logical_plan(stmt.clone(), qtrace, span_id.clone())
+            .create_df_logical_plan(stmt.clone(), qtrace, span_id.clone(), compiler_id)
             .await?;
 
         if let Some(span_id) = span_id.as_ref() {
@@ -230,6 +248,7 @@ impl QueryRouter {
         variable: &Vec<ast::Ident>,
         span_id: Option<Arc<SpanId>>,
     ) -> CompilationResult<QueryPlan> {
+        let compiler_id = self.get_compiler_id_and_refresh_cache_if_needed().await?;
         let name = variable.to_vec()[0].value.clone();
         if self.state.protocol == DatabaseProtocol::PostgreSQL {
             let full_variable = variable.iter().map(|v| v.value.to_lowercase()).join("_");
@@ -256,7 +275,7 @@ impl QueryRouter {
                 )?
             };
 
-            self.create_df_logical_plan(stmt, &mut None, span_id.clone())
+            self.create_df_logical_plan(stmt, &mut None, span_id.clone(), compiler_id)
                 .await
         } else if name.eq_ignore_ascii_case("databases") || name.eq_ignore_ascii_case("schemas") {
             Ok(QueryPlan::MetaTabular(
@@ -289,7 +308,7 @@ impl QueryRouter {
                 &mut None,
             )?;
 
-            self.create_df_logical_plan(stmt, &mut None, span_id.clone())
+            self.create_df_logical_plan(stmt, &mut None, span_id.clone(), compiler_id)
                 .await
         } else if name.eq_ignore_ascii_case("warnings") {
             Ok(QueryPlan::MetaTabular(
@@ -322,6 +341,7 @@ impl QueryRouter {
                 },
                 &mut None,
                 span_id.clone(),
+                compiler_id,
             )
             .await
         }
@@ -693,8 +713,8 @@ impl QueryRouter {
         stmt: ast::Statement,
         qtrace: &mut Option<Qtrace>,
         span_id: Option<Arc<SpanId>>,
+        compiler_id: Uuid,
     ) -> CompilationResult<QueryPlan> {
-        self.reauthenticate_if_needed().await?;
         match &stmt {
             ast::Statement::Query(query) => match &query.body {
                 ast::SetExpr::Select(select) if select.into.is_some() => {
@@ -709,7 +729,14 @@ impl QueryRouter {
 
         let sql_query_engine = SqlQueryEngine::new(self.session_manager.clone());
         sql_query_engine
-            .plan(stmt, qtrace, span_id, self.meta.clone(), self.state.clone())
+            .plan(
+                stmt,
+                qtrace,
+                span_id,
+                self.meta.clone(),
+                self.state.clone(),
+                compiler_id,
+            )
             .await
     }
 }