diff --git a/crates/runtime/src/datafusion/execution.rs b/crates/runtime/src/datafusion/execution.rs index 7a16e775e..e49d1f2b7 100644 --- a/crates/runtime/src/datafusion/execution.rs +++ b/crates/runtime/src/datafusion/execution.rs @@ -87,6 +87,12 @@ impl SqlExecutor { // statement = self.update_statement_references(statement, warehouse_name); // query = statement.to_string(); + // TODO: Code should be organized in a better way + // 1. Single place to parse SQL strings into AST + // 2. Single place to update AST + // 3. Single place to construct Logical plan from this AST + // 4. Single place to rewrite-optimize-adjust logical plan + // etc if let DFStatement::Statement(s) = statement { match *s { Statement::CreateTable { .. } => { @@ -122,6 +128,7 @@ impl SqlExecutor { | Statement::Insert { .. } | Statement::ShowSchemas { .. } | Statement::ShowVariable { .. } + | Statement::ShowObjects { .. } | Statement::Update { .. } => { return Box::pin(self.execute_with_custom_plan(&query, warehouse_name)).await; } diff --git a/crates/runtime/src/datafusion/planner.rs b/crates/runtime/src/datafusion/planner.rs index db46ca319..5170e1b8a 100644 --- a/crates/runtime/src/datafusion/planner.rs +++ b/crates/runtime/src/datafusion/planner.rs @@ -62,13 +62,6 @@ where } } - pub fn statement_to_plan(&self, statement: DFStatement) -> Result { - match statement { - DFStatement::Statement(s) => self.sql_statement_to_plan(*s), - _ => self.inner.statement_to_plan(statement), - } - } - /// Custom implementation of `sql_statement_to_plan` pub fn sql_statement_to_plan(&self, statement: Statement) -> Result { // Check for a custom statement type @@ -85,6 +78,7 @@ where } /// Handle custom statements not supported by the original `SqlToRel` + #[allow(clippy::too_many_lines)] fn handle_custom_statement(&self, statement: Statement) -> Result { let planner_context: &mut PlannerContext = &mut PlannerContext::new(); // Example: Custom handling for a specific statement @@ -95,6 +89,18 @@ where | Statement::Update { .. } => Ok(LogicalPlan::default()), Statement::ShowSchemas { .. } => self.show_variable_to_plan(&["schemas".into()]), Statement::ShowVariable { variable } => self.show_variable_to_plan(&variable), + Statement::ShowObjects { + terse: _, + show_options, + } => { + let Some(show_in) = show_options.show_in else { + return plan_err!("Unsupported show statement: missing show_in"); + }; + let Some(parent_name) = show_in.parent_name else { + return plan_err!("Unsupported show statement: missing parent_name"); + }; + self.show_objects_to_plan(&parent_name) + } Statement::Drop { object_type, if_exists, @@ -232,6 +238,40 @@ where } } + fn show_objects_to_plan(&self, parent: &ObjectName) -> Result { + // Only support listing objects in schema for now + match parent.0.len() { + 2 => { + let (catalog, schema) = (parent.0[0].value.clone(), parent.0[1].value.clone()); + + // Create query to list objects in schema + let columns = [ + "table_catalog as 'database_name'", + "table_schema as 'schema_name'", + "table_name as 'name'", + "case when table_type='BASE TABLE' then 'TABLE' else table_type end as 'kind'", + "null as 'comment'", + ] + .join(", "); + // TODO: views? + // TODO: Return programmatically constructed plan + let query = format!("SELECT {columns} FROM information_schema.tables where table_schema = '{schema}' and table_catalog = '{catalog}'"); + let mut statements = DFParser::parse_sql(query.as_str())?; + statements.pop_front().map_or_else( + || plan_err!("Failed to parse SQL statement"), + |statement| { + if let DFStatement::Statement(s) = statement { + self.sql_statement_to_plan(*s) + } else { + plan_err!("Failed to parse SQL statement") + } + }, + ) + } + _ => plan_err!("Unsupported show objects: {:?}", parent), + } + } + fn show_variable_to_plan(&self, variable: &[Ident]) -> Result { //println!("SHOW variable: {:?}", variable); if !self.inner.has_table("information_schema", "df_settings") {