From 9f9b35c270885acdec1381711f489bdffb626f4d Mon Sep 17 00:00:00 2001 From: Chenglong Zheng Date: Wed, 21 Oct 2026 07:57:16 +0000 Subject: [PATCH] feat: support wait and timeout options for TRUNCATE TABLE (#9753) * feat: support wait and timeout options for TRUNCATE TABLE * fix: address review feedback for truncate DDL options * fix: address review feedback for truncate DDL options --- src/common/meta/src/ddl_manager.rs | 143 +++++++++-- src/common/meta/src/error.rs | 12 + src/common/meta/src/key.rs | 1 + src/common/meta/src/kv_backend/test_util.rs | 12 +- src/operator/src/statement.rs | 3 +- src/operator/src/statement/ddl.rs | 78 +++++- src/sql/src/parsers/truncate_parser.rs | 83 ++++++- src/sql/src/parsers/utils.rs | 23 +- src/sql/src/statements/truncate.rs | 47 +++- tests-integration/tests/main.rs | 1 + tests-integration/tests/truncate.rs | 225 ++++++++++++++++++ .../common/truncate/truncate.result | 28 ++- .../standalone/common/truncate/truncate.sql | 14 +- 13 files changed, 620 insertions(+), 50 deletions(-) create mode 100644 tests-integration/tests/truncate.rs diff --git a/src/common/meta/src/ddl_manager.rs b/src/common/meta/src/ddl_manager.rs index 7b490895a4b..05100341112 100644 --- a/src/common/meta/src/ddl_manager.rs +++ b/src/common/meta/src/ddl_manager.rs @@ -247,19 +247,21 @@ pub trait RepartitionProcedureFactory: Send + Sync { /// The options for DDL tasks. /// /// Note: These options may not be utilized by all procedures. -/// At present, they are specifically applied in `RepartitionProcedure`. +/// At present, they are specifically applied in `RepartitionProcedure` and truncate tasks. #[derive(Debug, Clone, Copy)] pub struct DdlOptions { /// The timeout will be passed to the procedure. /// /// Note: Each procedure may implement its own timeout handling mechanism. + /// For truncate tasks, the timeout bounds submission and waiting rather than procedure execution. + /// A timeout does not cancel a submitted procedure. pub timeout: Duration, /// The flag that controls whether to wait for the procedure to complete. /// /// If wait is `true`, the procedure will wait for completion(success or failure) and the result will be returned. /// Otherwise, the procedure will be submitted and return the [ProcedureId](common_procedure::ProcedureId) immediately. /// - /// Note: The value of `wait` is independent of the `timeout` option. If a procedure ignores the `timeout` and `wait` is set to true, the operation returns until the procedure completes. + /// Note: Except for truncate tasks, the value of `wait` is independent of the `timeout` option. If a procedure ignores the `timeout` and `wait` is set to true, the operation returns until the procedure completes. pub wait: bool, } @@ -454,13 +456,8 @@ impl DdlManager { .context(PersistRepartitionGcRequirementSnafu)?; let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); - if wait { - self.execute_procedure_and_wait(procedure_with_id).await - } else { - self.submit_procedure(procedure_with_id) - .await - .map(|p| (p, None)) - } + self.execute_procedure_with_options(procedure_with_id, DdlOptions { wait, timeout }) + .await } /// Submits and executes an alter table task. @@ -870,13 +867,14 @@ impl DdlManager { self.execute_procedure_and_wait(procedure_with_id).await } - /// Submits and executes a truncate table task. + /// Submits a truncate table task and optionally waits for completion. #[tracing::instrument(skip_all)] pub async fn submit_truncate_table_task( &self, truncate_table_task: TruncateTableTask, table_info_value: DeserializedValueWithBytes, procedure_context: ProcedureContext, + ddl_options: DdlOptions, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); let procedure = TruncateTableProcedure::new(truncate_table_task, table_info_value, context); @@ -884,7 +882,8 @@ impl DdlManager { let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); - self.execute_procedure_and_wait(procedure_with_id).await + self.execute_procedure_with_options(procedure_with_id, ddl_options) + .await } /// Submits and executes a comment on task. @@ -908,6 +907,24 @@ impl DdlManager { self.execute_procedure_and_wait(procedure_with_id).await } + /// Submits a procedure and optionally waits for completion. + /// + /// The truncate task handler bounds its whole request, including metadata reads. + /// The manager drives submitted procedures independently of this future. + async fn execute_procedure_with_options( + &self, + procedure_with_id: ProcedureWithId, + ddl_options: DdlOptions, + ) -> Result<(ProcedureId, Option)> { + if ddl_options.wait { + self.execute_procedure_and_wait(procedure_with_id).await + } else { + self.submit_procedure(procedure_with_id) + .await + .map(|id| (id, None)) + } + } + /// Executes a procedure and waits for the result. async fn execute_procedure_and_wait( &self, @@ -1001,7 +1018,13 @@ impl DdlManager { .await } TruncateTable(truncate_table_task) => { - handle_truncate_table_task(self, truncate_table_task, procedure_context).await + handle_truncate_table_task( + self, + truncate_table_task, + procedure_context, + ddl_options, + ) + .await } CreateLogicalTables(create_table_tasks) => { handle_create_logical_table_tasks(self, create_table_tasks, procedure_context) @@ -1084,6 +1107,33 @@ async fn handle_truncate_table_task( ddl_manager: &DdlManager, truncate_table_task: TruncateTableTask, procedure_context: ProcedureContext, + ddl_options: DdlOptions, +) -> Result { + // Share one timeout across metadata checks, submission and optional waiting. + // Dropping this future does not cancel an already submitted procedure. + tokio::time::timeout( + ddl_options.timeout, + handle_truncate_table_task_inner( + ddl_manager, + truncate_table_task, + procedure_context, + ddl_options, + ), + ) + .await + .map_err(|_| { + error::DdlTaskTimeoutSnafu { + timeout: ddl_options.timeout, + } + .build() + })? +} + +async fn handle_truncate_table_task_inner( + ddl_manager: &DdlManager, + truncate_table_task: TruncateTableTask, + procedure_context: ProcedureContext, + ddl_options: DdlOptions, ) -> Result { let table_id = truncate_table_task.table_id; let table_metadata_manager = &ddl_manager.table_metadata_manager(); @@ -1108,10 +1158,18 @@ async fn handle_truncate_table_task( ); let (id, _) = ddl_manager - .submit_truncate_table_task(truncate_table_task, table_info_value, procedure_context) + .submit_truncate_table_task( + truncate_table_task, + table_info_value, + procedure_context, + ddl_options, + ) .await?; - info!("Table: {table_id} is truncated via procedure_id {id:?}"); + info!( + "Submitted truncate for table: {table_id} via procedure_id {id:?}, wait: {}", + ddl_options.wait + ); Ok(SubmitDdlTaskResponse { key: id.to_string().into(), @@ -1621,9 +1679,7 @@ mod tests { #[cfg(feature = "enterprise")] use common_base::protocol::Channel; use common_error::ext::BoxedError; - #[cfg(feature = "enterprise")] use common_error::ext::ErrorExt; - #[cfg(feature = "enterprise")] use common_error::status_code::StatusCode; #[cfg(feature = "enterprise")] use common_event_recorder::{PersistentEventContext, ProcedureEventInput, TriggerReason}; @@ -1656,6 +1712,7 @@ mod tests { use crate::key::TableMetadataManager; use crate::key::flow::FlowMetadataManager; use crate::kv_backend::memory::MemoryKvBackend; + use crate::kv_backend::test_util::MockKvBackendBuilder; use crate::node_manager::{DatanodeManager, DatanodeRef, FlownodeManager, FlownodeRef}; use crate::peer::Peer; use crate::procedure_executor::ExecutorContext; @@ -2221,6 +2278,60 @@ mod tests { } } + #[tokio::test] + async fn test_truncate_timeout_covers_metadata_checks() { + let mut context = new_ddl_context(Arc::new(DummyDatanodeManager)); + let kv_backend = MockKvBackendBuilder::default() + .range_async_fn(Arc::new(|_| Box::pin(std::future::pending()))) + .build() + .unwrap(); + context.table_metadata_manager = Arc::new(TableMetadataManager::new(Arc::new(kv_backend))); + let ddl_manager = ddl_manager_with_context(context).await; + + for wait in [true, false] { + let timeout = Duration::from_millis(20); + let request = SubmitDdlTaskRequest { + task: DdlTask::new_truncate_table( + "greptime".to_string(), + "public".to_string(), + "metrics".to_string(), + 1024, + Vec::new(), + ), + wait, + timeout, + }; + let err = tokio::time::timeout( + Duration::from_secs(1), + ddl_manager.submit_ddl_task( + ExecutorContext { + query_context: Some(QueryContext::default()), + ..Default::default() + }, + request, + ), + ) + .await + .expect("metadata reads must be bounded by the request timeout") + .unwrap_err(); + assert!(matches!( + err, + crate::error::Error::DdlTaskTimeout { timeout: actual, .. } if actual == timeout + )); + assert_eq!(err.status_code(), StatusCode::Cancelled); + assert!( + ddl_manager + .procedure_manager + .list_procedures() + .await + .unwrap() + .is_empty() + ); + } + + ddl_manager.procedure_manager.stop().await.unwrap(); + } + async fn ddl_manager_with_context(ddl_context: DdlContext) -> DdlManager { let kv_backend = Arc::new(MemoryKvBackend::new()); let state_store = Arc::new(KvStateStore::new(kv_backend.clone())); diff --git a/src/common/meta/src/error.rs b/src/common/meta/src/error.rs index 3b8f55e8a4c..c048f499dff 100644 --- a/src/common/meta/src/error.rs +++ b/src/common/meta/src/error.rs @@ -209,6 +209,15 @@ pub enum Error { source: common_procedure::Error, }, + #[snafu(display( + "DDL request timed out after {timeout:?}; a submitted procedure may still be running" + ))] + DdlTaskTimeout { + timeout: std::time::Duration, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Failed to start procedure manager"))] StartProcedureManager { #[snafu(implicit)] @@ -1343,6 +1352,9 @@ impl ErrorExt for Error { | InvalidRole { .. } | EmptyDdlTasks { .. } => StatusCode::InvalidArguments, + // DeadlineExceeded is treated as unreachable by MetaClient and + // retried. A wait timeout must not resubmit a running procedure. + DdlTaskTimeout { .. } => StatusCode::Cancelled, LoadTlsCertificate { .. } => StatusCode::Internal, #[cfg(feature = "pg_kvbackend")] diff --git a/src/common/meta/src/key.rs b/src/common/meta/src/key.rs index 644ba472233..1ed8f8df6f0 100644 --- a/src/common/meta/src/key.rs +++ b/src/common/meta/src/key.rs @@ -3318,6 +3318,7 @@ mod tests { Ok(RangeResponse { kvs, more: false }) }) }), + range_async_fn: None, batch_get_fn: Some({ let batch_get_calls = batch_get_calls.clone(); Arc::new(move |req| { diff --git a/src/common/meta/src/kv_backend/test_util.rs b/src/common/meta/src/kv_backend/test_util.rs index ce502c33321..589c6c699ef 100644 --- a/src/common/meta/src/kv_backend/test_util.rs +++ b/src/common/meta/src/kv_backend/test_util.rs @@ -16,6 +16,7 @@ use std::any::Any; use std::sync::Arc; use derive_builder::Builder; +use futures::future::BoxFuture; use crate::error::Result; use crate::kv_backend::txn::{Txn, TxnResponse}; @@ -27,11 +28,18 @@ use crate::kv_backend::{ pub type MockFn = Arc Result + Send + Sync>; +/// An asynchronous callback for mock backend operations. +pub type AsyncMockFn = + Arc BoxFuture<'static, Result> + Send + Sync>; + /// A mock kv backend for testing. #[derive(Builder)] pub struct MockKvBackend { #[builder(setter(strip_option), default)] pub range_fn: Option>, + /// Takes precedence over the synchronous range callback when set. + #[builder(setter(strip_option), default)] + pub range_async_fn: Option>, #[builder(setter(strip_option), default)] pub put_fn: Option>, #[builder(setter(strip_option), default)] @@ -76,7 +84,9 @@ impl KvBackend for MockKvBackend { } async fn range(&self, req: RangeRequest) -> Result { - if let Some(f) = &self.range_fn { + if let Some(f) = &self.range_async_fn { + f(req).await + } else if let Some(f) = &self.range_fn { f(req) } else { unimplemented!() diff --git a/src/operator/src/statement.rs b/src/operator/src/statement.rs index 6122a9f4281..b52cb814cbb 100644 --- a/src/operator/src/statement.rs +++ b/src/operator/src/statement.rs @@ -421,6 +421,7 @@ impl StatementExecutor { .await } Statement::TruncateTable(stmt) => { + let ddl_options = ddl::parse_ddl_options(stmt.options())?; let (catalog, schema, table) = table_idents_to_full_name(stmt.table_name(), &query_ctx) .map_err(BoxedError::new) @@ -429,7 +430,7 @@ impl StatementExecutor { let time_ranges = self .convert_truncate_time_ranges(&table_name, stmt.time_ranges(), &query_ctx) .await?; - self.truncate_table(table_name, time_ranges, query_ctx) + self.truncate_table_with_options(table_name, time_ranges, ddl_options, query_ctx) .await } Statement::CreateDatabase(stmt) => { diff --git a/src/operator/src/statement/ddl.rs b/src/operator/src/statement/ddl.rs index fa9da9f322e..19cc294d1bf 100644 --- a/src/operator/src/statement/ddl.rs +++ b/src/operator/src/statement/ddl.rs @@ -124,12 +124,22 @@ use crate::statement::StatementExecutor; use crate::statement::show::create_partitions_stmt; use crate::utils::{to_executor_context, to_executor_context_with_origin_frontend}; +/// Options shared by SQL DDL submission paths. #[derive(Debug, Clone, Copy)] -struct DdlSubmitOptions { +pub(crate) struct DdlSubmitOptions { wait: bool, timeout: Duration, } +impl Default for DdlSubmitOptions { + fn default() -> Self { + Self { + wait: SubmitDdlTaskRequest::default_wait(), + timeout: SubmitDdlTaskRequest::default_timeout(), + } + } +} + #[cfg(not(feature = "enterprise"))] const ALLOWED_FLOW_OPTIONS: [&str; 2] = [ DEFER_ON_MISSING_SOURCE_KEY, @@ -151,7 +161,9 @@ fn build_procedure_id_output(procedure_id: Vec) -> Result { Ok(Output::new_with_record_batches(batches)) } -fn parse_ddl_options(options: &OptionMap) -> Result { +/// Parses DDL submission options using the existing request defaults. +pub(crate) fn parse_ddl_options(options: &OptionMap) -> Result { + let defaults = DdlSubmitOptions::default(); let wait = match options.get(DDL_WAIT) { Some(value) => value.parse::().map_err(|_| { InvalidSqlSnafu { @@ -159,7 +171,7 @@ fn parse_ddl_options(options: &OptionMap) -> Result { } .build() })?, - None => SubmitDdlTaskRequest::default_wait(), + None => defaults.wait, }; let timeout = match options.get(DDL_TIMEOUT) { @@ -169,7 +181,7 @@ fn parse_ddl_options(options: &OptionMap) -> Result { } .build() })?, - None => SubmitDdlTaskRequest::default_timeout(), + None => defaults.timeout, }; Ok(DdlSubmitOptions { wait, timeout }) @@ -1717,12 +1729,29 @@ impl StatementExecutor { } } - #[tracing::instrument(skip_all)] pub async fn truncate_table( &self, table_name: TableName, time_ranges: Vec<(Timestamp, Timestamp)>, query_context: QueryContextRef, + ) -> Result { + self.truncate_table_with_options( + table_name, + time_ranges, + DdlSubmitOptions::default(), + query_context, + ) + .await + } + + /// Truncates a table, optionally waiting for the submitted procedure to finish. + #[tracing::instrument(name = "truncate_table", skip_all)] + pub(crate) async fn truncate_table_with_options( + &self, + table_name: TableName, + time_ranges: Vec<(Timestamp, Timestamp)>, + ddl_options: DdlSubmitOptions, + query_context: QueryContextRef, ) -> Result { ensure_table_writable(&table_name.schema_name, &table_name.table_name)?; @@ -1740,9 +1769,20 @@ impl StatementExecutor { table_name: table_name.to_string(), })?; let table_id = table.table_info().table_id(); - self.truncate_table_procedure(&table_name, table_id, time_ranges, query_context) + let response = self + .truncate_table_procedure( + &table_name, + table_id, + time_ranges, + ddl_options, + query_context, + ) .await?; + if !ddl_options.wait { + return build_procedure_id_output(response.key); + } + Ok(Output::new_with_affected_rows(0)) } @@ -2353,16 +2393,19 @@ impl StatementExecutor { table_name: &TableName, table_id: TableId, time_ranges: Vec<(Timestamp, Timestamp)>, + ddl_options: DdlSubmitOptions, query_context: QueryContextRef, ) -> Result { let executor_context = to_executor_context(query_context, TriggerReason::Manual); - let request = SubmitDdlTaskRequest::new(DdlTask::new_truncate_table( + let mut request = SubmitDdlTaskRequest::new(DdlTask::new_truncate_table( table_name.catalog_name.clone(), table_name.schema_name.clone(), table_name.table_name.clone(), table_id, time_ranges, )); + request.wait = ddl_options.wait; + request.timeout = ddl_options.timeout; self.procedure_executor .submit_ddl_task(executor_context, request) @@ -3272,6 +3315,10 @@ mod test { #[test] fn test_parse_ddl_options() { + let defaults = parse_ddl_options(&OptionMap::default()).unwrap(); + assert!(defaults.wait); + assert_eq!(defaults.timeout, SubmitDdlTaskRequest::default_timeout()); + let options = OptionMap::from([ ("timeout".to_string(), "5m".to_string()), ("wait".to_string(), "false".to_string()), @@ -3279,6 +3326,23 @@ mod test { let ddl_options = parse_ddl_options(&options).unwrap(); assert!(!ddl_options.wait); assert_eq!(Duration::from_secs(300), ddl_options.timeout); + + let options = OptionMap::from([("timeout".to_string(), "1500ms".to_string())]); + assert_eq!( + parse_ddl_options(&options).unwrap().timeout, + Duration::from_millis(1500) + ); + for (key, value) in [ + ("wait", "yes"), + ("timeout", "not-a-duration"), + ("timeout", "-1s"), + ] { + let options = OptionMap::from([(key.to_string(), value.to_string())]); + assert!(matches!( + parse_ddl_options(&options), + Err(error::Error::InvalidSql { .. }) + )); + } } #[test] diff --git a/src/sql/src/parsers/truncate_parser.rs b/src/sql/src/parsers/truncate_parser.rs index 1fe04865751..a5e33ef486a 100644 --- a/src/sql/src/parsers/truncate_parser.rs +++ b/src/sql/src/parsers/truncate_parser.rs @@ -18,10 +18,11 @@ use sqlparser::tokenizer::Token; use crate::error::{self, InvalidSqlSnafu, InvalidTableNameSnafu, Result, UnexpectedTokenSnafu}; use crate::parser::ParserContext; +use crate::parsers::utils; use crate::statements::statement::Statement; use crate::statements::truncate::TruncateTable; -/// `TRUNCATE [TABLE] table_name;` +/// `TRUNCATE [TABLE] table_name [FILE RANGE ...] [WITH (wait = ..., timeout = ...)];` impl ParserContext<'_> { pub(crate) fn parse_truncate(&mut self) -> Result { let _ = self.parser.next_token(); @@ -46,7 +47,10 @@ impl ParserContext<'_> { // if no range is specified, we just truncate the table if !have_range { - return Ok(Statement::TruncateTable(TruncateTable::new(table_ident))); + let options = utils::parse_ddl_options(&mut self.parser)?; + return Ok(Statement::TruncateTable( + TruncateTable::new(table_ident).with_options(options), + )); } // parse a list of time ranges consist of (Timestamp, Timestamp),? @@ -100,6 +104,7 @@ impl ParserContext<'_> { let peek = self.parser.peek_token().token; match peek { + Token::Word(word) if word.keyword == Keyword::WITH => break, sqlparser::tokenizer::Token::EOF | Token::SemiColon => { if time_ranges.is_empty() { return Err(InvalidSqlSnafu { @@ -112,24 +117,26 @@ impl ParserContext<'_> { Token::Comma => { self.parser.next_token(); // Consume the comma let next_peek = self.parser.peek_token().token; // Peek the token after the comma - if matches!(next_peek, Token::EOF | Token::SemiColon) { + if matches!(next_peek, Token::EOF | Token::SemiColon) + || matches!(next_peek, Token::Word(word) if word.keyword == Keyword::WITH) + { break; // Trailing comma, end of statement } // Otherwise, continue to parse next range continue; } _ => UnexpectedTokenSnafu { - expected: "a comma or end of statement", + expected: "a comma, WITH or end of statement", actual: self.peek_token_as_string(), } .fail()?, } } - Ok(Statement::TruncateTable(TruncateTable::new_with_ranges( - table_ident, - time_ranges, - ))) + let options = utils::parse_ddl_options(&mut self.parser)?; + Ok(Statement::TruncateTable( + TruncateTable::new_with_ranges(table_ident, time_ranges).with_options(options), + )) } } @@ -141,6 +148,64 @@ mod tests { use crate::dialect::GreptimeDbDialect; use crate::parser::ParseOptions; + #[test] + fn test_parse_truncate_ddl_options() { + for prefix in [ + "TRUNCATE foo", + "TRUNCATE TABLE foo", + "TRUNCATE TABLE foo FILE RANGE (0, 20)", + "TRUNCATE TABLE foo FILE RANGE (0, 20), (30, 40),", + ] { + let sql = format!("{prefix} WITH (WAIT = false, TIMEOUT = '30s');"); + let statements = ParserContext::create_with_dialect( + &sql, + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .unwrap(); + let Statement::TruncateTable(truncate) = &statements[0] else { + unreachable!() + }; + assert_eq!(truncate.options().get("wait"), Some("false")); + assert_eq!(truncate.options().get("timeout"), Some("30s")); + + let formatted = truncate.to_string(); + let reparsed = ParserContext::create_with_dialect( + &formatted, + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .unwrap(); + let Statement::TruncateTable(reparsed) = &reparsed[0] else { + unreachable!() + }; + assert_eq!(truncate.table_name(), reparsed.table_name()); + assert_eq!(truncate.time_ranges(), reparsed.time_ranges()); + assert_eq!( + truncate.options().to_str_map(), + reparsed.options().to_str_map() + ); + } + + for options in [ + "ttl = '1d'", + "unknown = true", + "wait = NULL", + "timeout = ['30s']", + ] { + let sql = format!("TRUNCATE TABLE foo WITH ({options})"); + assert!( + ParserContext::create_with_dialect( + &sql, + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .is_err(), + "{sql}" + ); + } + } + #[test] pub fn test_parse_truncate_with_ranges() { let sql = r#"TRUNCATE foo FILE RANGE (0, 20)"#; @@ -451,7 +516,7 @@ mod tests { assert!( result.is_err() && format!("{result:?}") - .contains("expected: 'a comma or end of statement', found: ("), + .contains("expected: 'a comma, WITH or end of statement', found: ("), "result is: {result:?}" ); diff --git a/src/sql/src/parsers/utils.rs b/src/sql/src/parsers/utils.rs index 01b136a997a..093bff02a0d 100644 --- a/src/sql/src/parsers/utils.rs +++ b/src/sql/src/parsers/utils.rs @@ -39,7 +39,9 @@ use snafu::{ResultExt, ensure}; use sqlparser::dialect::Dialect; use sqlparser::keywords::Keyword; use sqlparser::parser::Parser; -use table::requests::{SEMANTIC_PREFIX, validate_semantic_option, validate_table_option}; +use table::requests::{ + SEMANTIC_PREFIX, VALID_DDL_OPTION_KEYS, validate_semantic_option, validate_table_option, +}; use crate::error::{ ConvertToLogicalExpressionSnafu, InvalidSqlSnafu, InvalidTableOptionSnafu, ParseSqlValueSnafu, @@ -423,15 +425,32 @@ pub fn convert_month_day_nano_to_duration( Ok(std::time::Duration::new(adjusted_seconds, nanos_remainder)) } +/// Parses submission options, accepting only scalar `wait` and `timeout` values. +pub(crate) fn parse_ddl_options(parser: &mut Parser) -> Result { + parse_with_options_inner(parser, true) +} + pub fn parse_with_options(parser: &mut Parser) -> Result { + parse_with_options_inner(parser, false) +} + +fn parse_with_options_inner(parser: &mut Parser, ddl_only: bool) -> Result { let options = parser .parse_options(Keyword::WITH) .context(SyntaxSnafu)? .into_iter() .map(parse_option_string) .collect::>>()?; + // Validate before OptionMap can discard non-scalar secret options. for (key, value) in &options { - if key.starts_with(SEMANTIC_PREFIX) { + if ddl_only { + ensure!( + VALID_DDL_OPTION_KEYS.contains(&key.as_str()) && value.as_string().is_some(), + InvalidSqlSnafu { + msg: format!("invalid DDL option '{key}': '{value}'") + } + ); + } else if key.starts_with(SEMANTIC_PREFIX) { // Semantic keys are whitelisted and value-checked against their domain, // so a user cannot set an unknown key or an out-of-range value. let value = value.as_string().unwrap_or_default(); diff --git a/src/sql/src/statements/truncate.rs b/src/sql/src/statements/truncate.rs index 307326d9c71..9e5ae5b1814 100644 --- a/src/sql/src/statements/truncate.rs +++ b/src/sql/src/statements/truncate.rs @@ -18,11 +18,15 @@ use itertools::Itertools; use serde::Serialize; use sqlparser::ast::{ObjectName, Visit, VisitMut, Visitor, VisitorMut}; +use crate::statements::OptionMap; + /// TRUNCATE TABLE statement. #[derive(Debug, Clone, PartialEq, Eq, Serialize)] pub struct TruncateTable { table_name: ObjectName, time_ranges: Vec<(sqlparser::ast::Value, sqlparser::ast::Value)>, + /// Options for this DDL submission, never persisted as table options. + options: OptionMap, } impl Visit for TruncateTable { @@ -32,6 +36,7 @@ impl Visit for TruncateTable { start.visit(visitor)?; end.visit(visitor)?; } + Visit::visit(&self.options, visitor)?; ::std::ops::ControlFlow::Continue(()) } } @@ -43,6 +48,7 @@ impl VisitMut for TruncateTable { start.visit(visitor)?; end.visit(visitor)?; } + VisitMut::visit(&mut self.options, visitor)?; ::std::ops::ControlFlow::Continue(()) } } @@ -53,6 +59,7 @@ impl TruncateTable { Self { table_name, time_ranges: Vec::new(), + options: OptionMap::default(), } } @@ -64,9 +71,21 @@ impl TruncateTable { Self { table_name, time_ranges, + options: OptionMap::default(), } } + /// Sets the options for this DDL submission. + pub fn with_options(mut self, options: OptionMap) -> Self { + self.options = options; + self + } + + /// Returns the options for this DDL submission. + pub fn options(&self) -> &OptionMap { + &self.options + } + pub fn table_name(&self) -> &ObjectName { &self.table_name } @@ -80,19 +99,23 @@ impl Display for TruncateTable { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { let table_name = self.table_name(); write!(f, r#"TRUNCATE TABLE {table_name}"#)?; - if self.time_ranges.is_empty() { - return Ok(()); + if !self.time_ranges.is_empty() { + write!(f, " FILE RANGE ")?; + write!( + f, + "{}", + self.time_ranges + .iter() + .map(|(start, end)| format!("({}, {})", start, end)) + .join(", ") + )?; } - - write!(f, " FILE RANGE ")?; - write!( - f, - "{}", - self.time_ranges - .iter() - .map(|(start, end)| format!("({}, {})", start, end)) - .join(", ") - ) + if !self.options.is_empty() { + f.write_str(" WITH (")?; + self.options.write_sql(f, ", ", "", false)?; + f.write_str(")")?; + } + Ok(()) } } diff --git a/tests-integration/tests/main.rs b/tests-integration/tests/main.rs index 72d0a1850bd..714c438344d 100644 --- a/tests-integration/tests/main.rs +++ b/tests-integration/tests/main.rs @@ -32,6 +32,7 @@ mod sql; #[macro_use] mod region_migration; mod table_ddl_event; +mod truncate; #[macro_use] mod repartition; mod repartition_event; diff --git a/tests-integration/tests/truncate.rs b/tests-integration/tests/truncate.rs new file mode 100644 index 00000000000..c1e7735b633 --- /dev/null +++ b/tests-integration/tests/truncate.rs @@ -0,0 +1,225 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; +use std::time::Duration; + +use client::OutputData; +use common_error::ext::ErrorExt; +use common_error::status_code::StatusCode; +use common_meta::lock_key::TableLock; +use common_procedure::{ + LockKey, Procedure, ProcedureId, ProcedureState, ProcedureWithId, Status, watcher, +}; +use datatypes::arrow::array::AsArray; +use tests_integration::cluster::GreptimeDbClusterBuilder; +use tests_integration::standalone::GreptimeDbStandaloneBuilder; +use tests_integration::test_util::{execute_sql, execute_sql_and_expect, try_execute_sql}; +use tokio::sync::Notify; + +/// Holds the table lock until explicitly released, keeping truncate pending. +struct HoldTruncateTableLock { + table_id: u32, + started: Arc, + release: Arc, +} + +#[async_trait::async_trait] +impl Procedure for HoldTruncateTableLock { + fn type_name(&self) -> &str { + "test::HoldTruncateTableLock" + } + + async fn execute(&mut self, _: &common_procedure::Context) -> common_procedure::Result { + self.started.notify_one(); + self.release.notified().await; + Ok(Status::done()) + } + + fn dump(&self) -> common_procedure::Result { + Ok(String::new()) + } + + fn lock_key(&self) -> LockKey { + LockKey::new(vec![TableLock::Write(self.table_id).into()]) + } +} + +#[rstest_reuse::apply(crate::both_deployment_cases)] +async fn test_truncate_wait_and_timeout(distributed: bool) { + let cluster = if distributed { + Some( + GreptimeDbClusterBuilder::new("truncate_wait_and_timeout") + .await + .with_datanodes(1) + .build(false) + .await, + ) + } else { + None + }; + let standalone = if distributed { + None + } else { + Some( + GreptimeDbStandaloneBuilder::new("truncate_wait_and_timeout") + .build() + .await, + ) + }; + let (frontend, manager) = if let Some(cluster) = &cluster { + ( + cluster.fe_instance().clone(), + cluster.metasrv.procedure_manager().clone(), + ) + } else { + let standalone = standalone.as_ref().unwrap(); + ( + standalone.fe_instance().clone(), + standalone.procedure_manager.clone(), + ) + }; + + execute_sql( + &frontend, + "CREATE TABLE truncate_wait_test (v INT, ts TIMESTAMP TIME INDEX)", + ) + .await; + execute_sql(&frontend, "INSERT INTO truncate_wait_test VALUES (1, 1)").await; + let definition = execute_sql(&frontend, "SHOW CREATE TABLE truncate_wait_test") + .await + .data + .pretty_print() + .await; + let table_id = frontend + .catalog_manager() + .table("greptime", "public", "truncate_wait_test", None) + .await + .unwrap() + .unwrap() + .table_info() + .table_id(); + let started = Arc::new(Notify::new()); + let release = Arc::new(Notify::new()); + let mut blocker = manager + .submit(ProcedureWithId::with_random_id(Box::new( + HoldTruncateTableLock { + table_id, + started: started.clone(), + release: release.clone(), + }, + ))) + .await + .unwrap(); + tokio::time::timeout(Duration::from_secs(10), started.notified()) + .await + .unwrap(); + + let output = execute_sql( + &frontend, + "TRUNCATE TABLE truncate_wait_test WITH (wait=false, timeout='10s')", + ) + .await; + let OutputData::RecordBatches(batches) = output.data else { + panic!("expected Procedure ID result"); + }; + let batch = batches.iter().next().unwrap(); + let id = batch + .column_by_name("Procedure ID") + .unwrap() + .as_string::() + .value(0); + let id = ProcedureId::parse_str(id).unwrap(); + let state = execute_sql(&frontend, &format!("SELECT procedure_state('{id}')")) + .await + .data + .pretty_print() + .await; + assert!(state.contains(r#""status":"Running""#), "{state}"); + + // The same table lock holds a synchronous truncate past its wait deadline. + // Give metadata reads and submission a generous budget. Using the same budget + // also verifies that the async procedure keeps running after it expires. + let err = try_execute_sql( + &frontend, + "TRUNCATE TABLE truncate_wait_test WITH (wait=true, timeout='10s')", + ) + .await + .unwrap_err(); + if distributed { + assert!(matches!( + err.status_code(), + StatusCode::DeadlineExceeded | StatusCode::Cancelled + )); + } else { + assert_eq!(err.status_code(), StatusCode::Cancelled); + assert!(err.output_msg().contains("DDL request timed out"), "{err}"); + } + let pending = manager + .list_procedures() + .await + .unwrap() + .into_iter() + .filter(|procedure| procedure.type_name == "metasrv-procedure::TruncateTable") + .collect::>(); + assert_eq!(pending.len(), 2); + let mut watchers = pending + .into_iter() + .map(|procedure| { + assert!(matches!(procedure.state, ProcedureState::Running)); + manager.procedure_watcher(procedure.id).unwrap() + }) + .collect::>(); + execute_sql_and_expect( + &frontend, + "SELECT * FROM truncate_wait_test", + r#" ++---+-------------------------+ +| v | ts | ++---+-------------------------+ +| 1 | 1970-01-01T00:00:00.001 | ++---+-------------------------+ +"#, + ) + .await; + + release.notify_one(); + tokio::time::timeout(Duration::from_secs(10), async { + watcher::wait(&mut blocker).await.unwrap(); + for watcher in &mut watchers { + watcher::wait(watcher).await.unwrap(); + } + }) + .await + .unwrap(); + execute_sql_and_expect( + &frontend, + "SELECT * FROM truncate_wait_test", + r#" ++---+----+ +| v | ts | ++---+----+ ++---+----+ +"#, + ) + .await; + assert_eq!( + definition, + execute_sql(&frontend, "SHOW CREATE TABLE truncate_wait_test") + .await + .data + .pretty_print() + .await + ); +} diff --git a/tests/cases/standalone/common/truncate/truncate.result b/tests/cases/standalone/common/truncate/truncate.result index 27fcbd90265..36239ed80e5 100644 --- a/tests/cases/standalone/common/truncate/truncate.result +++ b/tests/cases/standalone/common/truncate/truncate.result @@ -126,7 +126,7 @@ SELECT ts, host, cpu, memory FROM monitor ORDER BY ts; | 2023-09-20T13:47:44 | host3 | 88.8 | 4096.0 | +---------------------+-------+------+--------+ -TRUNCATE monitor; +TRUNCATE monitor WITH (wait = true, timeout = '30s'); Affected Rows: 0 @@ -137,6 +137,32 @@ SELECT ts, host, cpu, memory FROM monitor ORDER BY ts; +----+------+-----+--------+ +----+------+-----+--------+ +TRUNCATE monitor WITH (wait = 'invalid'); + +Error: 1004(InvalidArguments), Invalid SQL, error: invalid DDL option 'wait': 'invalid' + +TRUNCATE monitor WITH (timeout = 'invalid'); + +Error: 1004(InvalidArguments), Invalid SQL, error: invalid DDL option 'timeout': 'invalid': expected number at 0 + +TRUNCATE monitor WITH (ttl = '1d'); + +Error: 2000(InvalidSyntax), Invalid SQL, error: invalid DDL option 'ttl': ''1d'' + +-- SQLNESS REPLACE ([0-9a-fA-F]{8}-[0-9a-fA-F]{4}-4[0-9a-fA-F]{3}-[89abAB][0-9a-fA-F]{3}-[0-9a-fA-F]{12}) PROC_ID +TRUNCATE monitor WITH (wait = false, timeout = '30s'); + ++--------------------------------------+ +| Procedure ID | ++--------------------------------------+ +| PROC_ID | ++--------------------------------------+ + +-- Keep a synchronous truncate in the same case to check the original result. +TRUNCATE monitor; + +Affected Rows: 0 + DROP TABLE monitor; Affected Rows: 0 diff --git a/tests/cases/standalone/common/truncate/truncate.sql b/tests/cases/standalone/common/truncate/truncate.sql index a116664f401..8ed86547e56 100644 --- a/tests/cases/standalone/common/truncate/truncate.sql +++ b/tests/cases/standalone/common/truncate/truncate.sql @@ -51,10 +51,22 @@ INSERT INTO monitor(ts, host, cpu, memory) VALUES SELECT ts, host, cpu, memory FROM monitor ORDER BY ts; -TRUNCATE monitor; +TRUNCATE monitor WITH (wait = true, timeout = '30s'); SELECT ts, host, cpu, memory FROM monitor ORDER BY ts; +TRUNCATE monitor WITH (wait = 'invalid'); + +TRUNCATE monitor WITH (timeout = 'invalid'); + +TRUNCATE monitor WITH (ttl = '1d'); + +-- SQLNESS REPLACE ([0-9a-fA-F]{8}-[0-9a-fA-F]{4}-4[0-9a-fA-F]{3}-[89abAB][0-9a-fA-F]{3}-[0-9a-fA-F]{12}) PROC_ID +TRUNCATE monitor WITH (wait = false, timeout = '30s'); + +-- Keep a synchronous truncate in the same case to check the original result. +TRUNCATE monitor; + DROP TABLE monitor;