From 297a343fd45f07ea267eb88d79eb592f5a912240 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Mon, 9 Nov 2026 11:08:47 +0000 Subject: [PATCH] fix(query): use current partition columns in prepared plans (#10244) A prepared statement keeps its logical plan, and that plan holds the table as it was at prepare time. The analyzer reads the partition columns from that table while the regions to scan come from the current route. After a table's partition columns change, the analyzer can consider a query safe to push down, so each region runs the whole query and the frontend concatenates their results. The same happens to the source of a prepared INSERT ... SELECT, which then stores wrong values. Before executing a kept plan, compare each scanned table's partition columns with the current table found by id. Scans of a table whose columns changed get the current table when the schema is unchanged. Otherwise they keep their schema, because a scan's projection holds indices into it, and report the current columns as absent from it, so the query keeps merging on the frontend. Signed-off-by: jeremyhi --- src/frontend/src/instance.rs | 5 + src/query/src/plan.rs | 260 +++++++++++++++++++-- src/table/src/table.rs | 8 +- tests-integration/src/test_util.rs | 27 ++- tests-integration/tests/repartition.rs | 300 ++++++++++++++++++++++++- 5 files changed, 578 insertions(+), 22 deletions(-) diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index 4580b611352..e48642f3d5a 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -1109,6 +1109,11 @@ impl SqlQueryHandler for Instance { stmt: Option, query_ctx: QueryContextRef, ) -> server_error::Result { + // Prepared statements reach here with a plan built before their current execution. + let plan = query::plan::refresh_partition_columns(plan, &self.catalog_manager, &query_ctx) + .await + .map_err(BoxedError::new) + .context(server_error::ExecutePlanSnafu)?; self.do_exec_plan_inner(plan, stmt, query_ctx) .await .map_err(BoxedError::new) diff --git a/src/query/src/plan.rs b/src/query/src/plan.rs index 26635a9d990..d69cda86da9 100644 --- a/src/query/src/plan.rs +++ b/src/query/src/plan.rs @@ -12,18 +12,25 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; -use datafusion::datasource::DefaultTableSource; +use catalog::CatalogManagerRef; +use datafusion::datasource::{DefaultTableSource, provider_as_source}; use datafusion_common::TableReference; -use datafusion_common::tree_node::{Transformed, TreeNodeRewriter}; +use datafusion_common::tree_node::{Transformed, TreeNodeRecursion, TreeNodeRewriter}; +use datafusion_expr::logical_plan::TableScan; use datafusion_expr::{Expr, LogicalPlan}; -use session::context::QueryContextRef; +use session::context::{QueryContext, QueryContextRef}; +use snafu::ResultExt; +use table::TableRef; +use table::dist_table::DistTable; pub use table::metadata::TableType; +use table::table::PartitionRules; use table::table::adapter::DfTableProviderAdapter; use table::table_name::TableName; -use crate::error::Result; +use crate::error::{CatalogSnafu, Result}; struct TableNamesExtractAndRewriter { pub(crate) table_names: HashSet, @@ -40,13 +47,8 @@ impl TreeNodeRewriter for TableNamesExtractAndRewriter { ) -> datafusion::error::Result> { match node { LogicalPlan::TableScan(mut scan) => { - if let Some(source) = scan.source.downcast_ref::() - && let Some(provider) = source - .table_provider - .downcast_ref::() - && provider.table().table_type() == TableType::Base - { - let info = provider.table().table_info(); + if let Some(table) = scanned_base_table(&scan) { + let info = table.table_info(); self.table_names.insert(TableName::new( info.catalog_name.clone(), info.schema_name.clone(), @@ -119,6 +121,128 @@ pub fn extract_and_rewrite_full_table_names( Ok((extractor.table_names, plan.data)) } +/// Updates the scans in `plan` whose table changed its partition columns after the plan was built. +/// +/// Distributed planning reads the partition columns from the table a scan holds, and the regions +/// from the current route. A plan kept across executions, such as a prepared statement's, must +/// therefore report the current partition columns, or a query that needs merging on the frontend +/// runs entirely on each region. +/// +/// Returns `plan` as it is when no scanned table changed its partition columns. +pub async fn refresh_partition_columns( + plan: LogicalPlan, + catalog_manager: &CatalogManagerRef, + query_ctx: &QueryContext, +) -> Result { + let mut scanned = HashMap::new(); + plan.apply_with_subqueries(|node| { + if let LogicalPlan::TableScan(scan) = node + && let Some(table) = scanned_base_table(scan) + { + scanned.insert(table.table_info().table_id(), table); + } + Ok(TreeNodeRecursion::Continue) + })?; + + let mut refreshed = HashMap::new(); + for (table_id, table) in scanned { + // Looks up by id to follow a rename and to not bind a table that replaced the scanned one. + let Some(info) = catalog_manager + .table_info_by_id(table_id) + .await + .context(CatalogSnafu)? + else { + continue; + }; + let current = catalog_manager + .table( + &info.catalog_name, + &info.schema_name, + &info.name, + Some(query_ctx), + ) + .await + .context(CatalogSnafu)?; + let Some(current) = current.filter(|t| t.table_info().table_id() == table_id) else { + continue; + }; + if has_same_partition_columns(&table, ¤t) { + continue; + } + let table = if table.schema() == current.schema() { + current + } else { + with_unmapped_partition_columns(&table, ¤t) + }; + refreshed.insert(table_id, table); + } + if refreshed.is_empty() { + return Ok(plan); + } + + let plan = plan.transform_with_subqueries(|node| { + let LogicalPlan::TableScan(mut scan) = node else { + return Ok(Transformed::no(node)); + }; + let Some(table) = scanned_base_table(&scan) + .and_then(|table| refreshed.get(&table.table_info().table_id())) + else { + return Ok(Transformed::no(LogicalPlan::TableScan(scan))); + }; + scan.source = provider_as_source(Arc::new(DfTableProviderAdapter::new(table.clone()))); + Ok(Transformed::yes(LogicalPlan::TableScan(scan))) + })?; + Ok(plan.data) +} + +fn scanned_base_table(scan: &TableScan) -> Option { + let source = scan.source.downcast_ref::()?; + let provider = source + .table_provider + .downcast_ref::()?; + let table = provider.table(); + (table.table_type() == TableType::Base).then_some(table) +} + +fn has_same_partition_columns(left: &TableRef, right: &TableRef) -> bool { + left.table_info() + .meta + .partition_column_names() + .eq(right.table_info().meta.partition_column_names()) + && extra_partition_columns(left) == extra_partition_columns(right) +} + +fn extra_partition_columns(table: &TableRef) -> &[String] { + table + .partition_rules() + .map(|rules| rules.extra_phy_cols_not_in_logical_table.as_slice()) + .unwrap_or_default() +} + +/// Builds a table that keeps the schema of `table` and reports the partition columns of `current` +/// as absent from that schema. +/// +/// A scan's projection holds indices into the schema it was built on, and Substrait sends that +/// schema with those indices, so a plan cannot move to a changed schema. Reporting the columns as +/// absent keeps distributed planning from relying on any of them. +fn with_unmapped_partition_columns(table: &TableRef, current: &TableRef) -> TableRef { + let mut info = (*table.table_info()).clone(); + info.meta.partition_key_indices.clear(); + let current_info = current.table_info(); + let extra_phy_cols_not_in_logical_table = current_info + .meta + .partition_column_names() + .chain(extra_partition_columns(current)) + .cloned() + .collect(); + DistTable::table_partitioned( + Arc::new(info), + Some(PartitionRules { + extra_phy_cols_not_in_logical_table, + }), + ) +} + /// A trait to extract expressions from a logical plan. pub trait ExtractExpr { /// Gets expressions from a logical plan. @@ -139,13 +263,20 @@ pub(crate) mod tests { use std::sync::Arc; use arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit}; - use common_catalog::consts::DEFAULT_CATALOG_NAME; + use catalog::RegisterTableRequest; + use catalog::memory::MemoryCatalogManager; + use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; use datafusion::logical_expr::builder::LogicalTableSource; use datafusion::logical_expr::{LogicalPlan, LogicalPlanBuilder, col, lit, scalar_subquery}; + use datatypes::prelude::ConcreteDataType; + use datatypes::schema::{ColumnSchema, Schema as TableSchema}; use session::context::QueryContextBuilder; + use table::test_util::table_info::test_table_info; use super::*; + const TABLE_ID: u32 = 1024; + fn mock_table_source() -> Arc { let schema = Schema::new(vec![ Field::new("id", DataType::Int32, true), @@ -307,4 +438,107 @@ pub(crate) mod tests { TableReference::full("qp031_external_catalog", "qp031_external_schema", "lookup"), ); } + + fn partitioned_table( + columns: &[&str], + partition_key_indices: Vec, + extras: &[&str], + ) -> TableRef { + let columns = columns + .iter() + .map(|name| ColumnSchema::new(*name, ConcreteDataType::string_datatype(), true)) + .collect(); + let mut info = test_table_info( + TABLE_ID, + "t", + DEFAULT_SCHEMA_NAME, + DEFAULT_CATALOG_NAME, + Arc::new(TableSchema::new(columns)), + ); + info.meta.partition_key_indices = partition_key_indices; + let rules = PartitionRules { + extra_phy_cols_not_in_logical_table: extras.iter().map(ToString::to_string).collect(), + }; + DistTable::table_partitioned(Arc::new(info), Some(rules)) + } + + /// Returns the plan that [refresh_partition_columns] builds for a scan of `captured` when + /// the catalog holds `current`. + async fn refresh_scan(captured: TableRef, current: TableRef) -> (TableScan, TableScan) { + let source = provider_as_source(Arc::new(DfTableProviderAdapter::new(captured))); + let filters = vec![col("a").gt(lit("x"))]; + let plan = LogicalPlanBuilder::scan_with_filters_fetch( + "t", + source, + Some(vec![1]), + filters, + Some(5), + ) + .unwrap() + .build() + .unwrap(); + let catalog_manager = MemoryCatalogManager::with_default_setup(); + catalog_manager + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: "t".to_string(), + table_id: TABLE_ID, + table: current, + }) + .unwrap(); + let catalog_manager: CatalogManagerRef = catalog_manager; + let refreshed = + refresh_partition_columns(plan.clone(), &catalog_manager, &QueryContext::arc()) + .await + .unwrap(); + let (LogicalPlan::TableScan(original), LogicalPlan::TableScan(refreshed)) = + (plan, refreshed) + else { + panic!("expected table scans"); + }; + (original, refreshed) + } + + #[tokio::test] + async fn test_refresh_partition_columns_compares_names() { + let captured = partitioned_table(&["a", "b"], vec![1], &[]); + let current = partitioned_table(&["b", "a", "c"], vec![0], &[]); + + let (original, refreshed) = refresh_scan(captured, current).await; + + assert!(Arc::ptr_eq(&original.source, &refreshed.source)); + } + + #[tokio::test] + async fn test_refresh_partition_columns_with_equal_schema() { + let captured = partitioned_table(&["a", "b"], vec![], &[]); + let current = partitioned_table(&["a", "b"], vec![1], &[]); + + let (original, refreshed) = refresh_scan(captured, current.clone()).await; + + let table = scanned_base_table(&refreshed).unwrap(); + assert!(Arc::ptr_eq(&table, ¤t)); + assert_eq!(original.projection, refreshed.projection); + assert_eq!(original.projected_schema, refreshed.projected_schema); + assert_eq!(original.filters, refreshed.filters); + assert_eq!(original.fetch, refreshed.fetch); + } + + #[tokio::test] + async fn test_refresh_partition_columns_with_changed_schema() { + let captured = partitioned_table(&["a", "b"], vec![1], &[]); + let current = partitioned_table(&["a", "b", "c"], vec![2], &["d"]); + + let (original, refreshed) = refresh_scan(captured.clone(), current).await; + + let table = scanned_base_table(&refreshed).unwrap(); + assert_eq!(table.schema(), captured.schema()); + assert!(table.table_info().meta.partition_key_indices.is_empty()); + assert_eq!(extra_partition_columns(&table), ["c", "d"]); + assert_eq!(original.projection, refreshed.projection); + assert_eq!(original.projected_schema, refreshed.projected_schema); + assert_eq!(original.filters, refreshed.filters); + assert_eq!(original.fetch, refreshed.fetch); + } } diff --git a/src/table/src/table.rs b/src/table/src/table.rs index cdc06a50be2..7b126b3fdfb 100644 --- a/src/table/src/table.rs +++ b/src/table/src/table.rs @@ -56,10 +56,10 @@ lazy_static! { /// Defines partition rules for a table. /// TODO(discord9): add the entire partition exprs rules here later pub struct PartitionRules { - /// The physical partition columns that are not in the logical table. - /// only used in kvbackend manager to store the physical partition columns that are not in the logical table. - /// This is used to avoid the partition columns in the physical table that are not in the logical table - /// to prevent certain optimizations, if table is not a logical table, this should be empty + /// The partition columns that are not in this table's schema: physical partition columns + /// missing from a logical table, or the current partition columns of a table that a plan + /// keeps with an older schema. They prevent optimizations that rely on knowing the partition + /// columns. pub extra_phy_cols_not_in_logical_table: Vec, } diff --git a/tests-integration/src/test_util.rs b/tests-integration/src/test_util.rs index 6da25d659b6..33f937ab484 100644 --- a/tests-integration/src/test_util.rs +++ b/tests-integration/src/test_util.rs @@ -1133,7 +1133,18 @@ pub async fn setup_mysql_server_with_user_provider( } let instance = setup_standalone_instance_with_plugins(name, store_type, plugins).await; + let mysql_server = + setup_mysql_server_for_frontend_instance(instance.fe_instance().clone(), user_provider) + .await; + (instance.guard, Arc::new(mysql_server)) +} + +/// Builds and starts a MySQL server on an existing frontend instance. +pub async fn setup_mysql_server_for_frontend_instance( + fe_instance_ref: Arc, + user_provider: Option, +) -> Box { let runtime = RuntimeBuilder::default() .worker_threads(2) .thread_name("mysql-runtime") @@ -1142,7 +1153,6 @@ pub async fn setup_mysql_server_with_user_provider( let fe_mysql_addr = format!("127.0.0.1:{}", ports::get_port()); - let fe_instance_ref = instance.fe_instance().clone(); let opts = MysqlOptions { addr: fe_mysql_addr.clone(), ..Default::default() @@ -1168,7 +1178,7 @@ pub async fn setup_mysql_server_with_user_provider( .await .unwrap(); - (instance.guard, Arc::new(mysql_server)) + mysql_server } pub async fn setup_pg_server( @@ -1234,7 +1244,17 @@ pub async fn setup_pg_server_with_user_provider( user_provider: Option, ) -> (TestGuard, Arc>) { let instance = setup_standalone_instance(name, store_type).await; + let pg_server = + setup_pg_server_for_frontend_instance(instance.fe_instance().clone(), user_provider).await; + (instance.guard, Arc::new(pg_server)) +} + +/// Builds and starts a Postgres server on an existing frontend instance. +pub async fn setup_pg_server_for_frontend_instance( + fe_instance_ref: Arc, + user_provider: Option, +) -> Box { let runtime = RuntimeBuilder::default() .worker_threads(2) .thread_name("pg-runtime") @@ -1243,7 +1263,6 @@ pub async fn setup_pg_server_with_user_provider( let fe_pg_addr = format!("127.0.0.1:{}", ports::get_port()); - let fe_instance_ref = instance.fe_instance().clone(); let opts = PostgresOptions { addr: fe_pg_addr.clone(), ..Default::default() @@ -1273,7 +1292,7 @@ pub async fn setup_pg_server_with_user_provider( .await .unwrap(); - (instance.guard, Arc::new(pg_server)) + pg_server } /// Sets up a standalone instance with both a Prometheus remote-write HTTP app diff --git a/tests-integration/tests/repartition.rs b/tests-integration/tests/repartition.rs index 880bd715c57..358b5d8406d 100644 --- a/tests-integration/tests/repartition.rs +++ b/tests-integration/tests/repartition.rs @@ -35,10 +35,14 @@ use servers::error::Result as ServerResult; use servers::query_handler::sql::SqlQueryHandler; use session::Session; use session::context::{Channel, QueryContext, QueryContextRef}; +use sqlx::Row; use store_api::codec::PrimaryKeyEncoding; use store_api::storage::{RegionId, TableId}; use tests_integration::cluster::{GreptimeDbCluster, GreptimeDbClusterBuilder}; -use tests_integration::test_util::{StorageType, get_test_store_config}; +use tests_integration::test_util::{ + StorageType, get_test_store_config, setup_mysql_server_for_frontend_instance, + setup_pg_server_for_frontend_instance, +}; use tokio::sync::oneshot; #[macro_export] @@ -682,6 +686,300 @@ async fn test_repartition_multi_hop_fulltext_index_gc() { ); } +const DEVICES: &str = "CREATE TABLE {t} (ts TIMESTAMP TIME INDEX, city STRING, device_id BIGINT, \ + PRIMARY KEY(city, device_id))"; +const DEVICES_ON_DEVICE_ID: &str = "CREATE TABLE {t} (ts TIMESTAMP TIME INDEX, city STRING, device_id BIGINT, \ + PRIMARY KEY(city, device_id)) PARTITION ON COLUMNS (device_id) (device_id < 10, device_id >= 10)"; +const DEVICE_ROWS: &str = "INSERT INTO {t} VALUES (1000, 'Aachen', 1), (1000, 'Zurich', 1), \ + (1000, 'Paris', 2), (1000, 'Aachen', 10), (1000, 'Zurich', 10), (1000, 'Beijing', 20)"; +const PARTITION_DEVICES: &str = "ALTER TABLE {t} PARTITION ON COLUMNS (device_id) \ + (device_id < 10, device_id >= 10 AND device_id < 20, device_id >= 20)"; +const METRIC_PHYSICAL: &str = r#"CREATE TABLE {t}_phy (ts TIMESTAMP TIME INDEX, val DOUBLE, host STRING PRIMARY KEY) ENGINE = metric WITH ("physical_metric_table" = "true")"#; +const METRIC_LOGICAL: &str = r#"CREATE TABLE {t} (ts TIMESTAMP TIME INDEX, val DOUBLE, host STRING PRIMARY KEY) ENGINE = metric WITH ("on_physical_table" = "{t}_phy")"#; +const METRIC_LOGICAL_NO_HOST: &str = r#"CREATE TABLE {t}_job (ts TIMESTAMP TIME INDEX, val DOUBLE, job STRING PRIMARY KEY) ENGINE = metric WITH ("on_physical_table" = "{t}_phy")"#; +const METRIC_ROWS: &str = "INSERT INTO {t} (host, ts, val) VALUES ('a_host', 1000, 1), \ + ('b_host', 1000, 2), ('n_host', 1000, 2), ('z_host', 1000, 2)"; +const METRIC_NO_HOST_ROWS: &str = + "INSERT INTO {t}_job (job, ts, val) VALUES ('a', 1000, 1), ('b', 1000, 2)"; +const PARTITION_METRIC_PHYSICAL: &str = "ALTER TABLE {t}_phy PARTITION ON COLUMNS (host) \ + (host < 'g', host >= 'g' AND host < 'p', host >= 'p')"; +const COUNT: &str = "SELECT COUNT(1) FROM {t}"; +const COUNT_BY_CITY: &str = "SELECT COUNT(1) FROM {t} GROUP BY city ORDER BY city"; +const COUNT_BY_DEVICE: &str = "SELECT COUNT(1) FROM {t} GROUP BY device_id ORDER BY device_id"; +const COUNT_BY_VAL: &str = "SELECT COUNT(1) FROM {t} GROUP BY val ORDER BY val"; +const FIRST_DEVICES: &str = "SELECT device_id FROM {t} ORDER BY device_id LIMIT 2"; +const FIRST_DEVICES_ALL_COLUMNS: &str = "SELECT * FROM {t} ORDER BY device_id LIMIT 2"; + +/// What a statement prepared before a case's DDL returns after it. +enum Expected { + /// The same as the same text prepared on a new connection. + NewStatement, + /// These values of the last column. + Rows(&'static [i64]), + Error, +} + +/// `{t}` in the statements stands for a table name unique to the case. +struct PreparedStatementCase { + name: &'static str, + setup: &'static [&'static str], + queries: &'static [(&'static str, Expected)], + ddl: &'static [&'static str], +} + +const PREPARED_STATEMENT_CASES: &[PreparedStatementCase] = &[ + PreparedStatementCase { + name: "first partitioning", + setup: &[ + DEVICES, + DEVICE_ROWS, + "CREATE VIEW {t}_view AS SELECT city, COUNT(1) AS c FROM {t} GROUP BY city", + "CREATE TABLE {t}_sink (city STRING PRIMARY KEY, c BIGINT, ts TIMESTAMP TIME INDEX)", + ], + queries: &[ + (COUNT, Expected::NewStatement), + (COUNT_BY_CITY, Expected::NewStatement), + (FIRST_DEVICES, Expected::NewStatement), + ( + "SELECT COUNT(1) FROM {t} WHERE device_id >= ?", + Expected::NewStatement, + ), + ("SELECT (SELECT COUNT(1) FROM {t})", Expected::NewStatement), + ( + "SELECT c FROM {t}_view ORDER BY city", + Expected::NewStatement, + ), + ( + "INSERT INTO {t}_sink SELECT city, COUNT(1), MAX(ts) FROM {t} GROUP BY city", + Expected::NewStatement, + ), + ( + "SELECT c FROM {t}_sink ORDER BY city", + Expected::Rows(&[2, 1, 1, 2]), + ), + ], + ddl: &[PARTITION_DEVICES], + }, + PreparedStatementCase { + name: "split adding a partition column", + setup: &[DEVICES_ON_DEVICE_ID, DEVICE_ROWS], + queries: &[(COUNT_BY_DEVICE, Expected::NewStatement)], + ddl: &[ + "ALTER TABLE {t} SPLIT PARTITION (device_id >= 10) ON COLUMNS (device_id, city) \ + INTO (device_id >= 10 AND city < 'M', device_id >= 10 AND city >= 'M')", + ], + }, + PreparedStatementCase { + name: "repartition replacing the partition column", + setup: &[DEVICES_ON_DEVICE_ID, DEVICE_ROWS], + queries: &[(COUNT_BY_DEVICE, Expected::NewStatement)], + ddl: &[ + "ALTER TABLE {t} REPARTITION (device_id < 10, device_id >= 10) ON COLUMNS (city) \ + INTO (city < 'M', city >= 'M')", + ], + }, + PreparedStatementCase { + name: "metric logical tables with and without the physical partition column", + setup: &[ + METRIC_PHYSICAL, + METRIC_LOGICAL, + METRIC_LOGICAL_NO_HOST, + METRIC_ROWS, + METRIC_NO_HOST_ROWS, + ], + queries: &[ + (COUNT, Expected::NewStatement), + (COUNT_BY_VAL, Expected::NewStatement), + ("TQL EVAL (0, 10, '5s') sum({t})", Expected::NewStatement), + ("SELECT COUNT(1) FROM {t}_job", Expected::NewStatement), + ], + ddl: &[PARTITION_METRIC_PHYSICAL], + }, + PreparedStatementCase { + name: "metric logical table gains the physical partition column", + setup: &[ + METRIC_PHYSICAL, + METRIC_LOGICAL_NO_HOST, + METRIC_NO_HOST_ROWS, + PARTITION_METRIC_PHYSICAL, + ], + queries: &[("SELECT COUNT(1) FROM {t}_job", Expected::NewStatement)], + ddl: &["ALTER TABLE {t}_job ADD COLUMN host STRING PRIMARY KEY"], + }, + PreparedStatementCase { + name: "schema and partition columns change", + setup: &[DEVICES, DEVICE_ROWS], + queries: &[ + (COUNT, Expected::NewStatement), + (COUNT_BY_CITY, Expected::NewStatement), + (FIRST_DEVICES, Expected::NewStatement), + (FIRST_DEVICES_ALL_COLUMNS, Expected::Rows(&[1, 1])), + ], + ddl: &["ALTER TABLE {t} ADD COLUMN zone STRING", PARTITION_DEVICES], + }, + PreparedStatementCase { + name: "metric logical table schema and physical partition columns change", + setup: &[METRIC_PHYSICAL, METRIC_LOGICAL, METRIC_ROWS], + queries: &[ + (COUNT, Expected::NewStatement), + (COUNT_BY_VAL, Expected::NewStatement), + ], + ddl: &[ + "ALTER TABLE {t} ADD COLUMN job STRING PRIMARY KEY", + PARTITION_METRIC_PHYSICAL, + ], + }, + PreparedStatementCase { + name: "column comment and partition columns change", + setup: &[DEVICES, DEVICE_ROWS], + queries: &[ + (COUNT, Expected::NewStatement), + (COUNT_BY_CITY, Expected::NewStatement), + ], + ddl: &["COMMENT ON COLUMN {t}.city IS 'x'", PARTITION_DEVICES], + }, + PreparedStatementCase { + name: "rename and partition columns change", + setup: &[DEVICES, DEVICE_ROWS], + queries: &[(COUNT, Expected::Rows(&[6]))], + ddl: &[ + "ALTER TABLE {t} RENAME {t}_renamed", + "ALTER TABLE {t}_renamed PARTITION ON COLUMNS (device_id) \ + (device_id < 10, device_id >= 10 AND device_id < 20, device_id >= 20)", + ], + }, + PreparedStatementCase { + name: "schema change alone", + setup: &[DEVICES_ON_DEVICE_ID, DEVICE_ROWS], + queries: &[ + (COUNT, Expected::NewStatement), + (FIRST_DEVICES_ALL_COLUMNS, Expected::Rows(&[1, 1])), + ], + ddl: &["ALTER TABLE {t} ADD COLUMN zone STRING"], + }, + PreparedStatementCase { + name: "dropped and recreated", + setup: &[DEVICES, DEVICE_ROWS], + queries: &[(COUNT, Expected::Error)], + ddl: &["DROP TABLE {t}", DEVICES], + }, +]; + +macro_rules! fetch_last_column { + ($pool:expr, $sql:expr, $param:expr) => {{ + let mut query = sqlx::query(sqlx::AssertSqlSafe($sql)); + if let Some(param) = $param { + query = query.bind(param); + } + let rows = query.fetch_all($pool).await.map_err(|e| e.to_string())?; + rows.iter() + .map(|row| { + let last = row.len() - 1; + row.try_get::(last) + .or_else(|_| row.try_get::(last).map(|v| v as i64)) + .map_err(|e| e.to_string()) + }) + .collect() + }}; +} + +/// A connection that keeps the statements it prepares. +enum PreparedClient { + Mysql(sqlx::MySqlPool), + Postgres(sqlx::PgPool), +} + +impl PreparedClient { + async fn connect_each(mysql_addr: &str, pg_addr: &str) -> [Self; 2] { + [ + Self::Mysql( + sqlx::mysql::MySqlPoolOptions::new() + .max_connections(1) + .connect(&format!("mysql://{mysql_addr}/public")) + .await + .unwrap(), + ), + Self::Postgres( + sqlx::postgres::PgPoolOptions::new() + .max_connections(1) + .connect(&format!("postgres://{pg_addr}/public")) + .await + .unwrap(), + ), + ] + } + + fn protocol(&self) -> &'static str { + match self { + Self::Mysql(_) => "mysql", + Self::Postgres(_) => "postgres", + } + } + + /// Returns the last column of each row, binding 2 to a `?` in `sql`. + async fn fetch(&self, sql: &str) -> Result, String> { + let param = sql.contains('?').then_some(2i64); + match self { + Self::Mysql(pool) => fetch_last_column!(pool, sql, param), + Self::Postgres(pool) => fetch_last_column!(pool, sql.replace('?', "$1"), param), + } + } +} + +#[tokio::test(flavor = "multi_thread")] +async fn test_prepared_statement_across_partition_column_change() { + let (cluster, _home_guard) = repartition_cluster("prepared_statement_partition").await; + let instance = cluster.fe_instance(); + let mysql = setup_mysql_server_for_frontend_instance(instance.clone(), None).await; + let pg = setup_pg_server_for_frontend_instance(instance.clone(), None).await; + let mysql_addr = mysql.bind_addr().unwrap().to_string(); + let pg_addr = pg.bind_addr().unwrap().to_string(); + + let mut failures = Vec::new(); + for (index, case) in PREPARED_STATEMENT_CASES.iter().enumerate() { + let table = format!("prepared_{index}"); + let sql = |template: &str| template.replace("{t}", &table); + for setup in case.setup { + run_sql(instance, &sql(setup), QueryContext::arc()) + .await + .unwrap(); + } + let held = PreparedClient::connect_each(&mysql_addr, &pg_addr).await; + for client in &held { + for (query, _) in case.queries { + client.fetch(&sql(query)).await.unwrap(); + } + } + for ddl in case.ddl { + run_sql(instance, &sql(ddl), QueryContext::arc()) + .await + .unwrap(); + } + + let new = PreparedClient::connect_each(&mysql_addr, &pg_addr).await; + for (held, new) in held.iter().zip(&new) { + for (query, expected) in case.queries { + let query = sql(query); + let expected = match expected { + // Runs first, so that the held statement of an `INSERT` writes last. + Expected::NewStatement => Some(new.fetch(&query).await.unwrap()), + Expected::Rows(rows) => Some(rows.to_vec()), + Expected::Error => None, + }; + let actual = held.fetch(&query).await; + if actual.as_ref().ok() != expected.as_ref() { + failures.push(format!( + "{}, {}, {query}: got {actual:?}, expected {expected:?}", + case.name, + held.protocol(), + )); + } + } + } + } + assert!(failures.is_empty(), "{failures:#?}"); +} + pub async fn test_partition_unpartitioned_mito(store_type: StorageType) { info!( "test_partition_unpartitioned_mito: store_type: {:?}",