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

cube-js · Oct 9, 2024 · 81edfd1 · 81edfd1
1 parent 583e566
commit 81edfd1
Show file tree

Hide file tree

Showing 5 changed files with 180 additions and 108 deletions.
diff --git a/rust/cubesql/cubesql/src/compile/query_engine.rs b/rust/cubesql/cubesql/src/compile/query_engine.rs
@@ -1,5 +1,8 @@
 use crate::compile::engine::df::planner::CubeQueryPlanner;
-use std::{backtrace::Backtrace, collections::HashMap, future::Future, pin::Pin, sync::Arc};
+use std::{
+    backtrace::Backtrace, collections::HashMap, future::Future, pin::Pin, sync::Arc,
+    time::SystemTime,
+};
 
 use crate::{
     compile::{
@@ -43,6 +46,7 @@ use datafusion::{
     sql::{parser::Statement as DFStatement, planner::SqlToRel},
     variable::VarType,
 };
+use uuid::Uuid;
 
 #[async_trait::async_trait]
 pub trait QueryEngine {
@@ -74,6 +78,11 @@ pub trait QueryEngine {
 
     fn sanitize_statement(&self, stmt: &Self::AstStatementType) -> Self::AstStatementType;
 
+    async fn get_compiler_id_and_refresh_cache_if_needed(
+        &self,
+        state: Arc<SessionState>,
+    ) -> Result<Uuid, CompilationError>;
+
     async fn plan(
         &self,
         stmt: Self::AstStatementType,
@@ -82,6 +91,28 @@ pub trait QueryEngine {
         meta: Arc<MetaContext>,
         state: Arc<SessionState>,
     ) -> CompilationResult<QueryPlan> {
+        let compiler_id = self
+            .get_compiler_id_and_refresh_cache_if_needed(state.clone())
+            .await?;
+
+        let planning_start = SystemTime::now();
+        if let Some(span_id) = span_id.as_ref() {
+            if let Some(auth_context) = state.auth_context() {
+                self.transport_ref()
+                    .log_load_state(
+                        Some(span_id.clone()),
+                        auth_context,
+                        state.get_load_request_meta(),
+                        "SQL API Query Planning".to_string(),
+                        serde_json::json!({
+                            "query": span_id.query_key.clone(),
+                        }),
+                    )
+                    .await
+                    .map_err(|e| CompilationError::internal(e.to_string()))?;
+            }
+        }
+
         let ctx = self.create_session_ctx(state.clone())?;
         let cube_ctx = self.create_cube_ctx(state.clone(), meta.clone(), ctx.clone())?;
 
@@ -140,7 +171,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 +217,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(
@@ -233,12 +270,31 @@ pub trait QueryEngine {
                 // TODO: We should find what optimizers will be safety to use for OLAP queries
                 guard.optimizer.rules = vec![];
             }
-            if let Some(span_id) = span_id {
+            if let Some(span_id) = &span_id {
                 span_id.set_is_data_query(true).await;
             }
         };
 
         log::debug!("Rewrite: {:#?}", rewrite_plan);
+
+        if let Some(span_id) = span_id.as_ref() {
+            if let Some(auth_context) = state.auth_context() {
+                self.transport_ref()
+                    .log_load_state(
+                        Some(span_id.clone()),
+                        auth_context,
+                        state.get_load_request_meta(),
+                        "SQL API Query Planning Success".to_string(),
+                        serde_json::json!({
+                            "query": span_id.query_key.clone(),
+                            "duration": planning_start.elapsed().unwrap().as_millis() as u64,
+                        }),
+                    )
+                    .await
+                    .map_err(|e| CompilationError::internal(e.to_string()))?;
+            }
+        }
+
         let rewrite_plan = Self::evaluate_wrapped_sql(
             self.transport_ref().clone(),
             Arc::new(state.get_load_request_meta()),
@@ -493,6 +549,21 @@ impl QueryEngine for SqlQueryEngine {
     fn sanitize_statement(&self, stmt: &Self::AstStatementType) -> Self::AstStatementType {
         SensitiveDataSanitizer::new().replace(stmt.clone())
     }
+
+    async fn get_compiler_id_and_refresh_cache_if_needed(
+        &self,
+        state: Arc<SessionState>,
+    ) -> Result<Uuid, CompilationError> {
+        self.compiler_cache_ref()
+            .get_compiler_id_and_refresh_if_needed(
+                state.auth_context().ok_or_else(|| {
+                    CompilationError::internal("Unable to get auth context".to_string())
+                })?,
+                state.protocol.clone(),
+            )
+            .await
+            .map_err(|e| CompilationError::internal(e.to_string()))
+    }
 }
 
 fn is_olap_query(parent: &LogicalPlan) -> Result<bool, CompilationError> {

diff --git a/rust/cubesql/cubesql/src/compile/rewrite/rewriter.rs b/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,
             )

diff --git a/rust/cubesql/cubesql/src/compile/router.rs b/rust/cubesql/cubesql/src/compile/router.rs
@@ -3,7 +3,7 @@ use crate::compile::{
     StatusFlags,
 };
 use sqlparser::ast;
-use std::{collections::HashMap, future::Future, pin::Pin, sync::Arc, time::SystemTime};
+use std::{collections::HashMap, future::Future, pin::Pin, sync::Arc};
 
 use crate::{
     compile::{
@@ -61,50 +61,8 @@ impl QueryRouter {
         qtrace: &mut Option<Qtrace>,
         span_id: Option<Arc<SpanId>>,
     ) -> CompilationResult<QueryPlan> {
-        let planning_start = SystemTime::now();
-        if let Some(span_id) = span_id.as_ref() {
-            if let Some(auth_context) = self.state.auth_context() {
-                self.session_manager
-                    .server
-                    .transport
-                    .log_load_state(
-                        Some(span_id.clone()),
-                        auth_context,
-                        self.state.get_load_request_meta(),
-                        "SQL API Query Planning".to_string(),
-                        serde_json::json!({
-                            "query": span_id.query_key.clone(),
-                        }),
-                    )
-                    .await
-                    .map_err(|e| CompilationError::internal(e.to_string()))?;
-            }
-        }
-        let result = self
-            .create_df_logical_plan(stmt.clone(), qtrace, span_id.clone())
-            .await?;
-
-        if let Some(span_id) = span_id.as_ref() {
-            if let Some(auth_context) = self.state.auth_context() {
-                self.session_manager
-                    .server
-                    .transport
-                    .log_load_state(
-                        Some(span_id.clone()),
-                        auth_context,
-                        self.state.get_load_request_meta(),
-                        "SQL API Query Planning Success".to_string(),
-                        serde_json::json!({
-                            "query": span_id.query_key.clone(),
-                            "duration": planning_start.elapsed().unwrap().as_millis() as u64,
-                        }),
-                    )
-                    .await
-                    .map_err(|e| CompilationError::internal(e.to_string()))?;
-            }
-        }
-
-        return Ok(result);
+        self.create_df_logical_plan(stmt.clone(), qtrace, span_id.clone())
+            .await
     }
 
     pub async fn plan(