diff --git a/src/common/grpc-expr/src/alter.rs b/src/common/grpc-expr/src/alter.rs index 85de11d5cd..d5214ae5dd 100644 --- a/src/common/grpc-expr/src/alter.rs +++ b/src/common/grpc-expr/src/alter.rs @@ -49,33 +49,14 @@ const LOCATION_TYPE_AFTER: i32 = LocationType::After as i32; /// mixed batch — annotation alters skip region dispatch, so they cannot share /// a statement with options that regions must see. fn annotation_family_of_keys<'a>( - mut keys: impl Iterator, + keys: impl Iterator, ) -> Result> { - let Some(first) = keys.next() else { - return Ok(None); - }; - let family = AnnotationFamily::of_key(first); - let mut count = 1usize; - for key in keys { - count += 1; - let this = AnnotationFamily::of_key(key); - if this != family { - return error::InvalidTableOptionRequestSnafu { - err_msg: family.or(this).unwrap().mixed_batch_error(), - } - .fail(); + table::requests::validate_annotation_keys(keys).map_err(|err| { + error::InvalidTableOptionRequestSnafu { + err_msg: err.to_string(), } - } - if let Some(family) = family - && family.requires_single_key() - && count > 1 - { - return error::InvalidTableOptionRequestSnafu { - err_msg: family.mixed_batch_error(), - } - .fail(); - } - Ok(family) + .build() + }) } /// Returns the annotation family when `kind` is a SET/UNSET whose keys all @@ -583,6 +564,131 @@ mod tests { assert_eq!("mem_usage".to_string(), drop_names.pop().unwrap()); } + #[test] + fn test_repartition_hints_together() { + let keys = [ + REPARTITION_COLUMN_HINT_KEY, + table::requests::REPARTITION_PARTITION_NUM_HINT_KEY, + ]; + let options = vec![ + (keys[0].to_string(), "host".to_string()), + (keys[1].to_string(), "10".to_string()), + ]; + let set = Kind::SetTableOptions(SetTableOptions { + table_options: options + .iter() + .map(|(key, value)| PbOption { + key: key.clone(), + value: value.clone(), + }) + .collect(), + }); + let unset = Kind::UnsetTableOptions(UnsetTableOptions { + keys: keys.map(str::to_string).to_vec(), + }); + for kind in [set, unset] { + assert_eq!( + annotation_alter_family(&kind).unwrap(), + Some(AnnotationFamily::RepartitionHint) + ); + let request = alter_expr_to_request( + 1, + AlterTableExpr { + kind: Some(kind), + ..Default::default() + }, + None, + ) + .unwrap(); + match request.alter_kind { + AlterKind::SetAnnotations { + family, + options: actual, + } => { + assert_eq!(family, AnnotationFamily::RepartitionHint); + assert_eq!(actual, options); + } + AlterKind::UnsetAnnotations { + family, + keys: actual, + } => { + assert_eq!(family, AnnotationFamily::RepartitionHint); + assert_eq!(actual, keys); + } + other => panic!("unexpected alter kind: {other:?}"), + } + } + } + + #[test] + fn test_repartition_partition_num_hint_expr() { + let key = table::requests::REPARTITION_PARTITION_NUM_HINT_KEY; + let set = Kind::SetTableOptions(SetTableOptions { + table_options: vec![PbOption { + key: key.to_string(), + value: "8".to_string(), + }], + }); + let unset = Kind::UnsetTableOptions(UnsetTableOptions { + keys: vec![key.to_string()], + }); + for kind in [set, unset] { + assert_eq!( + annotation_alter_family(&kind).unwrap(), + Some(AnnotationFamily::RepartitionHint) + ); + let request = alter_expr_to_request( + 1, + AlterTableExpr { + kind: Some(kind), + ..Default::default() + }, + None, + ) + .unwrap(); + match request.alter_kind { + AlterKind::SetAnnotations { family, options } => { + assert_eq!(family, AnnotationFamily::RepartitionHint); + assert_eq!(options, vec![(key.to_string(), "8".to_string())]); + } + AlterKind::UnsetAnnotations { family, keys } => { + assert_eq!(family, AnnotationFamily::RepartitionHint); + assert_eq!(keys, vec![key.to_string()]); + } + other => panic!("unexpected alter kind: {other:?}"), + } + } + for other in [key, table::requests::TTL_KEY] { + for kind in [ + Kind::SetTableOptions(SetTableOptions { + table_options: [key, other] + .into_iter() + .map(|key| PbOption { + key: key.to_string(), + value: "8".to_string(), + }) + .collect(), + }), + Kind::UnsetTableOptions(UnsetTableOptions { + keys: vec![key.to_string(), other.to_string()], + }), + ] { + assert!(annotation_alter_family(&kind).is_err()); + assert!( + alter_expr_to_request( + 1, + AlterTableExpr { + kind: Some(kind), + ..Default::default() + }, + None + ) + .is_err() + ); + } + } + } + #[test] fn test_set_repartition_column_hint_expr() { let expr = AlterTableExpr { @@ -633,7 +739,7 @@ mod tests { let err = alter_expr_to_request(1, expr, None).unwrap_err(); assert!( err.to_string() - .contains("repartition.column.hint must be altered separately") + .contains("repartition hints must be altered separately") ); // Duplicate hint entries are not a meaningful batch either. @@ -655,10 +761,7 @@ mod tests { })), }; let err = alter_expr_to_request(1, dup, None).unwrap_err(); - assert!( - err.to_string() - .contains("repartition.column.hint must be altered separately") - ); + assert!(err.to_string().contains("duplicate repartition hint keys")); } #[test] diff --git a/src/common/meta/src/ddl/tests/alter_table.rs b/src/common/meta/src/ddl/tests/alter_table.rs index 30910eb0ba..f0ceaf592e 100644 --- a/src/common/meta/src/ddl/tests/alter_table.rs +++ b/src/common/meta/src/ddl/tests/alter_table.rs @@ -1123,58 +1123,85 @@ async fn test_on_submit_alter_request_with_exist_poison() { } #[tokio::test] -async fn test_semantic_annotation_alter_is_metadata_only() { - let (tx, mut rx) = mpsc::channel(8); - let node_manager = Arc::new(MockDatanodeManager::new(DatanodeWatcher::new(tx))); - let ddl_context = new_ddl_context(node_manager); - let table_id = 1024; - let table_name = "foo"; - let task = test_create_table_task(table_name, table_id); - ddl_context - .table_metadata_manager - .create_table_metadata( - task.table_info.clone(), - prepare_table_route(table_id), - HashMap::new(), - ) - .await - .unwrap(); +async fn test_annotation_alter_is_metadata_only() { + for (key, value) in [ + ("greptime.semantic.signal_type", "metric"), + (table::requests::REPARTITION_PARTITION_NUM_HINT_KEY, "8"), + ] { + let (tx, mut rx) = mpsc::channel(8); + let node_manager = Arc::new(MockDatanodeManager::new(DatanodeWatcher::new(tx))); + let ddl_context = new_ddl_context(node_manager); + let table_id = 1024; + let table_name = "foo"; + let task = test_create_table_task(table_name, table_id); + ddl_context + .table_metadata_manager + .create_table_metadata( + task.table_info.clone(), + prepare_table_route(table_id), + HashMap::new(), + ) + .await + .unwrap(); - let alter_table_task = AlterTableTask { - alter_table: AlterTableExpr { - catalog_name: DEFAULT_CATALOG_NAME.to_string(), - schema_name: DEFAULT_SCHEMA_NAME.to_string(), - table_name: table_name.to_string(), - kind: Some(Kind::SetTableOptions(SetTableOptions { - table_options: vec![api::v1::Option { - key: "greptime.semantic.signal_type".to_string(), - value: "metric".to_string(), - }], - })), - }, - }; - let mut procedure = - AlterTableProcedure::new(table_id, alter_table_task, ddl_context.clone()).unwrap(); - execute_procedure_until_done(&mut procedure).await; + let alter_table_task = AlterTableTask { + alter_table: AlterTableExpr { + catalog_name: DEFAULT_CATALOG_NAME.to_string(), + schema_name: DEFAULT_SCHEMA_NAME.to_string(), + table_name: table_name.to_string(), + kind: Some(Kind::SetTableOptions(SetTableOptions { + table_options: vec![api::v1::Option { + key: key.to_string(), + value: value.to_string(), + }], + })), + }, + }; + let mut procedure = + AlterTableProcedure::new(table_id, alter_table_task, ddl_context.clone()).unwrap(); + execute_procedure_until_done(&mut procedure).await; - // Metadata-only: no region request reaches any datanode. - rx.try_recv().unwrap_err(); + // Metadata-only: no region request reaches any datanode. + rx.try_recv().unwrap_err(); - let table_info = ddl_context - .table_metadata_manager - .table_info_manager() - .get(table_id) - .await - .unwrap() - .unwrap() - .into_inner() - .table_info; - assert_eq!( - table_info - .meta - .options - .extra_options - .get("greptime.semantic.signal_type"), - Some(&"metric".to_string()) - ); + let table_info = ddl_context + .table_metadata_manager + .table_info_manager() + .get(table_id) + .await + .unwrap() + .unwrap() + .into_inner() + .table_info; + assert_eq!( + table_info.meta.options.extra_options.get(key), + Some(&value.to_string()) + ); + for _ in 0..2 { + let task = AlterTableTask { + alter_table: AlterTableExpr { + catalog_name: DEFAULT_CATALOG_NAME.to_string(), + schema_name: DEFAULT_SCHEMA_NAME.to_string(), + table_name: table_name.to_string(), + kind: Some(Kind::UnsetTableOptions(api::v1::UnsetTableOptions { + keys: vec![key.to_string()], + })), + }, + }; + let mut procedure = + AlterTableProcedure::new(table_id, task, ddl_context.clone()).unwrap(); + execute_procedure_until_done(&mut procedure).await; + rx.try_recv().unwrap_err(); + let info = ddl_context + .table_metadata_manager + .table_info_manager() + .get(table_id) + .await + .unwrap() + .unwrap() + .into_inner() + .table_info; + assert!(!info.meta.options.extra_options.contains_key(key)); + } + } } diff --git a/src/common/meta/src/ddl_manager.rs b/src/common/meta/src/ddl_manager.rs index d88f08918e..30911b7b27 100644 --- a/src/common/meta/src/ddl_manager.rs +++ b/src/common/meta/src/ddl_manager.rs @@ -2036,20 +2036,25 @@ mod tests { // The repartition hint drives physical repartitioning; on a logical // route it stays rejected by the guard. - let err = ddl_manager - .submit_ddl_task( - ExecutorContext { - query_context: Some(QueryContext::default()), - ..Default::default() - }, - SubmitDdlTaskRequest::new(DdlTask::new_alter_table(set_options_expr( - "logical", - &[("repartition.column.hint", "host")], - ))), - ) - .await - .unwrap_err(); - let msg = common_error::ext::ErrorExt::output_msg(&err); - assert!(msg.contains("non-physical TableRouteValue"), "{msg}"); + for (key, value) in [ + (table::requests::REPARTITION_COLUMN_HINT_KEY, "host"), + (table::requests::REPARTITION_PARTITION_NUM_HINT_KEY, "8"), + ] { + let err = ddl_manager + .submit_ddl_task( + ExecutorContext { + query_context: Some(QueryContext::default()), + ..Default::default() + }, + SubmitDdlTaskRequest::new(DdlTask::new_alter_table(set_options_expr( + "logical", + &[(key, value)], + ))), + ) + .await + .unwrap_err(); + let msg = common_error::ext::ErrorExt::output_msg(&err); + assert!(msg.contains("non-physical TableRouteValue"), "{msg}"); + } } } diff --git a/src/operator/src/statement/ddl.rs b/src/operator/src/statement/ddl.rs index 59695a6e7b..0e54ba4244 100644 --- a/src/operator/src/statement/ddl.rs +++ b/src/operator/src/statement/ddl.rs @@ -2587,6 +2587,7 @@ fn validate_and_normalize_annotations( CheckError::ColumnNotFound { column } => ColumnNotFoundSnafu { msg: column }.build(), e @ (CheckError::UnknownKey { .. } | CheckError::InvalidValue { .. } + | CheckError::InvalidPartitionNumHint { .. } | CheckError::ColumnNotStringForm { .. }) => InvalidSqlSnafu { err_msg: e.to_string(), } @@ -3525,6 +3526,36 @@ SELECT max(c1), min(c2) FROM schema_2.table_2;"; } } + #[test] + fn test_create_table_with_repartition_partition_num_hint() { + let key = table::requests::REPARTITION_PARTITION_NUM_HINT_KEY; + for (value, expected) in [ + (" 8 ", Some("8")), + ("0", None), + ("-1", None), + ("4294967296", None), + ] { + let expr = create_expr_from_sql(&format!( + "CREATE TABLE metrics (host STRING, ts TIMESTAMP TIME INDEX, PRIMARY KEY(host)) WITH ('{key}' = '{value}')" + )); + let result = create_table_info(&expr, vec![]); + if let Some(expected) = expected { + let info = result.unwrap(); + assert_eq!( + info.meta.options.extra_options.get(key).map(String::as_str), + Some(expected) + ); + } else { + assert!( + result + .unwrap_err() + .to_string() + .contains("expects a positive integer") + ); + } + } + } + #[test] fn test_create_table_with_repartition_column_hint() { let expr = create_expr_from_sql( diff --git a/src/table/src/metadata.rs b/src/table/src/metadata.rs index 69803c7855..b664c5cfd5 100644 --- a/src/table/src/metadata.rs +++ b/src/table/src/metadata.rs @@ -40,9 +40,10 @@ use store_api::storage::{ColumnDescriptor, ColumnDescriptorBuilder, ColumnId}; use crate::error::{self, Result}; use crate::requests::{ AddColumnRequest, AlterKind, AnnotationContext, AnnotationFamily, AnnotationValidationError, - ModifyColumnTypeRequest, REPARTITION_COLUMN_HINT_KEY, SetDefaultRequest, SetIndexOption, - TableOptions, UnsetIndexOption, has_stable_string_form, parse_entity_columns, - parse_entity_option_key, validate_and_normalize_annotation, + ModifyColumnTypeRequest, REPARTITION_COLUMN_HINT_KEY, REPARTITION_PARTITION_NUM_HINT_KEY, + SetDefaultRequest, SetIndexOption, TableOptions, UnsetIndexOption, has_stable_string_form, + parse_entity_columns, parse_entity_option_key, validate_and_normalize_annotation, + validate_annotation_keys, }; use crate::table_reference::TableReference; @@ -461,13 +462,13 @@ impl TableMeta { family: AnnotationFamily, options: &[(String, String)], ) -> Result { - ensure!( - !family.requires_single_key() || options.len() == 1, + validate_annotation_keys(options.iter().map(|(key, _)| key.as_str())).map_err(|err| { error::InvalidAlterRequestSnafu { table: table_name, - err: family.mixed_batch_error(), + err: err.to_string(), } - ); + .build() + })?; let cx = AnnotationContext { schema: &self.schema, partition_key_indices: &self.partition_key_indices, @@ -516,13 +517,13 @@ impl TableMeta { family: AnnotationFamily, keys: &[String], ) -> Result { - ensure!( - !family.requires_single_key() || keys.len() == 1, + validate_annotation_keys(keys.iter().map(String::as_str)).map_err(|err| { error::InvalidAlterRequestSnafu { table: table_name, - err: family.mixed_batch_error(), + err: err.to_string(), } - ); + .build() + })?; let mut new_options = self.options.clone(); for key in keys { ensure!( @@ -1491,6 +1492,7 @@ impl TableInfo { pub fn to_region_options(&self) -> HashMap { let mut options = HashMap::from(&self.meta.options); options.remove(REPARTITION_COLUMN_HINT_KEY); + options.remove(REPARTITION_PARTITION_NUM_HINT_KEY); options } @@ -1943,6 +1945,160 @@ mod tests { ); } + #[test] + fn test_set_unset_repartition_hints_together() { + let meta = TableMetaBuilder::empty() + .schema(Arc::new(new_test_schema())) + .primary_key_indices(vec![0]) + .engine("engine") + .next_column_id(3) + .build() + .unwrap(); + let options = vec![ + (REPARTITION_COLUMN_HINT_KEY.to_string(), "col1".to_string()), + ( + REPARTITION_PARTITION_NUM_HINT_KEY.to_string(), + "10".to_string(), + ), + ]; + for reverse in [false, true] { + let mut options = options.clone(); + if reverse { + options.reverse(); + } + let updated = meta + .builder_with_alter_kind( + "t", + &AlterKind::SetAnnotations { + family: AnnotationFamily::RepartitionHint, + options: options.clone(), + }, + ) + .unwrap() + .build() + .unwrap(); + assert_eq!( + updated.options.extra_options, + options.iter().cloned().collect() + ); + for (key, value) in [ + (REPARTITION_COLUMN_HINT_KEY, "missing"), + (REPARTITION_PARTITION_NUM_HINT_KEY, "0"), + ] { + let invalid = options + .iter() + .map(|(k, v)| { + ( + k.clone(), + if k == key { + value.to_string() + } else { + v.clone() + }, + ) + }) + .collect(); + assert!( + updated + .builder_with_alter_kind( + "t", + &AlterKind::SetAnnotations { + family: AnnotationFamily::RepartitionHint, + options: invalid, + } + ) + .is_err() + ); + assert_eq!( + updated.options.extra_options, + options.iter().cloned().collect() + ); + } + let cleared = updated + .builder_with_alter_kind( + "t", + &AlterKind::UnsetAnnotations { + family: AnnotationFamily::RepartitionHint, + keys: options.into_iter().map(|(key, _)| key).collect(), + }, + ) + .unwrap() + .build() + .unwrap(); + assert!(cleared.options.extra_options.is_empty()); + } + } + + #[test] + fn test_repartition_partition_num_hint() { + for partition_key_indices in [vec![], vec![0]] { + let mut meta = TableMetaBuilder::empty() + .schema(Arc::new(new_test_schema())) + .primary_key_indices(vec![0]) + .partition_key_indices(partition_key_indices) + .engine("engine") + .next_column_id(3) + .build() + .unwrap(); + for value in ["", " ", "0", "-1", "1.5", "abc", "4294967296"] { + let err = meta + .builder_with_alter_kind( + "t", + &AlterKind::SetAnnotations { + family: AnnotationFamily::RepartitionHint, + options: vec![( + REPARTITION_PARTITION_NUM_HINT_KEY.to_string(), + value.to_string(), + )], + }, + ) + .err() + .unwrap(); + assert!( + err.to_string().contains("expects a positive integer"), + "{err}" + ); + } + for (value, expected) in [(" 8 ", "8"), ("1", "1"), ("4294967295", "4294967295")] { + meta = meta + .builder_with_alter_kind( + "t", + &AlterKind::SetAnnotations { + family: AnnotationFamily::RepartitionHint, + options: vec![( + REPARTITION_PARTITION_NUM_HINT_KEY.to_string(), + value.to_string(), + )], + }, + ) + .unwrap() + .build() + .unwrap(); + assert_eq!( + meta.options.extra_options, + HashMap::from([( + REPARTITION_PARTITION_NUM_HINT_KEY.to_string(), + expected.to_string() + )]) + ); + } + for _ in 0..2 { + meta = meta + .builder_with_alter_kind( + "t", + &AlterKind::UnsetAnnotations { + family: AnnotationFamily::RepartitionHint, + keys: vec![REPARTITION_PARTITION_NUM_HINT_KEY.to_string()], + }, + ) + .unwrap() + .build() + .unwrap(); + assert!(meta.options.extra_options.is_empty()); + } + } + } + #[test] fn test_set_repartition_column_hint() { let meta = TableMetaBuilder::empty() @@ -2141,11 +2297,15 @@ mod tests { } #[test] - fn test_repartition_column_hint_is_not_region_option() { + fn test_repartition_hints_are_not_region_options() { let mut table_options = TableOptions::default(); table_options .extra_options .insert(REPARTITION_COLUMN_HINT_KEY.to_string(), "col1".to_string()); + table_options.extra_options.insert( + REPARTITION_PARTITION_NUM_HINT_KEY.to_string(), + "8".to_string(), + ); let table_info = TableInfoBuilder::default() .table_id(1) .table_version(0) @@ -2165,6 +2325,11 @@ mod tests { .build() .unwrap(); + assert!( + !table_info + .to_region_options() + .contains_key(REPARTITION_PARTITION_NUM_HINT_KEY) + ); assert!( !table_info .to_region_options() @@ -2972,7 +3137,7 @@ mod tests { } #[test] - fn test_repartition_hint_batch_must_be_single_key() { + fn test_repartition_hint_batch_rejects_duplicate_keys() { let meta = TableMetaBuilder::empty() .schema(Arc::new(new_test_schema())) .primary_key_indices(vec![0]) @@ -2995,7 +3160,7 @@ mod tests { .err() .unwrap(); assert!( - err.to_string().contains("must be altered separately"), + err.to_string().contains("duplicate repartition hint keys"), "{err}" ); @@ -3011,7 +3176,7 @@ mod tests { .err() .unwrap(); assert!( - err.to_string().contains("must be altered separately"), + err.to_string().contains("duplicate repartition hint keys"), "{err}" ); } diff --git a/src/table/src/requests.rs b/src/table/src/requests.rs index 3c9d0679ae..514d9f7800 100644 --- a/src/table/src/requests.rs +++ b/src/table/src/requests.rs @@ -77,7 +77,7 @@ pub fn is_trace_v1_table(table_info: &crate::metadata::TableInfo) -> bool { pub const OTLP_METRIC_COMPAT_KEY: &str = "otlp_metric_compat"; pub const OTLP_METRIC_COMPAT_PROM: &str = "prom"; -pub const VALID_TABLE_OPTION_KEYS: [&str; 14] = [ +pub const VALID_TABLE_OPTION_KEYS: [&str; 15] = [ // common keys: WRITE_BUFFER_SIZE_KEY, TTL_KEY, @@ -96,6 +96,7 @@ pub const VALID_TABLE_OPTION_KEYS: [&str; 14] = [ TABLE_DATA_MODEL, OTLP_METRIC_COMPAT_KEY, REPARTITION_COLUMN_HINT_KEY, + REPARTITION_PARTITION_NUM_HINT_KEY, ]; pub const DDL_TIMEOUT: &str = "timeout"; @@ -209,6 +210,9 @@ pub const SKIP_WAL_KEY: &str = store_api::mito_engine_options::SKIP_WAL_KEY; pub const TRACE_TABLE_PARTITIONS_HINT_KEY: &str = "trace_table_partitions"; pub const REPARTITION_COLUMN_HINT_KEY: &str = "repartition.column.hint"; +/// Table-level partition count hint consumed by the auto-repartition planner. +pub const REPARTITION_PARTITION_NUM_HINT_KEY: &str = "repartition.partition.num.hint"; + impl TableOptions { pub fn try_from_iter>( iter: U, @@ -350,24 +354,26 @@ pub struct ModifyColumnTypeRequest { pub enum AnnotationFamily { /// `greptime.semantic.*` options (see the [`semantic`] module). Semantic, - /// `repartition.column.hint`, consumed by the auto-repartition planner. + /// Column and partition count hints consumed by the auto-repartition planner. RepartitionHint, } impl AnnotationFamily { - /// The key namespace: a prefix for [`Self::Semantic`], the exact key for - /// [`Self::RepartitionHint`]. + /// The key namespace used in diagnostics; accepted keys are classified by [`Self::of_key`]. pub fn namespace(self) -> &'static str { match self { Self::Semantic => SEMANTIC_PREFIX, - Self::RepartitionHint => REPARTITION_COLUMN_HINT_KEY, + Self::RepartitionHint => "repartition.", } } pub fn of_key(key: &str) -> Option { if key.starts_with(SEMANTIC_PREFIX) { Some(Self::Semantic) - } else if key == REPARTITION_COLUMN_HINT_KEY { + } else if matches!( + key, + REPARTITION_COLUMN_HINT_KEY | REPARTITION_PARTITION_NUM_HINT_KEY + ) { Some(Self::RepartitionHint) } else { None @@ -384,10 +390,8 @@ impl AnnotationFamily { } } - /// Whether this family's SET/UNSET batch must contain exactly one key. - /// The repartition hint is a single marker; a batch with several hint - /// entries (duplicates included) has no meaningful order. - pub fn requires_single_key(self) -> bool { + /// Whether duplicate keys are rejected in this family's SET/UNSET batch. + pub fn requires_unique_keys(self) -> bool { matches!(self, Self::RepartitionHint) } @@ -398,12 +402,60 @@ impl AnnotationFamily { "`{SEMANTIC_PREFIX}*` options must be altered separately from other table options" ), Self::RepartitionHint => { - format!("{REPARTITION_COLUMN_HINT_KEY} must be altered separately") + "repartition hints must be altered separately from other table options".to_string() } } } } +/// Why an annotation SET/UNSET key batch was rejected. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum AnnotationKeyError { + MixedFamilies { family: AnnotationFamily }, + DuplicateKey, +} + +impl fmt::Display for AnnotationKeyError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::MixedFamilies { family } => f.write_str(&family.mixed_batch_error()), + Self::DuplicateKey => f.write_str("duplicate repartition hint keys"), + } + } +} + +/// Validates a SET/UNSET batch and returns its annotation family. +/// +/// Empty batches and batches containing only non-annotation keys return `None`. +/// Annotation keys cannot share a batch with another family or region options. +/// Duplicate keys are rejected only for families that require unique keys. +pub fn validate_annotation_keys<'a>( + keys: impl IntoIterator, +) -> std::result::Result, AnnotationKeyError> { + let mut keys = keys.into_iter(); + let Some(first) = keys.next() else { + return Ok(None); + }; + let family = AnnotationFamily::of_key(first); + let reject_duplicates = family.is_some_and(|family| family.requires_unique_keys()); + let mut seen = HashSet::new(); + if reject_duplicates { + seen.insert(first); + } + for key in keys { + let this = AnnotationFamily::of_key(key); + if this != family + && let Some(family) = family.or(this) + { + return Err(AnnotationKeyError::MixedFamilies { family }); + } + if reject_duplicates && !seen.insert(key) { + return Err(AnnotationKeyError::DuplicateKey); + } + } + Ok(family) +} + /// Table shape an annotation option is validated against. pub struct AnnotationContext<'a> { pub schema: &'a Schema, @@ -431,6 +483,9 @@ pub enum AnnotationValidationError { column: String, ty: ConcreteDataType, }, + InvalidPartitionNumHint { + value: String, + }, NotSingleColumn, PartitionMetadataConflict, TimeIndexConflict, @@ -449,6 +504,10 @@ impl fmt::Display for AnnotationValidationError { "entity column `{column}` (option `{key}`) has type `{ty}`, \ which cannot render as a string" ), + Self::InvalidPartitionNumHint { value } => write!( + f, + "{REPARTITION_PARTITION_NUM_HINT_KEY} expects a positive integer within u32 range, got `{value}`" + ), Self::NotSingleColumn => write!( f, "{REPARTITION_COLUMN_HINT_KEY} expects exactly one column name" @@ -466,7 +525,7 @@ impl fmt::Display for AnnotationValidationError { } /// Validates one annotation option and returns the value to store — the -/// repartition hint is trimmed to the bare column name, semantic values pass +/// repartition hints are trimmed, semantic values pass /// through unchanged. pub(crate) fn validate_and_normalize_annotation( family: AnnotationFamily, @@ -505,6 +564,15 @@ pub(crate) fn validate_and_normalize_annotation( } Ok(value.to_string()) } + AnnotationFamily::RepartitionHint if key == REPARTITION_PARTITION_NUM_HINT_KEY => { + let value = value.trim(); + if !matches!(value.parse::(), Ok(1..)) { + return Err(AnnotationValidationError::InvalidPartitionNumHint { + value: value.to_string(), + }); + } + Ok(value.to_string()) + } AnnotationFamily::RepartitionHint => { let column_name = value.trim(); if column_name.is_empty() || column_name.contains(',') { @@ -763,6 +831,49 @@ mod tests { use super::*; + #[test] + fn test_validate_annotation_keys() { + let column = REPARTITION_COLUMN_HINT_KEY; + let count = REPARTITION_PARTITION_NUM_HINT_KEY; + for (keys, expected) in [ + (vec![], None), + (vec![TTL_KEY, TTL_KEY], None), + (vec!["repartition.unknown.hint"], None), + (vec![column], Some(AnnotationFamily::RepartitionHint)), + (vec![count], Some(AnnotationFamily::RepartitionHint)), + (vec![column, count], Some(AnnotationFamily::RepartitionHint)), + (vec![count, column], Some(AnnotationFamily::RepartitionHint)), + ( + vec!["greptime.semantic.source", "greptime.semantic.source"], + Some(AnnotationFamily::Semantic), + ), + ] { + assert_eq!(validate_annotation_keys(keys), Ok(expected)); + } + for keys in [ + vec![column, column], + vec![count, count], + vec![column, count, column], + ] { + assert_eq!( + validate_annotation_keys(keys), + Err(AnnotationKeyError::DuplicateKey) + ); + } + for keys in [ + vec![column, TTL_KEY], + vec![TTL_KEY, count], + vec![column, "greptime.semantic.source"], + ] { + assert_eq!( + validate_annotation_keys(keys), + Err(AnnotationKeyError::MixedFamilies { + family: AnnotationFamily::RepartitionHint, + }) + ); + } + } + #[test] fn test_validate_table_option() { assert!(validate_table_option(FILE_TABLE_LOCATION_KEY)); @@ -773,6 +884,8 @@ mod tests { assert!(validate_table_option(STORAGE_KEY)); assert!(validate_table_option(MEMTABLE_BULK_MERGE_THRESHOLD)); assert!(validate_table_option(REPARTITION_COLUMN_HINT_KEY)); + assert!(validate_table_option(REPARTITION_PARTITION_NUM_HINT_KEY)); + assert_eq!(AnnotationFamily::of_key("repartition.unknown.hint"), None); assert!(!validate_table_option("foo")); // Only whitelisted semantic keys are accepted. diff --git a/tests/cases/standalone/common/show/show_create.result b/tests/cases/standalone/common/show/show_create.result index 18879ed95e..ba2c39386d 100644 --- a/tests/cases/standalone/common/show/show_create.result +++ b/tests/cases/standalone/common/show/show_create.result @@ -180,6 +180,310 @@ DROP TABLE alter_hint_table; Affected Rows: 0 +CREATE TABLE partition_num_hint_table ( + host STRING PRIMARY KEY, + ts TIMESTAMP TIME INDEX +) +PARTITION ON COLUMNS (host) ( + host < 'm', + host >= 'm' +) +WITH ('repartition.partition.num.hint' = ' 8 '); + +Affected Rows: 0 + +SHOW CREATE TABLE partition_num_hint_table; + ++--------------------------+---------------------------------------------------------+ +| Table | Create Table | ++--------------------------+---------------------------------------------------------+ +| partition_num_hint_table | CREATE TABLE IF NOT EXISTS "partition_num_hint_table" ( | +| | "host" STRING NULL, | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | TIME INDEX ("ts"), | +| | PRIMARY KEY ("host") | +| | ) | +| | PARTITION ON COLUMNS ("host") ( | +| | host < 'm', | +| | host >= 'm' | +| | ) | +| | ENGINE=mito | +| | WITH( | +| | 'repartition.partition.num.hint' = '8' | +| | ) | ++--------------------------+---------------------------------------------------------+ + +ALTER TABLE partition_num_hint_table SET 'repartition.partition.num.hint' = 16; + +Affected Rows: 0 + +SHOW CREATE TABLE partition_num_hint_table; + ++--------------------------+---------------------------------------------------------+ +| Table | Create Table | ++--------------------------+---------------------------------------------------------+ +| partition_num_hint_table | CREATE TABLE IF NOT EXISTS "partition_num_hint_table" ( | +| | "host" STRING NULL, | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | TIME INDEX ("ts"), | +| | PRIMARY KEY ("host") | +| | ) | +| | PARTITION ON COLUMNS ("host") ( | +| | host < 'm', | +| | host >= 'm' | +| | ) | +| | ENGINE=mito | +| | WITH( | +| | 'repartition.partition.num.hint' = '16' | +| | ) | ++--------------------------+---------------------------------------------------------+ + +ALTER TABLE partition_num_hint_table SET 'repartition.partition.num.hint' = '0'; + +Error: 1004(InvalidArguments), Invalid alter table(partition_num_hint_table) request: repartition.partition.num.hint expects a positive integer within u32 range, got `0` + +ALTER TABLE partition_num_hint_table UNSET 'repartition.partition.num.hint'; + +Affected Rows: 0 + +ALTER TABLE partition_num_hint_table UNSET 'repartition.partition.num.hint'; + +Affected Rows: 0 + +SHOW CREATE TABLE partition_num_hint_table; + ++--------------------------+---------------------------------------------------------+ +| Table | Create Table | ++--------------------------+---------------------------------------------------------+ +| partition_num_hint_table | CREATE TABLE IF NOT EXISTS "partition_num_hint_table" ( | +| | "host" STRING NULL, | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | TIME INDEX ("ts"), | +| | PRIMARY KEY ("host") | +| | ) | +| | PARTITION ON COLUMNS ("host") ( | +| | host < 'm', | +| | host >= 'm' | +| | ) | +| | ENGINE=mito | +| | | ++--------------------------+---------------------------------------------------------+ + +DROP TABLE partition_num_hint_table; + +Affected Rows: 0 + +CREATE TABLE partition_num_hint_table ( + host STRING PRIMARY KEY, + ts TIMESTAMP TIME INDEX +) +WITH ('repartition.partition.num.hint' = '10'); + +Affected Rows: 0 + +SHOW CREATE TABLE partition_num_hint_table; + ++--------------------------+---------------------------------------------------------+ +| Table | Create Table | ++--------------------------+---------------------------------------------------------+ +| partition_num_hint_table | CREATE TABLE IF NOT EXISTS "partition_num_hint_table" ( | +| | "host" STRING NULL, | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | TIME INDEX ("ts"), | +| | PRIMARY KEY ("host") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | 'repartition.partition.num.hint' = '10' | +| | ) | ++--------------------------+---------------------------------------------------------+ + +DROP TABLE partition_num_hint_table; + +Affected Rows: 0 + +CREATE TABLE partition_num_hint_table ( + host STRING PRIMARY KEY, + ts TIMESTAMP TIME INDEX +) +WITH ('repartition.partition.num.hint' = 10); + +Affected Rows: 0 + +SHOW CREATE TABLE partition_num_hint_table; + ++--------------------------+---------------------------------------------------------+ +| Table | Create Table | ++--------------------------+---------------------------------------------------------+ +| partition_num_hint_table | CREATE TABLE IF NOT EXISTS "partition_num_hint_table" ( | +| | "host" STRING NULL, | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | TIME INDEX ("ts"), | +| | PRIMARY KEY ("host") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | 'repartition.partition.num.hint' = '10' | +| | ) | ++--------------------------+---------------------------------------------------------+ + +DROP TABLE partition_num_hint_table; + +Affected Rows: 0 + +CREATE TABLE combined_hint_table ( + host STRING PRIMARY KEY, + `service` STRING, + ts TIMESTAMP TIME INDEX +); + +Affected Rows: 0 + +ALTER TABLE combined_hint_table SET + 'repartition.column.hint' = 'host', + 'repartition.partition.num.hint' = 10; + +Affected Rows: 0 + +SHOW CREATE TABLE combined_hint_table; + ++---------------------+----------------------------------------------------+ +| Table | Create Table | ++---------------------+----------------------------------------------------+ +| combined_hint_table | CREATE TABLE IF NOT EXISTS "combined_hint_table" ( | +| | "host" STRING NULL, | +| | "service" STRING NULL, | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | TIME INDEX ("ts"), | +| | PRIMARY KEY ("host") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | 'repartition.column.hint' = 'host', | +| | 'repartition.partition.num.hint' = '10' | +| | ) | ++---------------------+----------------------------------------------------+ + +ALTER TABLE combined_hint_table SET + 'repartition.column.hint' = 'service', + 'repartition.partition.num.hint' = 0; + +Error: 1004(InvalidArguments), Invalid alter table(combined_hint_table) request: repartition.partition.num.hint expects a positive integer within u32 range, got `0` + +SHOW CREATE TABLE combined_hint_table; + ++---------------------+----------------------------------------------------+ +| Table | Create Table | ++---------------------+----------------------------------------------------+ +| combined_hint_table | CREATE TABLE IF NOT EXISTS "combined_hint_table" ( | +| | "host" STRING NULL, | +| | "service" STRING NULL, | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | TIME INDEX ("ts"), | +| | PRIMARY KEY ("host") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | 'repartition.column.hint' = 'host', | +| | 'repartition.partition.num.hint' = '10' | +| | ) | ++---------------------+----------------------------------------------------+ + +ALTER TABLE combined_hint_table SET + 'repartition.partition.num.hint' = 20, + 'repartition.column.hint' = 'missing'; + +Error: 4002(TableColumnNotFound), Column missing not exists in table combined_hint_table + +SHOW CREATE TABLE combined_hint_table; + ++---------------------+----------------------------------------------------+ +| Table | Create Table | ++---------------------+----------------------------------------------------+ +| combined_hint_table | CREATE TABLE IF NOT EXISTS "combined_hint_table" ( | +| | "host" STRING NULL, | +| | "service" STRING NULL, | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | TIME INDEX ("ts"), | +| | PRIMARY KEY ("host") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | 'repartition.column.hint' = 'host', | +| | 'repartition.partition.num.hint' = '10' | +| | ) | ++---------------------+----------------------------------------------------+ + +ALTER TABLE combined_hint_table SET + 'repartition.partition.num.hint' = 20, + 'repartition.column.hint' = 'service'; + +Affected Rows: 0 + +SHOW CREATE TABLE combined_hint_table; + ++---------------------+----------------------------------------------------+ +| Table | Create Table | ++---------------------+----------------------------------------------------+ +| combined_hint_table | CREATE TABLE IF NOT EXISTS "combined_hint_table" ( | +| | "host" STRING NULL, | +| | "service" STRING NULL, | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | TIME INDEX ("ts"), | +| | PRIMARY KEY ("host") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | 'repartition.column.hint' = 'service', | +| | 'repartition.partition.num.hint' = '20' | +| | ) | ++---------------------+----------------------------------------------------+ + +ALTER TABLE combined_hint_table SET + 'repartition.partition.num.hint' = 30, + 'repartition.partition.num.hint' = 40; + +Error: 1004(InvalidArguments), Invalid table option request: duplicate repartition hint keys + +ALTER TABLE combined_hint_table SET + 'repartition.partition.num.hint' = 30, + 'ttl' = '7d'; + +Error: 1004(InvalidArguments), Invalid table option request: repartition hints must be altered separately from other table options + +ALTER TABLE combined_hint_table UNSET + 'repartition.column.hint', + 'repartition.partition.num.hint'; + +Affected Rows: 0 + +SHOW CREATE TABLE combined_hint_table; + ++---------------------+----------------------------------------------------+ +| Table | Create Table | ++---------------------+----------------------------------------------------+ +| combined_hint_table | CREATE TABLE IF NOT EXISTS "combined_hint_table" ( | +| | "host" STRING NULL, | +| | "service" STRING NULL, | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | TIME INDEX ("ts"), | +| | PRIMARY KEY ("host") | +| | ) | +| | | +| | ENGINE=mito | +| | | ++---------------------+----------------------------------------------------+ + +DROP TABLE combined_hint_table; + +Affected Rows: 0 + CREATE TABLE not_supported_table_storage_option ( `id` INT UNSIGNED, host STRING, diff --git a/tests/cases/standalone/common/show/show_create.sql b/tests/cases/standalone/common/show/show_create.sql index fd8b1c16ea..5e86273091 100644 --- a/tests/cases/standalone/common/show/show_create.sql +++ b/tests/cases/standalone/common/show/show_create.sql @@ -61,6 +61,98 @@ SHOW CREATE TABLE alter_hint_table; DROP TABLE alter_hint_table; +CREATE TABLE partition_num_hint_table ( + host STRING PRIMARY KEY, + ts TIMESTAMP TIME INDEX +) +PARTITION ON COLUMNS (host) ( + host < 'm', + host >= 'm' +) +WITH ('repartition.partition.num.hint' = ' 8 '); + +SHOW CREATE TABLE partition_num_hint_table; + +ALTER TABLE partition_num_hint_table SET 'repartition.partition.num.hint' = 16; + +SHOW CREATE TABLE partition_num_hint_table; + +ALTER TABLE partition_num_hint_table SET 'repartition.partition.num.hint' = '0'; + +ALTER TABLE partition_num_hint_table UNSET 'repartition.partition.num.hint'; + +ALTER TABLE partition_num_hint_table UNSET 'repartition.partition.num.hint'; + +SHOW CREATE TABLE partition_num_hint_table; + +DROP TABLE partition_num_hint_table; + +CREATE TABLE partition_num_hint_table ( + host STRING PRIMARY KEY, + ts TIMESTAMP TIME INDEX +) +WITH ('repartition.partition.num.hint' = '10'); + +SHOW CREATE TABLE partition_num_hint_table; + +DROP TABLE partition_num_hint_table; + +CREATE TABLE partition_num_hint_table ( + host STRING PRIMARY KEY, + ts TIMESTAMP TIME INDEX +) +WITH ('repartition.partition.num.hint' = 10); + +SHOW CREATE TABLE partition_num_hint_table; + +DROP TABLE partition_num_hint_table; + +CREATE TABLE combined_hint_table ( + host STRING PRIMARY KEY, + `service` STRING, + ts TIMESTAMP TIME INDEX +); + +ALTER TABLE combined_hint_table SET + 'repartition.column.hint' = 'host', + 'repartition.partition.num.hint' = 10; + +SHOW CREATE TABLE combined_hint_table; + +ALTER TABLE combined_hint_table SET + 'repartition.column.hint' = 'service', + 'repartition.partition.num.hint' = 0; + +SHOW CREATE TABLE combined_hint_table; + +ALTER TABLE combined_hint_table SET + 'repartition.partition.num.hint' = 20, + 'repartition.column.hint' = 'missing'; + +SHOW CREATE TABLE combined_hint_table; + +ALTER TABLE combined_hint_table SET + 'repartition.partition.num.hint' = 20, + 'repartition.column.hint' = 'service'; + +SHOW CREATE TABLE combined_hint_table; + +ALTER TABLE combined_hint_table SET + 'repartition.partition.num.hint' = 30, + 'repartition.partition.num.hint' = 40; + +ALTER TABLE combined_hint_table SET + 'repartition.partition.num.hint' = 30, + 'ttl' = '7d'; + +ALTER TABLE combined_hint_table UNSET + 'repartition.column.hint', + 'repartition.partition.num.hint'; + +SHOW CREATE TABLE combined_hint_table; + +DROP TABLE combined_hint_table; + CREATE TABLE not_supported_table_storage_option ( `id` INT UNSIGNED, host STRING,