diff --git a/Cargo.lock b/Cargo.lock index 1b2a44d0e4e..073ae035256 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7301,7 +7301,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "07033963ba89ebaf1584d767badaa2e8fcec21aedea6b8c0346d487d49c28667" dependencies = [ "cfg-if", - "windows-targets 0.48.5", + "windows-targets 0.52.6", ] [[package]] @@ -9620,9 +9620,9 @@ dependencies = [ [[package]] name = "pgwire" -version = "0.38.0" +version = "0.38.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "89d5e5a60d3f6e40c91f6a2a7f8d09665e636272bd5611977253559b6651aabb" +checksum = "f2a798d130b8975a566c2cf6d8955746e1f09a9ee2c3ff2e6020a2c6528c5bd1" dependencies = [ "async-trait", "base64 0.22.1", diff --git a/src/catalog/src/system_schema/information_schema/region_peers.rs b/src/catalog/src/system_schema/information_schema/region_peers.rs index 5bc91d207ea..b1438ef53d4 100644 --- a/src/catalog/src/system_schema/information_schema/region_peers.rs +++ b/src/catalog/src/system_schema/information_schema/region_peers.rs @@ -267,7 +267,7 @@ impl InformationSchemaRegionPeersBuilder { ]; if !predicates.eval(&row) { - return; + continue; } self.table_catalogs.push(Some(table_catalog)); diff --git a/src/query/src/datafusion.rs b/src/query/src/datafusion.rs index dc84c4afac6..e2e577debf7 100644 --- a/src/query/src/datafusion.rs +++ b/src/query/src/datafusion.rs @@ -354,25 +354,6 @@ impl DatafusionQueryEngine { Ok(physical_plan) } - #[tracing::instrument(skip_all)] - pub fn optimize( - &self, - context: &QueryEngineContext, - plan: &LogicalPlan, - ) -> Result { - let _timer = metrics::OPTIMIZE_LOGICAL_ELAPSED.start_timer(); - - // Optimized by extension rules - let optimized_plan = self - .state - .optimize_by_extension_rules(plan.clone(), context)?; - - // Optimized by datafusion optimizer - let optimized_plan = self.state.session_state().optimize(&optimized_plan)?; - - Ok(optimized_plan) - } - #[tracing::instrument(skip_all)] fn optimize_physical_plan( &self, @@ -444,32 +425,17 @@ impl QueryEngine for DatafusionQueryEngine { async fn describe( &self, plan: LogicalPlan, - query_ctx: QueryContextRef, + _query_ctx: QueryContextRef, ) -> Result { - let ctx = self.engine_context(query_ctx); - if let Ok(optimised_plan) = self.optimize(&ctx, &plan) { - let schema = optimised_plan - .schema() - .clone() - .try_into() - .context(ConvertSchemaSnafu)?; - Ok(DescribeResult { - schema, - logical_plan: optimised_plan, - }) - } else { - // Table's like those in information_schema cannot be optimized when - // it contains parameters. So we fallback to original plans. - let schema = plan - .schema() - .clone() - .try_into() - .context(ConvertSchemaSnafu)?; - Ok(DescribeResult { - schema, - logical_plan: plan, - }) - } + let schema = plan + .schema() + .clone() + .try_into() + .context(ConvertSchemaSnafu)?; + Ok(DescribeResult { + schema, + logical_plan: plan, + }) } async fn execute(&self, plan: LogicalPlan, query_ctx: QueryContextRef) -> Result { @@ -924,7 +890,7 @@ mod tests { ) ); assert_eq!( - "Limit: skip=0, fetch=20\n Aggregate: groupBy=[[]], aggr=[[sum(CAST(numbers.number AS UInt64))]]\n TableScan: numbers projection=[number]", + "Limit: skip=0, fetch=20\n Projection: sum(numbers.number)\n Aggregate: groupBy=[[]], aggr=[[sum(numbers.number)]]\n TableScan: numbers", format!("{}", logical_plan.display_indent()) ); } diff --git a/src/query/src/planner.rs b/src/query/src/planner.rs index 44c9bc39561..f522dc567a7 100644 --- a/src/query/src/planner.rs +++ b/src/query/src/planner.rs @@ -28,6 +28,7 @@ use datafusion::execution::context::SessionState; use datafusion::sql::planner::PlannerContext; use datafusion_common::ToDFSchema; use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; +use datafusion_expr::expr::{Exists, InSubquery}; use datafusion_expr::{ Analyze, Explain, ExplainFormat, Expr as DfExpr, LogicalPlan, LogicalPlanBuilder, PlanType, ToStringifiedPlan, col, @@ -424,9 +425,20 @@ impl DfLogicalPlanner { let mut placeholder_types = HashMap::new(); let mut casted_placeholders = HashSet::new(); + Self::extract_from_plan(plan, &mut placeholder_types, &mut casted_placeholders)?; + + Ok(placeholder_types) + } + + fn extract_from_plan( + plan: &LogicalPlan, + placeholder_types: &mut HashMap>, + casted_placeholders: &mut HashSet, + ) -> Result<()> { plan.apply(|node| { for expr in node.expressions() { let _ = expr.apply(|e| { + // Handle casted placeholders if let DfExpr::Cast(cast) = e && let DfExpr::Placeholder(ph) = &*cast.expr { @@ -434,6 +446,7 @@ impl DfLogicalPlanner { casted_placeholders.insert(ph.id.clone()); } + // Handle bare (non-casted) placeholders if let DfExpr::Placeholder(ph) = e && !casted_placeholders.contains(&ph.id) && !placeholder_types.contains_key(&ph.id) @@ -441,13 +454,26 @@ impl DfLogicalPlanner { placeholder_types.insert(ph.id.clone(), None); } + // Recurse into subquery plans embedded in expressions + match e { + DfExpr::Exists(Exists { subquery, .. }) + | DfExpr::InSubquery(InSubquery { subquery, .. }) + | DfExpr::ScalarSubquery(subquery) => { + Self::extract_from_plan( + &subquery.subquery, + placeholder_types, + casted_placeholders, + )?; + } + _ => {} + } + Ok(TreeNodeRecursion::Continue) }); } Ok(TreeNodeRecursion::Continue) })?; - - Ok(placeholder_types) + Ok(()) } /// Gets inferred parameter types from a logical plan. @@ -619,4 +645,16 @@ mod tests { assert_eq!(type_2, &Some(DataType::Utf8)); assert_eq!(type_3, &Some(DataType::Int32)); } + + #[tokio::test] + async fn test_get_inferred_parameter_types_subquery() { + let plan = parse_sql_to_plan( + r#"SELECT * FROM test WHERE id = (SELECT id FROM test CROSS JOIN (SELECT parse_ident($1::TEXT) AS parts) p LIMIT 1)"#, + ).await; + let types = DfLogicalPlanner::get_inferred_parameter_types(&plan).unwrap(); + + assert_eq!(types.len(), 1); + let type_1 = types.get("$1").unwrap(); + assert_eq!(type_1, &Some(DataType::Utf8)); + } } diff --git a/src/servers/Cargo.toml b/src/servers/Cargo.toml index e75192c9bab..8b64a256e72 100644 --- a/src/servers/Cargo.toml +++ b/src/servers/Cargo.toml @@ -89,7 +89,7 @@ operator.workspace = true otel-arrow-rust.workspace = true parking_lot.workspace = true pg_interval = { version = "0.5.2", package = "pg_interval_2" } -pgwire = { version = "0.38", default-features = false, features = [ +pgwire = { version = "0.38.1", default-features = false, features = [ "server-api-ring", "pg-ext-types", ] } diff --git a/tests-fuzz/src/utils/partition.rs b/tests-fuzz/src/utils/partition.rs index d3dc30061dc..89a684326b4 100644 --- a/tests-fuzz/src/utils/partition.rs +++ b/tests-fuzz/src/utils/partition.rs @@ -36,7 +36,7 @@ pub struct PartitionCount { } pub async fn count_partitions(db: &MySqlPool, datanode_id: u64) -> Result { - let sql = "select count(1) as count from information_schema.region_peers where peer_id == ?"; + let sql = "select count(1) as count from information_schema.region_peers where peer_id = ?"; sqlx::query_as::<_, PartitionCount>(sql) .bind(datanode_id) .fetch_one(db) diff --git a/tests-fuzz/targets/migration/fuzz_migrate_mito_regions.rs b/tests-fuzz/targets/migration/fuzz_migrate_mito_regions.rs index c8ebbb54afa..17cbfb9251d 100644 --- a/tests-fuzz/targets/migration/fuzz_migrate_mito_regions.rs +++ b/tests-fuzz/targets/migration/fuzz_migrate_mito_regions.rs @@ -261,13 +261,18 @@ async fn migrate_regions(ctx: &FuzzContext, migrations: &[Migration]) -> Result< { let output = procedure_state(&greptime, &procedure_id).await; info!("Checking procedure: {procedure_id}, output: {output}"); - (fetch_partition(&greptime, region_id).await.unwrap(), output) + (fetch_partition(&greptime, region_id).await.ok(), output) } }) }, |(partition, output)| { - info!("Region: {region_id}, datanode: {}", partition.datanode_id); - partition.datanode_id == migration.to_peer && output.contains("Done") + if let Some(partition) = partition { + info!("Region: {region_id}, datanode: {}", partition.datanode_id); + partition.datanode_id == migration.to_peer && output.contains("Done") + } else { + info!("Region: {region_id}, partition not found yet"); + false + } }, Duration::from_secs(5), ) diff --git a/tests/cases/standalone/common/prepare/mysql_prepare.result b/tests/cases/standalone/common/prepare/mysql_prepare.result index abc267b50e4..5ef242a891b 100644 --- a/tests/cases/standalone/common/prepare/mysql_prepare.result +++ b/tests/cases/standalone/common/prepare/mysql_prepare.result @@ -42,7 +42,7 @@ affected_rows: 0 -- SQLNESS PROTOCOL MYSQL EXECUTE stmt USING 'a'; -Failed to execute query, err: MySqlError { ERROR 1815 (HY000): (EngineExecuteQuery): Cast error: Cannot cast string 'a' to value of Int32 type } +Failed to execute query, err: MySqlError { ERROR 1210 (HY000): (InvalidArguments): Invalid request parameter: Unable to convert a to datatype Int32(Int32Type) } -- SQLNESS PROTOCOL MYSQL DEALLOCATE stmt; @@ -124,6 +124,25 @@ DEALLOCATE stmt; affected_rows: 0 +-- SQLNESS PROTOCOL MYSQL +PREPARE stmt FROM 'SELECT table_name, table_schema FROM information_schema.tables WHERE table_name = ?'; + +affected_rows: 0 + +-- SQLNESS PROTOCOL MYSQL +EXECUTE stmt USING 'cake'; + ++------------+--------------+ +| table_name | table_schema | ++------------+--------------+ +| cake | public | ++------------+--------------+ + +-- SQLNESS PROTOCOL MYSQL +DEALLOCATE stmt; + +affected_rows: 0 + -- SQLNESS PROTOCOL MYSQL DROP TABLE cake; diff --git a/tests/cases/standalone/common/prepare/mysql_prepare.sql b/tests/cases/standalone/common/prepare/mysql_prepare.sql index 8e80a0a867a..e96e945f881 100644 --- a/tests/cases/standalone/common/prepare/mysql_prepare.sql +++ b/tests/cases/standalone/common/prepare/mysql_prepare.sql @@ -72,5 +72,14 @@ EXECUTE stmt USING 'happy', 42, 0; -- SQLNESS PROTOCOL MYSQL DEALLOCATE stmt; +-- SQLNESS PROTOCOL MYSQL +PREPARE stmt FROM 'SELECT table_name, table_schema FROM information_schema.tables WHERE table_name = ?'; + +-- SQLNESS PROTOCOL MYSQL +EXECUTE stmt USING 'cake'; + +-- SQLNESS PROTOCOL MYSQL +DEALLOCATE stmt; + -- SQLNESS PROTOCOL MYSQL DROP TABLE cake;