chore(cubesql): QueryEngine - support custom metadata for QueryPlan (#8817)

ovr · web-flow · commit 8be688c63552 · 2024-10-16T12:29:32.000+02:00
diff --git a/rust/cubesql/cubesql/src/compile/plan.rs b/rust/cubesql/cubesql/src/compile/plan.rs
@@ -45,15 +45,9 @@ pub enum QueryPlan {
     MetaOk(StatusFlags, CommandCompletion),
     MetaTabular(StatusFlags, Box<dataframe::DataFrame>),
     // Query will be executed via Data Fusion
-    DataFusionSelect(StatusFlags, LogicalPlan, DFSessionContext),
+    DataFusionSelect(LogicalPlan, DFSessionContext),
     // Query will be executed via DataFusion and saved to session
-    CreateTempTable(
-        StatusFlags,
-        LogicalPlan,
-        DFSessionContext,
-        String,
-        Arc<TempTableManager>,
-    ),
+    CreateTempTable(LogicalPlan, DFSessionContext, String, Arc<TempTableManager>),
 }
 
 impl fmt::Debug for QueryPlan {
@@ -70,16 +64,13 @@ impl fmt::Debug for QueryPlan {
                     flags
                 ))
             },
-            QueryPlan::DataFusionSelect(flags, _, _) => {
-                f.write_str(&format!(
-                    "DataFusionSelect(StatusFlags: {:?}, LogicalPlan: hidden, DFSessionContext: hidden)",
-                    flags
-                ))
+            QueryPlan::DataFusionSelect(_, _) => {
+                f.write_str(&"DataFusionSelect(LogicalPlan: hidden, DFSessionContext: hidden)")
             },
-            QueryPlan::CreateTempTable(flags, _, _, name, _) => {
+            QueryPlan::CreateTempTable(_, _, name, _) => {
                 f.write_str(&format!(
-                    "CreateTempTable(StatusFlags: {:?}, LogicalPlan: hidden, DFSessionContext: hidden, Name: {:?}, SessionState: hidden",
-                    flags, name
+                    "CreateTempTable(LogicalPlan: hidden, DFSessionContext: hidden, Name: {}, SessionState: hidden",
+                    name
                 ))
             },
         }
@@ -89,8 +80,9 @@ impl fmt::Debug for QueryPlan {
 impl QueryPlan {
     pub fn as_logical_plan(&self) -> LogicalPlan {
         match self {
-            QueryPlan::DataFusionSelect(_, plan, _)
-            | QueryPlan::CreateTempTable(_, plan, _, _, _) => plan.clone(),
+            QueryPlan::DataFusionSelect(plan, _) | QueryPlan::CreateTempTable(plan, _, _, _) => {
+                plan.clone()
+            }
             QueryPlan::MetaOk(_, _) | QueryPlan::MetaTabular(_, _) => {
                 panic!("This query doesnt have a plan, because it already has values for response")
             }
@@ -99,8 +91,8 @@ impl QueryPlan {
 
     pub async fn as_physical_plan(&self) -> Result<Arc<dyn ExecutionPlan>, CubeError> {
         match self {
-            QueryPlan::DataFusionSelect(_, plan, ctx)
-            | QueryPlan::CreateTempTable(_, plan, ctx, _, _) => {
+            QueryPlan::DataFusionSelect(plan, ctx)
+            | QueryPlan::CreateTempTable(plan, ctx, _, _) => {
                 DataFrame::new(ctx.state.clone(), plan)
                     .create_physical_plan()
                     .await
@@ -114,8 +106,7 @@ impl QueryPlan {
 
     pub fn print(&self, pretty: bool) -> Result<String, CubeError> {
         match self {
-            QueryPlan::DataFusionSelect(_, plan, _)
-            | QueryPlan::CreateTempTable(_, plan, _, _, _) => {
+            QueryPlan::DataFusionSelect(plan, _) | QueryPlan::CreateTempTable(plan, _, _, _) => {
                 if pretty {
                     Ok(plan.display_indent().to_string())
                 } else {
@@ -134,7 +125,7 @@ pub async fn get_df_batches(
     plan: &QueryPlan,
 ) -> Result<Pin<Box<dyn RecordBatchStream + Send>>, CubeError> {
     match plan {
-        QueryPlan::DataFusionSelect(_, plan, ctx) => {
+        QueryPlan::DataFusionSelect(plan, ctx) => {
             let df = DataFrame::new(ctx.state.clone(), &plan);
             let safe_stream = async move {
                 std::panic::AssertUnwindSafe(df.execute_stream())
diff --git a/rust/cubesql/cubesql/src/compile/query_engine.rs b/rust/cubesql/cubesql/src/compile/query_engine.rs
@@ -17,7 +17,7 @@ use crate::{
             analysis::LogicalPlanAnalysis,
             converter::{LogicalPlanToLanguageContext, LogicalPlanToLanguageConverter},
         },
-        CompilationError, CompilationResult, DatabaseProtocol, QueryPlan, Rewriter, StatusFlags,
+        CompilationError, CompilationResult, DatabaseProtocol, QueryPlan, Rewriter,
     },
     config::ConfigObj,
     sql::{
@@ -46,8 +46,12 @@ use datafusion::{
 
 #[async_trait::async_trait]
 pub trait QueryEngine {
+    /// Custom type for AST statement type, It allows to use any parsers for SQL
     type AstStatementType: std::fmt::Display + Send;
 
+    /// Additional metadata for results of plan method instead of extending query plan
+    type PlanMetadataType: std::fmt::Debug + Send;
+
     fn compiler_cache_ref(&self) -> &Arc<dyn CompilerCache>;
 
     fn transport_ref(&self) -> &Arc<dyn TransportService>;
@@ -70,7 +74,7 @@ pub trait QueryEngine {
         &self,
         cube_ctx: &CubeContext,
         stmt: &Self::AstStatementType,
-    ) -> Result<LogicalPlan, DataFusionError>;
+    ) -> Result<(LogicalPlan, Self::PlanMetadataType), DataFusionError>;
 
     fn sanitize_statement(&self, stmt: &Self::AstStatementType) -> Self::AstStatementType;
 
@@ -81,11 +85,11 @@ pub trait QueryEngine {
         span_id: Option<Arc<SpanId>>,
         meta: Arc<MetaContext>,
         state: Arc<SessionState>,
-    ) -> CompilationResult<QueryPlan> {
+    ) -> CompilationResult<(QueryPlan, Self::PlanMetadataType)> {
         let ctx = self.create_session_ctx(state.clone())?;
         let cube_ctx = self.create_cube_ctx(state.clone(), meta.clone(), ctx.clone())?;
 
-        let plan = self.create_logical_plan(&cube_ctx, &stmt).map_err(|err| {
+        let (plan, metadata) = self.create_logical_plan(&cube_ctx, &stmt).map_err(|err| {
             let message = format!("Initial planning error: {}", err,);
             let meta = Some(HashMap::from([
                 ("query".to_string(), stmt.to_string()),
@@ -249,11 +253,7 @@ pub trait QueryEngine {
             qtrace.set_best_plan_and_cube_scans(&rewrite_plan);
         }
 
-        Ok(QueryPlan::DataFusionSelect(
-            StatusFlags::empty(),
-            rewrite_plan,
-            ctx,
-        ))
+        Ok((QueryPlan::DataFusionSelect(rewrite_plan, ctx), metadata))
     }
 
     fn evaluate_wrapped_sql(
@@ -308,6 +308,8 @@ impl SqlQueryEngine {
 impl QueryEngine for SqlQueryEngine {
     type AstStatementType = sqlparser::ast::Statement;
 
+    type PlanMetadataType = ();
+
     fn create_cube_ctx(
         &self,
         state: Arc<SessionState>,
@@ -470,12 +472,12 @@ impl QueryEngine for SqlQueryEngine {
         &self,
         cube_ctx: &CubeContext,
         stmt: &Self::AstStatementType,
-    ) -> Result<LogicalPlan, DataFusionError> {
+    ) -> Result<(LogicalPlan, Self::PlanMetadataType), DataFusionError> {
         let df_query_planner = SqlToRel::new_with_options(cube_ctx, true);
         let plan =
-            df_query_planner.statement_to_plan(DFStatement::Statement(Box::new(stmt.clone())));
+            df_query_planner.statement_to_plan(DFStatement::Statement(Box::new(stmt.clone())))?;
 
-        plan
+        Ok((plan, ()))
     }
 
     fn compiler_cache_ref(&self) -> &Arc<dyn CompilerCache> {
diff --git a/rust/cubesql/cubesql/src/compile/router.rs b/rust/cubesql/cubesql/src/compile/router.rs
@@ -356,8 +356,8 @@ impl QueryRouter {
                         )])],
                     )),
                 )),
-                QueryPlan::DataFusionSelect(flags, plan, context)
-                | QueryPlan::CreateTempTable(flags, plan, context, _, _) => {
+                QueryPlan::DataFusionSelect(plan, context)
+                | QueryPlan::CreateTempTable(plan, context, _, _) => {
                     // EXPLAIN over CREATE TABLE AS shows the SELECT query plan
                     let plan = Arc::new(plan);
                     let schema = LogicalPlan::explain_schema();
@@ -385,7 +385,7 @@ impl QueryRouter {
                         })
                     };
 
-                    Ok(QueryPlan::DataFusionSelect(flags, explain_plan, context))
+                    Ok(QueryPlan::DataFusionSelect(explain_plan, context))
                 }
             }
         })
@@ -596,7 +596,7 @@ impl QueryRouter {
         span_id: Option<Arc<SpanId>>,
     ) -> Result<QueryPlan, CompilationError> {
         let plan = self.select_to_plan(stmt, qtrace, span_id).await?;
-        let QueryPlan::DataFusionSelect(flags, plan, ctx) = plan else {
+        let QueryPlan::DataFusionSelect(plan, ctx) = plan else {
             return Err(CompilationError::internal(
                 "unable to build DataFusion plan from Query".to_string(),
             ));
@@ -608,8 +608,8 @@ impl QueryRouter {
                 "table name contains no ident parts".to_string(),
             ));
         };
+
         Ok(QueryPlan::CreateTempTable(
-            flags,
             plan,
             ctx,
             table_name.value.to_string(),
@@ -708,9 +708,11 @@ impl QueryRouter {
         }
 
         let sql_query_engine = SqlQueryEngine::new(self.session_manager.clone());
-        sql_query_engine
+        let (plan, _) = sql_query_engine
             .plan(stmt, qtrace, span_id, self.meta.clone(), self.state.clone())
-            .await
+            .await?;
+
+        Ok(plan)
     }
 }
 
diff --git a/rust/cubesql/cubesql/src/compile/test/mod.rs b/rust/cubesql/cubesql/src/compile/test/mod.rs
@@ -933,19 +933,21 @@ impl TestContext {
                 .await
                 .map_err(|e| CubeError::internal(format!("Error during planning: {}", e)))?;
             match query {
-                QueryPlan::DataFusionSelect(flags, plan, ctx) => {
+                QueryPlan::DataFusionSelect(plan, ctx) => {
                     let df = DFDataFrame::new(ctx.state, &plan);
                     let batches = df.collect().await?;
                     let frame = batches_to_dataframe(&df.schema().into(), batches)?;
 
                     output.push(frame.print());
-                    output_flags = flags;
                 }
                 QueryPlan::MetaTabular(flags, frame) => {
                     output.push(frame.print());
                     output_flags = flags;
                 }
-                QueryPlan::MetaOk(flags, _) | QueryPlan::CreateTempTable(flags, _, _, _, _) => {
+                QueryPlan::CreateTempTable(_, _, _, _) => {
+                    // nothing to do
+                }
+                QueryPlan::MetaOk(flags, _) => {
                     output_flags = flags;
                 }
             }
diff --git a/rust/cubesql/cubesql/src/sql/postgres/extended.rs b/rust/cubesql/cubesql/src/sql/postgres/extended.rs
@@ -492,7 +492,7 @@ impl Portal {
 
                             return;
                         }
-                        QueryPlan::DataFusionSelect(_, plan, ctx) => {
+                        QueryPlan::DataFusionSelect(plan, ctx) => {
                             let df = DFDataFrame::new(ctx.state.clone(), &plan);
                             let safe_stream = async move {
                                 std::panic::AssertUnwindSafe(df.execute_stream())
@@ -511,7 +511,7 @@ impl Portal {
                                 Err(err) => return yield Err(CubeError::panic(err).into()),
                             }
                         }
-                        QueryPlan::CreateTempTable(_, plan, ctx, name, temp_tables) => {
+                        QueryPlan::CreateTempTable(plan, ctx, name, temp_tables) => {
                             let df = DFDataFrame::new(ctx.state.clone(), &plan);
                             let record_batch = df.collect();
                             let row_count = match record_batch.await {
diff --git a/rust/cubesql/cubesql/src/sql/postgres/shim.rs b/rust/cubesql/cubesql/src/sql/postgres/shim.rs
@@ -72,7 +72,7 @@ impl QueryPlanExt for QueryPlan {
         required_format: protocol::Format,
     ) -> Result<Option<protocol::RowDescription>, ConnectionError> {
         match &self {
-            QueryPlan::MetaOk(_, _) | QueryPlan::CreateTempTable(_, _, _, _, _) => Ok(None),
+            QueryPlan::MetaOk(_, _) | QueryPlan::CreateTempTable(_, _, _, _) => Ok(None),
             QueryPlan::MetaTabular(_, frame) => {
                 let mut result = vec![];
 
@@ -86,7 +86,7 @@ impl QueryPlanExt for QueryPlan {
 
                 Ok(Some(protocol::RowDescription::new(result)))
             }
-            QueryPlan::DataFusionSelect(_, logical_plan, _) => {
+            QueryPlan::DataFusionSelect(logical_plan, _) => {
                 let mut result = vec![];
 
                 for field in logical_plan.schema().fields() {

Original file line number	Diff line number	Diff line change
`@@ -933,19 +933,21 @@ impl TestContext {`
`933`	`933`	`.await`
`934`	`934`	`.map_err(\|e\| CubeError::internal(format!("Error during planning: {}", e)))?;`
`935`	`935`	`match query {`
`936`		`- QueryPlan::DataFusionSelect(flags, plan, ctx) => {`
	`936`	`+ QueryPlan::DataFusionSelect(plan, ctx) => {`
`937`	`937`	`let df = DFDataFrame::new(ctx.state, &plan);`
`938`	`938`	`let batches = df.collect().await?;`
`939`	`939`	`let frame = batches_to_dataframe(&df.schema().into(), batches)?;`
`940`	`940`
`941`	`941`	`output.push(frame.print());`
`942`		`- output_flags = flags;`
`943`	`942`	`}`
`944`	`943`	`QueryPlan::MetaTabular(flags, frame) => {`
`945`	`944`	`output.push(frame.print());`
`946`	`945`	`output_flags = flags;`
`947`	`946`	`}`
`948`		`- QueryPlan::MetaOk(flags, _) \| QueryPlan::CreateTempTable(flags, _, _, _, _) => {`
	`947`	`+ QueryPlan::CreateTempTable(_, _, _, _) => {`
	`948`	`+ // nothing to do`
	`949`	`+ }`
	`950`	`+ QueryPlan::MetaOk(flags, _) => {`
`949`	`951`	`output_flags = flags;`
`950`	`952`	`}`
`951`	`953`	`}`
Original file line number	Diff line number	Diff line change
`@@ -492,7 +492,7 @@ impl Portal {`
`492`	`492`
`493`	`493`	`return;`
`494`	`494`	`}`
`495`		`- QueryPlan::DataFusionSelect(_, plan, ctx) => {`
	`495`	`+ QueryPlan::DataFusionSelect(plan, ctx) => {`
`496`	`496`	`let df = DFDataFrame::new(ctx.state.clone(), &plan);`
`497`	`497`	`let safe_stream = async move {`
`498`	`498`	`std::panic::AssertUnwindSafe(df.execute_stream())`
`@@ -511,7 +511,7 @@ impl Portal {`
`511`	`511`	`Err(err) => return yield Err(CubeError::panic(err).into()),`
`512`	`512`	`}`
`513`	`513`	`}`
`514`		`- QueryPlan::CreateTempTable(_, plan, ctx, name, temp_tables) => {`
	`514`	`+ QueryPlan::CreateTempTable(plan, ctx, name, temp_tables) => {`
`515`	`515`	`let df = DFDataFrame::new(ctx.state.clone(), &plan);`
`516`	`516`	`let record_batch = df.collect();`
`517`	`517`	`let row_count = match record_batch.await {`