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: {:?}",