1use std::collections::{HashMap, HashSet};
16use std::sync::Arc;
17use std::time::Duration;
18
19use api::helper::ColumnDataTypeWrapper;
20use api::v1::alter_table_expr::Kind;
21use api::v1::meta::CreateFlowTask as PbCreateFlowTask;
22use api::v1::repartition::Source;
23use api::v1::{
24 AlterDatabaseExpr, AlterTableExpr, CreateFlowExpr, CreateTableExpr, CreateViewExpr,
25 PartitionedSource, Repartition, TargetPartitionColumns, UnpartitionedSource, column_def,
26};
27#[cfg(feature = "enterprise")]
28use api::v1::{
29 CreateTriggerExpr as PbCreateTriggerExpr, meta::CreateTriggerTask as PbCreateTriggerTask,
30};
31use catalog::CatalogManagerRef;
32use chrono::Utc;
33use common_base::regex_pattern::NAME_PATTERN_REG;
34use common_catalog::consts::{
35 DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, is_ddl_reserved_table, is_readonly_schema,
36 is_readonly_table,
37};
38use common_catalog::{format_full_flow_name, format_full_table_name};
39use common_error::ext::BoxedError;
40#[cfg(feature = "enterprise")]
41use common_meta::cache_invalidator::CacheInvalidatorRef;
42use common_meta::cache_invalidator::Context;
43use common_meta::ddl::create_flow::{
44 DEFER_ON_MISSING_SOURCE_KEY, FLOW_EXPERIMENTAL_ENABLE_INCREMENTAL_READ_KEY, FlowType,
45 INTERNAL_EVAL_OFFSET_KEY, INTERNAL_EVAL_SCHEDULE_KEY,
46};
47use common_meta::instruction::CacheIdent;
48#[cfg(feature = "enterprise")]
49use common_meta::key::TableMetadataManagerRef;
50use common_meta::key::schema_name::{SchemaName, SchemaNameKey};
51#[cfg(feature = "enterprise")]
52use common_meta::procedure_executor::ProcedureExecutorRef;
53#[cfg(feature = "enterprise")]
54use common_meta::rpc::ddl::trigger::CreateTriggerTask;
55#[cfg(feature = "enterprise")]
56use common_meta::rpc::ddl::trigger::DropTriggerTask;
57use common_meta::rpc::ddl::{
58 CreateFlowTask, CreatorGrantIntent, DdlTask, DropFlowTask, DropViewTask, SubmitDdlTaskRequest,
59 SubmitDdlTaskResponse, TriggerReason,
60};
61use common_query::Output;
62use common_recordbatch::{RecordBatch, RecordBatches};
63use common_sql::convert::sql_value_to_value;
64use common_telemetry::{debug, info, tracing, warn};
65use common_time::{Timestamp, Timezone};
66use datafusion_common::tree_node::TreeNodeVisitor;
67use datafusion_expr::LogicalPlan;
68use datafusion_expr::logical_plan::Distinct;
69use datatypes::prelude::ConcreteDataType;
70use datatypes::schema::{ColumnSchema, Schema};
71use datatypes::value::Value;
72use datatypes::vectors::{StringVector, VectorRef};
73use humantime::parse_duration;
74use partition::expr::{Operand, PartitionExpr, RestrictedOp};
75use partition::multi_dim::MultiDimPartitionRule;
76use query::parser::QueryStatement;
77use query::plan::extract_and_rewrite_full_table_names;
78use query::query_engine::DefaultSerializer;
79use query::sql::create_table_stmt;
80use session::context::QueryContextRef;
81use session::table_name::table_idents_to_full_name;
82use snafu::{OptionExt, ResultExt, ensure};
83use sql::parser::{ParseOptions, ParserContext};
84use sql::parsers::utils::is_tql;
85use sql::statements::OptionMap;
86#[cfg(feature = "enterprise")]
87use sql::statements::alter::trigger::AlterTrigger;
88use sql::statements::alter::{AlterDatabase, AlterTable, AlterTableOperation};
89#[cfg(feature = "enterprise")]
90use sql::statements::create::trigger::CreateTrigger;
91use sql::statements::create::{
92 CreateExternalTable, CreateFlow, CreateTable, CreateTableLike, CreateView, Partitions,
93};
94use sql::statements::statement::Statement;
95use sqlparser::ast::{Expr, Ident, UnaryOperator, Value as ParserValue};
96use store_api::metric_engine_consts::{LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME};
97use store_api::mito_engine_options::APPEND_MODE_KEY;
98use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan};
99use table::TableRef;
100use table::dist_table::DistTable;
101use table::metadata::{self, TableId, TableInfo, TableMeta, TableType};
102use table::requests::{
103 AlterKind, AlterTableRequest, AnnotationContext, COMMENT_KEY, DDL_TIMEOUT, DDL_WAIT,
104 TableOptions, validate_and_normalize_annotation_options,
105};
106use table::table_name::TableName;
107use table::table_reference::TableReference;
108
109use crate::error::{
110 self, AlterExprToRequestSnafu, BuildDfLogicalPlanSnafu, CatalogSnafu, ColumnDataTypeSnafu,
111 ColumnNotFoundSnafu, ConvertSchemaSnafu, CreateLogicalTablesSnafu,
112 DeserializePartitionExprSnafu, EmptyDdlExprSnafu, ExternalSnafu, ExtractTableNamesSnafu,
113 FlowNotFoundSnafu, InvalidPartitionRuleSnafu, InvalidPartitionSnafu, InvalidSqlSnafu,
114 InvalidTableNameSnafu, InvalidViewNameSnafu, InvalidViewStmtSnafu, NotSupportedSnafu,
115 PartitionExprToPbSnafu, Result, SchemaInUseSnafu, SchemaNotFoundSnafu, SchemaReadOnlySnafu,
116 SerializePartitionExprSnafu, SubstraitCodecSnafu, TableAlreadyExistsSnafu,
117 TableDdlReservedSnafu, TableMetadataManagerSnafu, TableNotFoundSnafu, TableReadOnlySnafu,
118 UnrecognizedTableOptionSnafu, ViewAlreadyExistsSnafu,
119};
120use crate::expr_helper::{self, RepartitionRequest, RepartitionSource};
121use crate::statement::StatementExecutor;
122use crate::statement::show::create_partitions_stmt;
123use crate::utils::{to_executor_context, to_executor_context_with_origin_frontend};
124
125#[derive(Debug, Clone, Copy)]
126struct DdlSubmitOptions {
127 wait: bool,
128 timeout: Duration,
129}
130
131#[cfg(not(feature = "enterprise"))]
132const ALLOWED_FLOW_OPTIONS: [&str; 2] = [
133 DEFER_ON_MISSING_SOURCE_KEY,
134 FLOW_EXPERIMENTAL_ENABLE_INCREMENTAL_READ_KEY,
135];
136
137fn build_procedure_id_output(procedure_id: Vec<u8>) -> Result<Output> {
138 let procedure_id = String::from_utf8_lossy(&procedure_id).to_string();
139 let vector: VectorRef = Arc::new(StringVector::from(vec![procedure_id]));
140 let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
141 "Procedure ID",
142 vector.data_type(),
143 false,
144 )]));
145 let batch =
146 RecordBatch::new(schema.clone(), vec![vector]).context(error::BuildRecordBatchSnafu)?;
147 let batches =
148 RecordBatches::try_new(schema, vec![batch]).context(error::BuildRecordBatchSnafu)?;
149 Ok(Output::new_with_record_batches(batches))
150}
151
152fn parse_ddl_options(options: &OptionMap) -> Result<DdlSubmitOptions> {
153 let wait = match options.get(DDL_WAIT) {
154 Some(value) => value.parse::<bool>().map_err(|_| {
155 InvalidSqlSnafu {
156 err_msg: format!("invalid DDL option '{DDL_WAIT}': '{value}'"),
157 }
158 .build()
159 })?,
160 None => SubmitDdlTaskRequest::default_wait(),
161 };
162
163 let timeout = match options.get(DDL_TIMEOUT) {
164 Some(value) => parse_duration(value).map_err(|err| {
165 InvalidSqlSnafu {
166 err_msg: format!("invalid DDL option '{DDL_TIMEOUT}': '{value}': {err}"),
167 }
168 .build()
169 })?,
170 None => SubmitDdlTaskRequest::default_timeout(),
171 };
172
173 Ok(DdlSubmitOptions { wait, timeout })
174}
175
176#[cfg(not(feature = "enterprise"))]
177fn supported_flow_options() -> String {
178 ALLOWED_FLOW_OPTIONS.join(", ")
179}
180
181fn normalize_flow_bool_option(key: &str, value: &str) -> Result<String> {
182 value
183 .trim()
184 .to_ascii_lowercase()
185 .parse::<bool>()
186 .map(|value| value.to_string())
187 .map_err(|_| {
188 InvalidSqlSnafu {
189 err_msg: format!("invalid flow option '{key}': '{value}'"),
190 }
191 .build()
192 })
193}
194
195fn validate_and_normalize_flow_options(
196 options: HashMap<String, String>,
197 eval_interval: Option<i64>,
198) -> Result<HashMap<String, String>> {
199 if let Some(secs) = eval_interval
201 && secs <= 0
202 {
203 return InvalidSqlSnafu {
204 err_msg: format!("EVAL INTERVAL must be positive, got {secs} seconds"),
205 }
206 .fail();
207 }
208
209 options
210 .into_iter()
211 .map(|(key, value)| {
212 if matches!(
213 key.as_str(),
214 FlowType::FLOW_TYPE_KEY | INTERNAL_EVAL_OFFSET_KEY | INTERNAL_EVAL_SCHEDULE_KEY
215 ) {
216 return InvalidSqlSnafu {
217 err_msg: format!("flow option '{key}' is reserved for internal use"),
218 }
219 .fail();
220 }
221
222 let normalized_value = match key.as_str() {
223 DEFER_ON_MISSING_SOURCE_KEY | FLOW_EXPERIMENTAL_ENABLE_INCREMENTAL_READ_KEY => {
224 normalize_flow_bool_option(&key, &value)?
225 }
226 #[cfg(feature = "enterprise")]
227 _ => value,
228 #[cfg(not(feature = "enterprise"))]
229 _ => {
230 return InvalidSqlSnafu {
231 err_msg: format!(
232 "unknown flow option '{key}', supported options: {}",
233 supported_flow_options()
234 ),
235 }
236 .fail();
237 }
238 };
239
240 Ok((key, normalized_value))
241 })
242 .collect()
243}
244
245fn validate_eval_offset(
249 eval_offset_secs: Option<i64>,
250 eval_interval_secs: Option<i64>,
251) -> Result<()> {
252 if let Some(offset_secs) = eval_offset_secs {
253 let Some(eval_interval_secs) = eval_interval_secs else {
254 return InvalidSqlSnafu {
255 err_msg: "EVAL OFFSET requires EVAL INTERVAL to be specified".to_string(),
256 }
257 .fail();
258 };
259 if !(0..eval_interval_secs).contains(&offset_secs) {
260 return InvalidSqlSnafu {
261 err_msg: format!(
262 "EVAL OFFSET must be in range [0, EVAL INTERVAL), got {offset_secs} seconds with EVAL INTERVAL {eval_interval_secs} seconds"
263 ),
264 }
265 .fail();
266 }
267 }
268 Ok(())
269}
270
271fn determine_flow_type_for_source_state(
272 flow_name: &str,
273 flow_options: &HashMap<String, String>,
274 has_missing_source_table: bool,
275 has_instant_ttl_source_table: bool,
276 force_batching: bool,
277) -> Result<Option<FlowType>> {
278 if has_instant_ttl_source_table && force_batching {
279 return InvalidSqlSnafu {
282 err_msg: format!(
283 "flow '{}' with EVAL INTERVAL requires the batching scheduler, but source tables with ttl=instant are not supported under batching mode; use a TTL longer than the flush interval",
284 flow_name
285 ),
286 }
287 .fail();
288 }
289
290 if has_missing_source_table {
291 let defer_on_missing_source = flow_options
292 .get(DEFER_ON_MISSING_SOURCE_KEY)
293 .is_some_and(|value| value == "true");
294 ensure!(
295 defer_on_missing_source,
296 InvalidSqlSnafu {
297 err_msg: format!(
298 "missing source tables for flow '{}'; use WITH ({DEFER_ON_MISSING_SOURCE_KEY} = true) to create a pending flow",
299 flow_name
300 )
301 }
302 );
303 info!(
304 "Flow `{}` is created as a pending batching flow because source tables are missing and defer_on_missing_source=true",
305 flow_name
306 );
307 return Ok(Some(FlowType::Batching));
308 }
309
310 Ok(None)
311}
312
313fn is_stateless_flow_plan(plan: &LogicalPlan) -> bool {
317 let mut scans = 0;
318 let mut supported_nodes = true;
319 let result = plan.apply_with_subqueries(|node| {
320 match node {
321 LogicalPlan::TableScan(_) => scans += 1,
322 LogicalPlan::Projection(_) | LogicalPlan::Filter(_) => {}
323 LogicalPlan::Distinct(Distinct::All(_)) => {}
324 _ => supported_nodes = false,
325 }
326 Ok(datafusion_common::tree_node::TreeNodeRecursion::Continue)
327 });
328
329 result.is_ok() && supported_nodes && scans == 1
330}
331
332const INSTANT_TTL_FLOW_QUERY_ERROR: &str = "instant-TTL flow sources support only stateless projection/filter queries over one source; non-stateless queries require a persisted source";
333
334const STATELESS_FLOW_QUERY_ERROR: &str =
335 "flow streaming supports only stateless projection/filter queries over one source";
336
337impl StatementExecutor {
338 pub fn catalog_manager(&self) -> CatalogManagerRef {
339 self.catalog_manager.clone()
340 }
341
342 #[tracing::instrument(skip_all)]
343 pub async fn create_table(&self, stmt: CreateTable, ctx: QueryContextRef) -> Result<TableRef> {
344 let (catalog, schema, _table) = table_idents_to_full_name(&stmt.name, &ctx)
345 .map_err(BoxedError::new)
346 .context(error::ExternalSnafu)?;
347
348 let schema_options = self
349 .table_metadata_manager
350 .schema_manager()
351 .get(SchemaNameKey {
352 catalog: &catalog,
353 schema: &schema,
354 })
355 .await
356 .context(TableMetadataManagerSnafu)?
357 .map(|v| v.into_inner());
358
359 let create_expr = &mut expr_helper::create_to_expr(&stmt, &ctx)?;
360 if let Some(schema_options) = schema_options {
363 for (key, value) in schema_options.extra_options.iter() {
364 if key.starts_with("compaction.") {
365 continue;
366 }
367 create_expr
368 .table_options
369 .entry(key.clone())
370 .or_insert(value.clone());
371 }
372 }
373
374 self.create_table_inner(create_expr, stmt.partitions, ctx, TriggerReason::Manual)
375 .await
376 }
377
378 #[tracing::instrument(skip_all)]
379 pub async fn create_table_like(
380 &self,
381 stmt: CreateTableLike,
382 ctx: QueryContextRef,
383 ) -> Result<TableRef> {
384 let (catalog, schema, table) = table_idents_to_full_name(&stmt.source_name, &ctx)
385 .map_err(BoxedError::new)
386 .context(error::ExternalSnafu)?;
387 let table_ref = self
388 .catalog_manager
389 .table(&catalog, &schema, &table, Some(&ctx))
390 .await
391 .context(CatalogSnafu)?
392 .context(TableNotFoundSnafu { table_name: &table })?;
393 let partition_info = self
394 .partition_manager
395 .find_physical_partition_info(table_ref.table_info().table_id())
396 .await
397 .context(error::FindTablePartitionRuleSnafu { table_name: table })?;
398
399 let schema_options = self
401 .table_metadata_manager
402 .schema_manager()
403 .get(SchemaNameKey {
404 catalog: &catalog,
405 schema: &schema,
406 })
407 .await
408 .context(TableMetadataManagerSnafu)?
409 .map(|v| v.into_inner());
410
411 let quote_style = ctx.quote_style();
412 let mut create_stmt =
413 create_table_stmt(&table_ref.table_info(), schema_options, quote_style)
414 .context(error::ParseQuerySnafu)?;
415 create_stmt.name = stmt.table_name;
416 create_stmt.if_not_exists = false;
417
418 let table_info = table_ref.table_info();
419 let partitions = create_partitions_stmt(&table_info, &partition_info.partitions)?.and_then(
420 |mut partitions| {
421 if !partitions.column_list.is_empty() {
422 partitions.set_quote(quote_style);
423 Some(partitions)
424 } else {
425 None
426 }
427 },
428 );
429
430 let create_expr = &mut expr_helper::create_to_expr(&create_stmt, &ctx)?;
431 self.create_table_inner(create_expr, partitions, ctx, TriggerReason::Manual)
432 .await
433 }
434
435 #[tracing::instrument(skip_all)]
436 pub async fn create_external_table(
437 &self,
438 create_expr: CreateExternalTable,
439 ctx: QueryContextRef,
440 ) -> Result<TableRef> {
441 let create_expr =
442 &mut expr_helper::create_external_expr(create_expr, &ctx, &self.local_file_access)
443 .await?;
444 self.create_table_inner(create_expr, None, ctx, TriggerReason::Manual)
445 .await
446 }
447
448 pub async fn create_declared_relationships_table(
452 &self,
453 catalog: &str,
454 query_ctx: QueryContextRef,
455 ) -> Result<TableRef> {
456 let mut expr = super::semantic_graph::build_declared_relationships_expr(catalog);
457 self.create_non_logic_table(&mut expr, None, query_ctx, TriggerReason::AutoCreate)
458 .await
459 }
460
461 #[tracing::instrument(skip_all)]
462 pub async fn create_table_inner(
463 &self,
464 create_table: &mut CreateTableExpr,
465 partitions: Option<Partitions>,
466 query_ctx: QueryContextRef,
467 trigger_reason: TriggerReason,
468 ) -> Result<TableRef> {
469 ensure_table_definition_writable(&create_table.schema_name, &create_table.table_name)?;
470
471 if create_table.engine == METRIC_ENGINE_NAME
472 && create_table
473 .table_options
474 .contains_key(LOGICAL_TABLE_METADATA_KEY)
475 {
476 if let Some(partitions) = partitions.as_ref()
477 && !partitions.exprs.is_empty()
478 {
479 self.validate_logical_table_partition_rule(create_table, partitions, &query_ctx)
480 .await?;
481 }
482 self.create_logical_tables(
484 std::slice::from_ref(create_table),
485 query_ctx,
486 trigger_reason,
487 )
488 .await?
489 .into_iter()
490 .next()
491 .context(error::UnexpectedSnafu {
492 violated: "expected to create logical tables",
493 })
494 } else {
495 self.create_non_logic_table(create_table, partitions, query_ctx, trigger_reason)
497 .await
498 }
499 }
500
501 #[tracing::instrument(skip_all)]
502 pub async fn create_non_logic_table(
503 &self,
504 create_table: &mut CreateTableExpr,
505 partitions: Option<Partitions>,
506 query_ctx: QueryContextRef,
507 trigger_reason: TriggerReason,
508 ) -> Result<TableRef> {
509 let _timer = crate::metrics::DIST_CREATE_TABLE.start_timer();
510
511 let schema = self
513 .table_metadata_manager
514 .schema_manager()
515 .get(SchemaNameKey::new(
516 &create_table.catalog_name,
517 &create_table.schema_name,
518 ))
519 .await
520 .context(TableMetadataManagerSnafu)?;
521 ensure!(
522 schema.is_some(),
523 SchemaNotFoundSnafu {
524 schema_info: &create_table.schema_name,
525 }
526 );
527
528 if let Some(table) = self
530 .catalog_manager
531 .table(
532 &create_table.catalog_name,
533 &create_table.schema_name,
534 &create_table.table_name,
535 Some(&query_ctx),
536 )
537 .await
538 .context(CatalogSnafu)?
539 {
540 return if create_table.create_if_not_exists {
541 Ok(table)
542 } else {
543 TableAlreadyExistsSnafu {
544 table: format_full_table_name(
545 &create_table.catalog_name,
546 &create_table.schema_name,
547 &create_table.table_name,
548 ),
549 }
550 .fail()
551 };
552 }
553
554 ensure!(
555 NAME_PATTERN_REG.is_match(&create_table.table_name),
556 InvalidTableNameSnafu {
557 table_name: &create_table.table_name,
558 }
559 );
560
561 let table_name = TableName::new(
562 &create_table.catalog_name,
563 &create_table.schema_name,
564 &create_table.table_name,
565 );
566
567 let (partitions, partition_cols) = parse_partitions(create_table, partitions, &query_ctx)?;
568 let mut table_info = create_table_info(create_table, partition_cols)?;
569
570 let resp = self
571 .create_table_procedure(
572 create_table.clone(),
573 partitions,
574 table_info.clone(),
575 query_ctx,
576 trigger_reason,
577 )
578 .await?;
579
580 let table_id = resp
581 .table_ids
582 .into_iter()
583 .next()
584 .context(error::UnexpectedSnafu {
585 violated: "expected table_id",
586 })?;
587 info!("Successfully created table '{table_name}' with table id {table_id}");
588
589 table_info.ident.table_id = table_id;
590
591 let table_info = Arc::new(table_info);
592 create_table.table_id = Some(api::v1::TableId { id: table_id });
593
594 let table = DistTable::table(table_info);
595
596 Ok(table)
597 }
598
599 #[tracing::instrument(skip_all)]
600 pub async fn create_logical_tables(
601 &self,
602 create_table_exprs: &[CreateTableExpr],
603 query_context: QueryContextRef,
604 trigger_reason: TriggerReason,
605 ) -> Result<Vec<TableRef>> {
606 let _timer = crate::metrics::DIST_CREATE_TABLES.start_timer();
607 ensure!(
608 !create_table_exprs.is_empty(),
609 EmptyDdlExprSnafu {
610 name: "create logic tables"
611 }
612 );
613
614 for create_table in create_table_exprs {
616 ensure!(
617 NAME_PATTERN_REG.is_match(&create_table.table_name),
618 InvalidTableNameSnafu {
619 table_name: &create_table.table_name,
620 }
621 );
622 }
623
624 let raw_tables_info = create_table_exprs
625 .iter()
626 .map(|create| create_table_info(create, vec![]))
627 .collect::<Result<Vec<_>>>()?;
628 let tables_data = create_table_exprs
629 .iter()
630 .cloned()
631 .zip(raw_tables_info.iter().cloned())
632 .collect::<Vec<_>>();
633
634 let resp = self
635 .create_logical_tables_procedure(tables_data, query_context.clone(), trigger_reason)
636 .await?;
637
638 let table_ids = resp.table_ids;
639 ensure!(
640 table_ids.len() == raw_tables_info.len(),
641 CreateLogicalTablesSnafu {
642 reason: format!(
643 "The number of tables is inconsistent with the expected number to be created, expected: {}, actual: {}",
644 raw_tables_info.len(),
645 table_ids.len()
646 )
647 }
648 );
649 info!("Successfully created logical tables: {:?}", table_ids);
650
651 let mut tables_info = Vec::with_capacity(table_ids.len());
655 for (table_id, create_table) in table_ids.iter().zip(create_table_exprs.iter()) {
656 let table = self
657 .catalog_manager
658 .table(
659 &create_table.catalog_name,
660 &create_table.schema_name,
661 &create_table.table_name,
662 Some(&query_context),
663 )
664 .await
665 .context(CatalogSnafu)?
666 .with_context(|| TableNotFoundSnafu {
667 table_name: format_full_table_name(
668 &create_table.catalog_name,
669 &create_table.schema_name,
670 &create_table.table_name,
671 ),
672 })?;
673
674 let table_info = table.table_info();
675 ensure!(
677 table_info.table_id() == *table_id,
678 CreateLogicalTablesSnafu {
679 reason: format!(
680 "Table id mismatch after creation, expected {}, got {} for table {}",
681 table_id,
682 table_info.table_id(),
683 format_full_table_name(
684 &create_table.catalog_name,
685 &create_table.schema_name,
686 &create_table.table_name
687 )
688 )
689 }
690 );
691
692 tables_info.push(table_info);
693 }
694
695 Ok(tables_info.into_iter().map(DistTable::table).collect())
696 }
697
698 async fn validate_logical_table_partition_rule(
699 &self,
700 create_table: &CreateTableExpr,
701 partitions: &Partitions,
702 query_ctx: &QueryContextRef,
703 ) -> Result<()> {
704 let (_, mut logical_partition_exprs) =
705 parse_partitions_for_logical_validation(create_table, partitions, query_ctx)?;
706
707 let physical_table_name = create_table
708 .table_options
709 .get(LOGICAL_TABLE_METADATA_KEY)
710 .with_context(|| CreateLogicalTablesSnafu {
711 reason: format!(
712 "expect `{LOGICAL_TABLE_METADATA_KEY}` option on creating logical table"
713 ),
714 })?;
715
716 let physical_table = self
717 .catalog_manager
718 .table(
719 &create_table.catalog_name,
720 &create_table.schema_name,
721 physical_table_name,
722 Some(query_ctx),
723 )
724 .await
725 .context(CatalogSnafu)?
726 .context(TableNotFoundSnafu {
727 table_name: physical_table_name.clone(),
728 })?;
729
730 let physical_table_info = physical_table.table_info();
731 let (partition_rule, _) = self
732 .partition_manager
733 .find_table_partition_rule(&physical_table_info)
734 .await
735 .context(error::FindTablePartitionRuleSnafu {
736 table_name: physical_table_name.clone(),
737 })?;
738
739 let multi_dim_rule = partition_rule
740 .as_ref()
741 .as_any()
742 .downcast_ref::<MultiDimPartitionRule>()
743 .context(InvalidPartitionRuleSnafu {
744 reason: "physical table partition rule is not range-based",
745 })?;
746
747 let mut physical_partition_exprs = multi_dim_rule.exprs().to_vec();
749 logical_partition_exprs.sort_unstable();
750 physical_partition_exprs.sort_unstable();
751
752 ensure!(
753 physical_partition_exprs == logical_partition_exprs,
754 InvalidPartitionRuleSnafu {
755 reason: format!(
756 "logical table partition rule must match the corresponding physical table's\n logical table partition exprs:\t\t {:?}\n physical table partition exprs:\t {:?}",
757 logical_partition_exprs, physical_partition_exprs
758 ),
759 }
760 );
761
762 Ok(())
763 }
764
765 #[cfg(feature = "enterprise")]
766 #[tracing::instrument(skip_all)]
767 pub async fn create_trigger(
768 &self,
769 stmt: CreateTrigger,
770 query_context: QueryContextRef,
771 ) -> Result<Output> {
772 let expr = expr_helper::to_create_trigger_task_expr(stmt, &query_context)?;
773 self.create_trigger_inner(expr, query_context).await
774 }
775
776 #[cfg(feature = "enterprise")]
777 pub async fn create_trigger_inner(
778 &self,
779 expr: PbCreateTriggerExpr,
780 query_context: QueryContextRef,
781 ) -> Result<Output> {
782 self.create_trigger_procedure(expr, query_context).await?;
783 Ok(Output::new_with_affected_rows(0))
784 }
785
786 #[cfg(feature = "enterprise")]
787 async fn create_trigger_procedure(
788 &self,
789 expr: PbCreateTriggerExpr,
790 query_context: QueryContextRef,
791 ) -> Result<SubmitDdlTaskResponse> {
792 let task = CreateTriggerTask::try_from(PbCreateTriggerTask {
793 create_trigger: Some(expr),
794 })
795 .context(error::InvalidExprSnafu)?;
796
797 let executor_context = to_executor_context(query_context, TriggerReason::Manual);
798 let request = SubmitDdlTaskRequest::new(DdlTask::new_create_trigger(task));
799
800 self.procedure_executor
801 .submit_ddl_task(executor_context, request)
802 .await
803 .context(error::ExecuteDdlSnafu)
804 }
805
806 #[tracing::instrument(skip_all)]
807 pub async fn create_flow(
808 &self,
809 stmt: CreateFlow,
810 query_context: QueryContextRef,
811 ) -> Result<Output> {
812 let eval_offset_secs = stmt.eval_offset;
814 let expr = expr_helper::to_create_flow_task_expr(stmt, &query_context)?;
815
816 self.create_flow_procedure(expr, eval_offset_secs, query_context)
819 .await?;
820 Ok(Output::new_with_affected_rows(0))
821 }
822
823 pub async fn create_flow_inner(
825 &self,
826 expr: CreateFlowExpr,
827 query_context: QueryContextRef,
828 ) -> Result<Output> {
829 self.create_flow_procedure(expr, None, query_context)
830 .await?;
831 Ok(Output::new_with_affected_rows(0))
832 }
833
834 async fn create_flow_procedure(
835 &self,
836 mut expr: CreateFlowExpr,
837 eval_offset_secs: Option<i64>,
838 query_context: QueryContextRef,
839 ) -> Result<SubmitDdlTaskResponse> {
840 let eval_interval_secs = expr.eval_interval.as_ref().map(|e| e.seconds);
841
842 if let Some(secs) = eval_interval_secs
844 && secs <= 0
845 {
846 return InvalidSqlSnafu {
847 err_msg: format!("EVAL INTERVAL must be positive, got {secs} seconds"),
848 }
849 .fail();
850 }
851
852 validate_eval_offset(eval_offset_secs, eval_interval_secs)?;
856
857 expr.flow_options =
861 validate_and_normalize_flow_options(expr.flow_options, eval_interval_secs)?;
862
863 let flow_type = self
864 .determine_flow_type(&expr, query_context.clone())
865 .await?;
866 info!("determined flow={} type: {:#?}", expr.flow_name, flow_type);
867
868 expr.flow_options
869 .insert(FlowType::FLOW_TYPE_KEY.to_string(), flow_type.to_string());
870
871 let mut task = CreateFlowTask::try_from(PbCreateFlowTask {
872 create_flow: Some(expr),
873 })
874 .context(error::InvalidExprSnafu)?;
875 task.eval_offset_secs = eval_offset_secs;
876 let executor_context = to_executor_context(query_context, TriggerReason::Manual);
877 let request = SubmitDdlTaskRequest::new(DdlTask::new_create_flow(task));
878
879 self.procedure_executor
880 .submit_ddl_task(executor_context, request)
881 .await
882 .context(error::ExecuteDdlSnafu)
883 }
884
885 async fn determine_flow_type(
891 &self,
892 expr: &CreateFlowExpr,
893 query_ctx: QueryContextRef,
894 ) -> Result<FlowType> {
895 let has_eval_interval = expr.eval_interval.is_some();
896
897 let mut has_missing_source_table = false;
898 let mut has_instant_ttl_source_table = false;
899
900 for src_table_name in &expr.source_table_names {
901 let table = self
902 .catalog_manager()
903 .table(
904 &src_table_name.catalog_name,
905 &src_table_name.schema_name,
906 &src_table_name.table_name,
907 Some(&query_ctx),
908 )
909 .await
910 .map_err(BoxedError::new)
911 .context(ExternalSnafu)?;
912
913 let Some(table) = table else {
914 has_missing_source_table = true;
915 continue;
916 };
917
918 if table.table_info().meta.options.ttl == Some(common_time::TimeToLive::Instant) {
919 warn!(
920 "Source table `{}` for flow `{}`'s ttl=instant, {}",
921 format_full_table_name(
922 &src_table_name.catalog_name,
923 &src_table_name.schema_name,
924 &src_table_name.table_name
925 ),
926 expr.flow_name,
927 if has_eval_interval {
928 "rejecting flow because EVAL INTERVAL requires the batching scheduler which does not support instant TTL"
929 } else {
930 "fallback to streaming mode"
931 }
932 );
933 has_instant_ttl_source_table = true;
934 }
935 }
936
937 if let Some(flow_type) = determine_flow_type_for_source_state(
938 &expr.flow_name,
939 &expr.flow_options,
940 has_missing_source_table,
941 has_instant_ttl_source_table,
942 has_eval_interval,
943 )? {
944 return Ok(flow_type);
945 }
946
947 if has_eval_interval {
953 return Ok(FlowType::Batching);
954 }
955
956 let engine = &self.query_engine;
957 let stmts = ParserContext::create_with_dialect(
958 &expr.sql,
959 query_ctx.sql_dialect(),
960 ParseOptions::default(),
961 )
962 .map_err(BoxedError::new)
963 .context(ExternalSnafu)?;
964
965 ensure!(
966 stmts.len() == 1,
967 InvalidSqlSnafu {
968 err_msg: format!("Expect only one statement, found {}", stmts.len())
969 }
970 );
971 let stmt = &stmts[0];
972
973 let is_tql_query = is_tql(query_ctx.sql_dialect(), &expr.sql)
974 .map_err(BoxedError::new)
975 .context(ExternalSnafu)?;
976 if is_tql_query {
977 if has_instant_ttl_source_table {
978 return InvalidSqlSnafu {
979 err_msg: INSTANT_TTL_FLOW_QUERY_ERROR.to_string(),
980 }
981 .fail();
982 }
983 return Ok(FlowType::Batching);
984 }
985
986 let plan = match stmt {
989 Statement::Tql(_) => {
990 if has_instant_ttl_source_table {
991 return InvalidSqlSnafu {
992 err_msg: INSTANT_TTL_FLOW_QUERY_ERROR.to_string(),
993 }
994 .fail();
995 }
996 return Ok(FlowType::Batching);
997 }
998 _ => engine
999 .planner()
1000 .plan(&QueryStatement::Sql(stmt.clone()), query_ctx)
1001 .await
1002 .map_err(BoxedError::new)
1003 .context(ExternalSnafu)?,
1004 };
1005
1006 struct FindQueryShape {
1009 is_aggr: bool,
1010 is_distinct: bool,
1011 has_distinct_on: bool,
1012 }
1013
1014 impl TreeNodeVisitor<'_> for FindQueryShape {
1015 type Node = LogicalPlan;
1016 fn f_down(
1017 &mut self,
1018 node: &Self::Node,
1019 ) -> datafusion_common::Result<datafusion_common::tree_node::TreeNodeRecursion>
1020 {
1021 match node {
1022 LogicalPlan::Aggregate(_) => self.is_aggr = true,
1023 LogicalPlan::Distinct(Distinct::All(_)) => self.is_distinct = true,
1024 LogicalPlan::Distinct(Distinct::On(_)) => self.has_distinct_on = true,
1025 _ => (),
1026 }
1027 Ok(datafusion_common::tree_node::TreeNodeRecursion::Continue)
1028 }
1029 }
1030
1031 let mut query_shape = FindQueryShape {
1032 is_aggr: false,
1033 is_distinct: false,
1034 has_distinct_on: false,
1035 };
1036
1037 plan.visit_with_subqueries(&mut query_shape)
1038 .context(BuildDfLogicalPlanSnafu)?;
1039 if query_shape.has_distinct_on {
1040 return InvalidSqlSnafu {
1041 err_msg: if has_instant_ttl_source_table {
1042 INSTANT_TTL_FLOW_QUERY_ERROR.to_string()
1043 } else {
1044 STATELESS_FLOW_QUERY_ERROR.to_string()
1045 },
1046 }
1047 .fail();
1048 }
1049 if query_shape.is_aggr {
1050 return if has_instant_ttl_source_table {
1051 InvalidSqlSnafu {
1052 err_msg: INSTANT_TTL_FLOW_QUERY_ERROR.to_string(),
1053 }
1054 .fail()
1055 } else {
1056 Ok(FlowType::Batching)
1057 };
1058 }
1059 if query_shape.is_distinct && !has_instant_ttl_source_table {
1060 return Ok(FlowType::Batching);
1061 }
1062
1063 ensure!(
1064 is_stateless_flow_plan(&plan),
1065 InvalidSqlSnafu {
1066 err_msg: if has_instant_ttl_source_table {
1067 INSTANT_TTL_FLOW_QUERY_ERROR.to_string()
1068 } else {
1069 STATELESS_FLOW_QUERY_ERROR.to_string()
1070 },
1071 }
1072 );
1073 Ok(FlowType::Streaming)
1074 }
1075
1076 #[tracing::instrument(skip_all)]
1077 pub async fn create_view(
1078 &self,
1079 create_view: CreateView,
1080 ctx: QueryContextRef,
1081 ) -> Result<TableRef> {
1082 let logical_plan = match &*create_view.query {
1084 Statement::Query(query) => {
1085 self.plan(
1086 &QueryStatement::Sql(Statement::Query(query.clone())),
1087 ctx.clone(),
1088 )
1089 .await?
1090 }
1091 Statement::Tql(query) => self.plan_tql(query.clone(), &ctx).await?,
1092 _ => {
1093 return InvalidViewStmtSnafu {}.fail();
1094 }
1095 };
1096 let definition = create_view.to_string();
1098
1099 let schema: Schema = logical_plan
1102 .schema()
1103 .clone()
1104 .try_into()
1105 .context(ConvertSchemaSnafu)?;
1106 let plan_columns: Vec<_> = schema
1107 .column_schemas()
1108 .iter()
1109 .map(|c| c.name.clone())
1110 .collect();
1111
1112 let columns: Vec<_> = create_view
1113 .columns
1114 .iter()
1115 .map(|ident| ident.to_string())
1116 .collect();
1117
1118 if !columns.is_empty() {
1120 ensure!(
1121 columns.len() == plan_columns.len(),
1122 error::ViewColumnsMismatchSnafu {
1123 view_name: create_view.name.to_string(),
1124 expected: plan_columns.len(),
1125 actual: columns.len(),
1126 }
1127 );
1128 }
1129
1130 let (table_names, plan) = extract_and_rewrite_full_table_names(logical_plan, ctx.clone())
1133 .context(ExtractTableNamesSnafu)?;
1134
1135 let table_names = table_names.into_iter().map(|t| t.into()).collect();
1136
1137 let encoded_plan = DFLogicalSubstraitConvertor
1144 .encode(&plan, DefaultSerializer)
1145 .context(SubstraitCodecSnafu)?;
1146
1147 let expr = expr_helper::to_create_view_expr(
1148 create_view,
1149 encoded_plan.to_vec(),
1150 table_names,
1151 columns,
1152 plan_columns,
1153 definition,
1154 ctx.clone(),
1155 )?;
1156
1157 self.create_view_by_expr(expr, ctx).await
1159 }
1160
1161 pub async fn create_view_by_expr(
1162 &self,
1163 expr: CreateViewExpr,
1164 ctx: QueryContextRef,
1165 ) -> Result<TableRef> {
1166 ensure_table_definition_writable(&expr.schema_name, &expr.view_name)?;
1169 ensure! {
1170 !(expr.create_if_not_exists & expr.or_replace),
1171 InvalidSqlSnafu {
1172 err_msg: "syntax error Create Or Replace and If Not Exist cannot be used together",
1173 }
1174 };
1175 let _timer = crate::metrics::DIST_CREATE_VIEW.start_timer();
1176
1177 let schema_exists = self
1178 .table_metadata_manager
1179 .schema_manager()
1180 .exists(SchemaNameKey::new(&expr.catalog_name, &expr.schema_name))
1181 .await
1182 .context(TableMetadataManagerSnafu)?;
1183
1184 ensure!(
1185 schema_exists,
1186 SchemaNotFoundSnafu {
1187 schema_info: &expr.schema_name,
1188 }
1189 );
1190
1191 if let Some(table) = self
1193 .catalog_manager
1194 .table(
1195 &expr.catalog_name,
1196 &expr.schema_name,
1197 &expr.view_name,
1198 Some(&ctx),
1199 )
1200 .await
1201 .context(CatalogSnafu)?
1202 {
1203 let table_type = table.table_info().table_type;
1204
1205 match (table_type, expr.create_if_not_exists, expr.or_replace) {
1206 (TableType::View, true, false) => {
1207 return Ok(table);
1208 }
1209 (TableType::View, false, false) => {
1210 return ViewAlreadyExistsSnafu {
1211 name: format_full_table_name(
1212 &expr.catalog_name,
1213 &expr.schema_name,
1214 &expr.view_name,
1215 ),
1216 }
1217 .fail();
1218 }
1219 (TableType::View, _, true) => {
1220 }
1222 _ => {
1223 return TableAlreadyExistsSnafu {
1224 table: format_full_table_name(
1225 &expr.catalog_name,
1226 &expr.schema_name,
1227 &expr.view_name,
1228 ),
1229 }
1230 .fail();
1231 }
1232 }
1233 }
1234
1235 ensure!(
1236 NAME_PATTERN_REG.is_match(&expr.view_name),
1237 InvalidViewNameSnafu {
1238 name: expr.view_name.clone(),
1239 }
1240 );
1241
1242 let view_name = TableName::new(&expr.catalog_name, &expr.schema_name, &expr.view_name);
1243
1244 let mut view_info = TableInfo {
1245 ident: metadata::TableIdent {
1246 table_id: 0,
1248 version: 0,
1249 },
1250 name: expr.view_name.clone(),
1251 desc: None,
1252 catalog_name: expr.catalog_name.clone(),
1253 schema_name: expr.schema_name.clone(),
1254 meta: TableMeta::empty(),
1256 table_type: TableType::View,
1257 };
1258
1259 let executor_context = to_executor_context(ctx, TriggerReason::Manual);
1260 let request = SubmitDdlTaskRequest::new(DdlTask::new_create_view(expr, view_info.clone()));
1261
1262 let resp = self
1263 .procedure_executor
1264 .submit_ddl_task(executor_context, request)
1265 .await
1266 .context(error::ExecuteDdlSnafu)?;
1267
1268 debug!(
1269 "Submit creating view '{view_name}' task response: {:?}",
1270 resp
1271 );
1272
1273 let view_id = resp
1274 .table_ids
1275 .into_iter()
1276 .next()
1277 .context(error::UnexpectedSnafu {
1278 violated: "expected table_id",
1279 })?;
1280 info!("Successfully created view '{view_name}' with view id {view_id}");
1281
1282 view_info.ident.table_id = view_id;
1283
1284 let view_info = Arc::new(view_info);
1285
1286 let table = DistTable::table(view_info);
1287
1288 self.cache_invalidator
1290 .invalidate(
1291 &Context::default(),
1292 &[
1293 CacheIdent::TableId(view_id),
1294 CacheIdent::TableName(view_name.clone()),
1295 ],
1296 )
1297 .await
1298 .context(error::InvalidateTableCacheSnafu)?;
1299
1300 Ok(table)
1301 }
1302
1303 #[tracing::instrument(skip_all)]
1304 pub async fn drop_flow(
1305 &self,
1306 catalog_name: String,
1307 flow_name: String,
1308 drop_if_exists: bool,
1309 query_context: QueryContextRef,
1310 ) -> Result<Output> {
1311 if let Some(flow) = self
1312 .flow_metadata_manager
1313 .flow_name_manager()
1314 .get(&catalog_name, &flow_name)
1315 .await
1316 .context(error::TableMetadataManagerSnafu)?
1317 {
1318 let flow_id = flow.flow_id();
1319 let task = DropFlowTask {
1320 catalog_name,
1321 flow_name,
1322 flow_id,
1323 drop_if_exists,
1324 };
1325 self.drop_flow_procedure(task, query_context).await?;
1326
1327 Ok(Output::new_with_affected_rows(0))
1328 } else if drop_if_exists {
1329 Ok(Output::new_with_affected_rows(0))
1330 } else {
1331 FlowNotFoundSnafu {
1332 flow_name: format_full_flow_name(&catalog_name, &flow_name),
1333 }
1334 .fail()
1335 }
1336 }
1337
1338 async fn drop_flow_procedure(
1339 &self,
1340 expr: DropFlowTask,
1341 query_context: QueryContextRef,
1342 ) -> Result<SubmitDdlTaskResponse> {
1343 let executor_context = to_executor_context(query_context, TriggerReason::Manual);
1344 let request = SubmitDdlTaskRequest::new(DdlTask::new_drop_flow(expr));
1345
1346 self.procedure_executor
1347 .submit_ddl_task(executor_context, request)
1348 .await
1349 .context(error::ExecuteDdlSnafu)
1350 }
1351
1352 #[cfg(feature = "enterprise")]
1353 #[tracing::instrument(skip_all)]
1354 pub(super) async fn drop_trigger(
1355 &self,
1356 catalog_name: String,
1357 trigger_name: String,
1358 drop_if_exists: bool,
1359 query_context: QueryContextRef,
1360 ) -> Result<Output> {
1361 let task = DropTriggerTask {
1362 catalog_name,
1363 trigger_name,
1364 drop_if_exists,
1365 };
1366 self.drop_trigger_procedure(task, query_context).await?;
1367 Ok(Output::new_with_affected_rows(0))
1368 }
1369
1370 #[cfg(feature = "enterprise")]
1371 async fn drop_trigger_procedure(
1372 &self,
1373 expr: DropTriggerTask,
1374 query_context: QueryContextRef,
1375 ) -> Result<SubmitDdlTaskResponse> {
1376 let executor_context = to_executor_context(query_context, TriggerReason::Manual);
1377 let request = SubmitDdlTaskRequest::new(DdlTask::new_drop_trigger(expr));
1378
1379 self.procedure_executor
1380 .submit_ddl_task(executor_context, request)
1381 .await
1382 .context(error::ExecuteDdlSnafu)
1383 }
1384
1385 #[tracing::instrument(skip_all)]
1387 pub async fn drop_view(
1388 &self,
1389 catalog: String,
1390 schema: String,
1391 view: String,
1392 drop_if_exists: bool,
1393 query_context: QueryContextRef,
1394 ) -> Result<Output> {
1395 let view_info = if let Some(view) = self
1396 .catalog_manager
1397 .table(&catalog, &schema, &view, None)
1398 .await
1399 .context(CatalogSnafu)?
1400 {
1401 view.table_info()
1402 } else if drop_if_exists {
1403 return Ok(Output::new_with_affected_rows(0));
1405 } else {
1406 return TableNotFoundSnafu {
1407 table_name: format_full_table_name(&catalog, &schema, &view),
1408 }
1409 .fail();
1410 };
1411
1412 ensure!(
1414 view_info.table_type == TableType::View,
1415 error::InvalidViewSnafu {
1416 msg: "not a view",
1417 view_name: format_full_table_name(&catalog, &schema, &view),
1418 }
1419 );
1420
1421 let view_id = view_info.table_id();
1422 let view_name = TableName::new(&catalog, &schema, &view);
1423
1424 let task = DropViewTask {
1425 catalog,
1426 schema,
1427 view,
1428 view_id,
1429 drop_if_exists,
1430 };
1431
1432 self.drop_view_procedure(task, query_context).await?;
1433
1434 self.cache_invalidator
1436 .invalidate(
1437 &Context::default(),
1438 &[
1439 CacheIdent::TableId(view_id),
1440 CacheIdent::TableName(view_name),
1441 ],
1442 )
1443 .await
1444 .context(error::InvalidateTableCacheSnafu)?;
1445
1446 Ok(Output::new_with_affected_rows(0))
1447 }
1448
1449 async fn drop_view_procedure(
1451 &self,
1452 expr: DropViewTask,
1453 query_context: QueryContextRef,
1454 ) -> Result<SubmitDdlTaskResponse> {
1455 let executor_context = to_executor_context(query_context, TriggerReason::Manual);
1456 let request = SubmitDdlTaskRequest::new(DdlTask::new_drop_view(expr));
1457
1458 self.procedure_executor
1459 .submit_ddl_task(executor_context, request)
1460 .await
1461 .context(error::ExecuteDdlSnafu)
1462 }
1463
1464 #[tracing::instrument(skip_all)]
1465 pub async fn alter_logical_tables(
1466 &self,
1467 alter_table_exprs: Vec<AlterTableExpr>,
1468 query_context: QueryContextRef,
1469 trigger_reason: TriggerReason,
1470 ) -> Result<Output> {
1471 let _timer = crate::metrics::DIST_ALTER_TABLES.start_timer();
1472 ensure!(
1473 !alter_table_exprs.is_empty(),
1474 EmptyDdlExprSnafu {
1475 name: "alter logical tables"
1476 }
1477 );
1478
1479 let mut groups: HashMap<TableId, Vec<AlterTableExpr>> = HashMap::new();
1481 for expr in alter_table_exprs {
1482 let catalog = if expr.catalog_name.is_empty() {
1484 query_context.current_catalog()
1485 } else {
1486 &expr.catalog_name
1487 };
1488 let schema = if expr.schema_name.is_empty() {
1489 query_context.current_schema()
1490 } else {
1491 expr.schema_name.clone()
1492 };
1493 let table_name = &expr.table_name;
1494 let table = self
1495 .catalog_manager
1496 .table(catalog, &schema, table_name, Some(&query_context))
1497 .await
1498 .context(CatalogSnafu)?
1499 .with_context(|| TableNotFoundSnafu {
1500 table_name: format_full_table_name(catalog, &schema, table_name),
1501 })?;
1502 let table_id = table.table_info().ident.table_id;
1503 let physical_table_id = self
1504 .table_metadata_manager
1505 .table_route_manager()
1506 .get_physical_table_id(table_id)
1507 .await
1508 .context(TableMetadataManagerSnafu)?;
1509 groups.entry(physical_table_id).or_default().push(expr);
1510 }
1511
1512 let mut handles = Vec::with_capacity(groups.len());
1514 for (_physical_table_id, exprs) in groups {
1515 let fut =
1516 self.alter_logical_tables_procedure(exprs, query_context.clone(), trigger_reason);
1517 handles.push(fut);
1518 }
1519 let _results = futures::future::try_join_all(handles).await?;
1520
1521 Ok(Output::new_with_affected_rows(0))
1522 }
1523
1524 #[tracing::instrument(skip_all)]
1525 pub async fn drop_table(
1526 &self,
1527 table_name: TableName,
1528 drop_if_exists: bool,
1529 query_context: QueryContextRef,
1530 ) -> Result<Output> {
1531 self.drop_tables(&[table_name], drop_if_exists, query_context)
1533 .await
1534 }
1535
1536 #[tracing::instrument(skip_all)]
1537 pub async fn drop_tables(
1538 &self,
1539 table_names: &[TableName],
1540 drop_if_exists: bool,
1541 query_context: QueryContextRef,
1542 ) -> Result<Output> {
1543 let mut tables = Vec::with_capacity(table_names.len());
1544 for table_name in table_names {
1545 ensure_table_writable(&table_name.schema_name, &table_name.table_name)?;
1546
1547 if let Some(table) = self
1548 .catalog_manager
1549 .table(
1550 &table_name.catalog_name,
1551 &table_name.schema_name,
1552 &table_name.table_name,
1553 Some(&query_context),
1554 )
1555 .await
1556 .context(CatalogSnafu)?
1557 {
1558 tables.push(table.table_info().table_id());
1559 } else if drop_if_exists {
1560 continue;
1562 } else {
1563 return TableNotFoundSnafu {
1564 table_name: table_name.to_string(),
1565 }
1566 .fail();
1567 }
1568 }
1569
1570 for (table_name, table_id) in table_names.iter().zip(tables.into_iter()) {
1571 self.drop_table_procedure(table_name, table_id, drop_if_exists, query_context.clone())
1572 .await?;
1573
1574 self.cache_invalidator
1576 .invalidate(
1577 &Context::default(),
1578 &[
1579 CacheIdent::TableId(table_id),
1580 CacheIdent::TableName(table_name.clone()),
1581 ],
1582 )
1583 .await
1584 .context(error::InvalidateTableCacheSnafu)?;
1585 }
1586 Ok(Output::new_with_affected_rows(0))
1587 }
1588
1589 #[cfg(feature = "enterprise")]
1590 #[tracing::instrument(skip_all)]
1591 pub async fn undrop_table(
1592 &self,
1593 table_name: TableName,
1594 query_context: QueryContextRef,
1595 ) -> Result<Output> {
1596 execute_undrop_table(
1597 &self.table_metadata_manager,
1598 &self.procedure_executor,
1599 &self.cache_invalidator,
1600 table_name,
1601 query_context,
1602 )
1603 .await
1604 }
1605
1606 #[tracing::instrument(skip_all)]
1607 pub async fn drop_database(
1608 &self,
1609 catalog: String,
1610 schema: String,
1611 drop_if_exists: bool,
1612 query_context: QueryContextRef,
1613 ) -> Result<Output> {
1614 ensure!(
1615 !is_readonly_schema(&schema),
1616 SchemaReadOnlySnafu { name: schema }
1617 );
1618
1619 if self
1620 .catalog_manager
1621 .schema_exists(&catalog, &schema, None)
1622 .await
1623 .context(CatalogSnafu)?
1624 {
1625 if schema == query_context.current_schema() {
1626 SchemaInUseSnafu { name: schema }.fail()
1627 } else {
1628 self.drop_database_procedure(catalog, schema, drop_if_exists, query_context)
1629 .await?;
1630
1631 Ok(Output::new_with_affected_rows(0))
1632 }
1633 } else if drop_if_exists {
1634 Ok(Output::new_with_affected_rows(0))
1636 } else {
1637 SchemaNotFoundSnafu {
1638 schema_info: schema,
1639 }
1640 .fail()
1641 }
1642 }
1643
1644 #[tracing::instrument(skip_all)]
1645 pub async fn truncate_table(
1646 &self,
1647 table_name: TableName,
1648 time_ranges: Vec<(Timestamp, Timestamp)>,
1649 query_context: QueryContextRef,
1650 ) -> Result<Output> {
1651 ensure_table_writable(&table_name.schema_name, &table_name.table_name)?;
1652
1653 let table = self
1654 .catalog_manager
1655 .table(
1656 &table_name.catalog_name,
1657 &table_name.schema_name,
1658 &table_name.table_name,
1659 Some(&query_context),
1660 )
1661 .await
1662 .context(CatalogSnafu)?
1663 .with_context(|| TableNotFoundSnafu {
1664 table_name: table_name.to_string(),
1665 })?;
1666 let table_id = table.table_info().table_id();
1667 self.truncate_table_procedure(&table_name, table_id, time_ranges, query_context)
1668 .await?;
1669
1670 Ok(Output::new_with_affected_rows(0))
1671 }
1672
1673 #[tracing::instrument(skip_all)]
1674 pub async fn alter_table(
1675 &self,
1676 alter_table: AlterTable,
1677 query_context: QueryContextRef,
1678 ) -> Result<Output> {
1679 if matches!(
1680 alter_table.alter_operation(),
1681 AlterTableOperation::Repartition { .. } | AlterTableOperation::Partition { .. }
1682 ) {
1683 let request = expr_helper::to_repartition_request(alter_table, &query_context)?;
1684 return self.repartition_table(request, &query_context).await;
1685 }
1686
1687 let expr = expr_helper::to_alter_table_expr(alter_table, &query_context)?;
1688 self.alter_table_inner(expr, query_context, TriggerReason::Manual)
1689 .await
1690 }
1691
1692 #[tracing::instrument(skip_all)]
1693 pub async fn repartition_table(
1694 &self,
1695 request: RepartitionRequest,
1696 query_context: &QueryContextRef,
1697 ) -> Result<Output> {
1698 ensure_table_definition_writable(&request.schema_name, &request.table_name)?;
1700
1701 let table_ref = TableReference::full(
1702 &request.catalog_name,
1703 &request.schema_name,
1704 &request.table_name,
1705 );
1706 let table = self
1708 .catalog_manager
1709 .table(
1710 &request.catalog_name,
1711 &request.schema_name,
1712 &request.table_name,
1713 Some(query_context),
1714 )
1715 .await
1716 .context(CatalogSnafu)?
1717 .with_context(|| TableNotFoundSnafu {
1718 table_name: table_ref.to_string(),
1719 })?;
1720 let table_id = table.table_info().ident.table_id;
1721 let (physical_table_id, physical_table_route) = self
1723 .table_metadata_manager
1724 .table_route_manager()
1725 .get_physical_table_route(table_id)
1726 .await
1727 .context(TableMetadataManagerSnafu)?;
1728
1729 ensure!(
1730 physical_table_id == table_id,
1731 NotSupportedSnafu {
1732 feat: "REPARTITION on logical tables"
1733 }
1734 );
1735
1736 let table_info = table.table_info();
1737 let existing_partition_columns = table_info.meta.partition_columns().collect::<Vec<_>>();
1738 let column_schemas = table_info.meta.schema.column_schemas();
1739 let target_partition_columns = match &request.source {
1745 RepartitionSource::Partitions {
1746 target_partition_columns,
1747 ..
1748 } => {
1749 ensure!(
1750 !existing_partition_columns.is_empty(),
1751 InvalidPartitionRuleSnafu {
1752 reason: format!(
1753 "table {} does not have partition columns, cannot repartition",
1754 table_ref
1755 )
1756 }
1757 );
1758
1759 if let Some(target_partition_columns) = target_partition_columns {
1760 ensure!(
1761 !target_partition_columns.is_empty(),
1762 InvalidPartitionRuleSnafu {
1763 reason: "ON COLUMNS requires at least one partition column"
1764 }
1765 );
1766 validate_and_collect_partition_columns(
1767 target_partition_columns,
1768 column_schemas,
1769 )?
1770 } else {
1771 existing_partition_columns.clone()
1772 }
1773 }
1774 RepartitionSource::Unpartitioned { partition_columns } => {
1775 ensure!(
1776 !partition_columns.is_empty(),
1777 InvalidPartitionRuleSnafu {
1778 reason: "PARTITION ON COLUMNS requires at least one partition column"
1779 }
1780 );
1781 ensure!(
1782 existing_partition_columns.is_empty(),
1783 InvalidPartitionRuleSnafu {
1784 reason: format!("table {} already has partition columns", table_ref)
1785 }
1786 );
1787 partition_columns
1788 .iter()
1789 .map(|column_name| {
1790 column_schemas
1791 .iter()
1792 .find(|column| &column.name == column_name)
1793 .with_context(|| ColumnNotFoundSnafu { msg: column_name })
1794 })
1795 .collect::<Result<Vec<_>>>()?
1796 }
1797 };
1798
1799 let from_column_name_and_type = column_name_and_type(&existing_partition_columns);
1800 let target_column_name_and_type = column_name_and_type(&target_partition_columns);
1801 let target_partition_column_names = target_partition_columns
1802 .iter()
1803 .map(|column| column.name.clone())
1804 .collect::<Vec<_>>();
1805 let timezone = query_context.timezone();
1806 let from_partition_exprs = match &request.source {
1808 RepartitionSource::Partitions { from_exprs, .. } => from_exprs
1809 .iter()
1810 .map(|expr| convert_one_expr(expr, &from_column_name_and_type, &timezone))
1811 .collect::<Result<Vec<_>>>()?,
1812 RepartitionSource::Unpartitioned { .. } => vec![],
1813 };
1814
1815 let mut into_partition_exprs = request
1816 .into_exprs
1817 .iter()
1818 .map(|expr| convert_one_expr(expr, &target_column_name_and_type, &timezone))
1819 .collect::<Result<Vec<_>>>()?;
1820
1821 if matches!(&request.source, RepartitionSource::Partitions { .. })
1824 && from_partition_exprs.len() > 1
1825 && into_partition_exprs.len() == 1
1826 && let Some(expr) = into_partition_exprs.pop()
1827 {
1828 into_partition_exprs.push(partition::simplify::simplify_merged_partition_expr(expr));
1829 }
1830
1831 let mut existing_partition_exprs =
1833 Vec::with_capacity(physical_table_route.region_routes.len());
1834 for route in &physical_table_route.region_routes {
1835 let expr_json = route.region.partition_expr();
1836 if !expr_json.is_empty() {
1837 match PartitionExpr::from_json_str(&expr_json) {
1838 Ok(Some(expr)) => existing_partition_exprs.push(expr),
1839 Ok(None) => {
1840 }
1842 Err(e) => {
1843 return Err(e).context(DeserializePartitionExprSnafu);
1844 }
1845 }
1846 }
1847 }
1848
1849 if matches!(&request.source, RepartitionSource::Partitions { .. }) {
1852 for from_expr in &from_partition_exprs {
1853 ensure!(
1854 existing_partition_exprs.contains(from_expr),
1855 InvalidPartitionRuleSnafu {
1856 reason: format!(
1857 "partition expression '{}' does not exist in table {}",
1858 from_expr, table_ref
1859 )
1860 }
1861 );
1862 }
1863 }
1864
1865 let new_partition_exprs: Vec<PartitionExpr> = match &request.source {
1868 RepartitionSource::Partitions { .. } => existing_partition_exprs
1869 .into_iter()
1870 .filter(|expr| !from_partition_exprs.contains(expr))
1871 .chain(into_partition_exprs.clone().into_iter())
1872 .collect(),
1873 RepartitionSource::Unpartitioned { .. } => into_partition_exprs.clone(),
1874 };
1875 ensure_partition_expr_columns_in_target(
1876 &new_partition_exprs,
1877 &target_partition_column_names.iter().collect(),
1878 )?;
1879 let new_partition_exprs_len = new_partition_exprs.len();
1880 let from_partition_exprs_len = from_partition_exprs.len();
1881
1882 let _ = MultiDimPartitionRule::try_new(
1884 target_partition_column_names,
1885 vec![],
1886 new_partition_exprs,
1887 true,
1888 )
1889 .context(InvalidPartitionSnafu)?;
1890
1891 let ddl_options = parse_ddl_options(&request.options)?;
1892 let serialize_exprs = |exprs: Vec<PartitionExpr>| -> Result<Vec<String>> {
1893 let mut json_exprs = Vec::with_capacity(exprs.len());
1894 for expr in exprs {
1895 json_exprs.push(expr.as_json_str().context(SerializePartitionExprSnafu)?);
1896 }
1897 Ok(json_exprs)
1898 };
1899 let from_partition_exprs_json = serialize_exprs(from_partition_exprs)?;
1900 let into_partition_exprs_json = serialize_exprs(into_partition_exprs)?;
1901 let source = match &request.source {
1902 RepartitionSource::Partitions {
1903 target_partition_columns,
1904 ..
1905 } => Source::PartitionExprs(PartitionedSource {
1906 exprs: from_partition_exprs_json,
1907 target_partition_columns: target_partition_columns
1908 .clone()
1909 .map(|columns| TargetPartitionColumns { columns }),
1910 }),
1911 RepartitionSource::Unpartitioned { partition_columns } => {
1912 Source::Unpartitioned(UnpartitionedSource {
1913 partition_columns: partition_columns.clone(),
1914 })
1915 }
1916 };
1917 let repartition = Repartition {
1918 into_partition_exprs: into_partition_exprs_json,
1919 source: Some(source),
1920 ..Default::default()
1921 };
1922 let executor_context = to_executor_context(query_context.clone(), TriggerReason::Manual);
1923 let mut req = SubmitDdlTaskRequest::new(DdlTask::new_alter_table(AlterTableExpr {
1924 catalog_name: request.catalog_name.clone(),
1925 schema_name: request.schema_name.clone(),
1926 table_name: request.table_name.clone(),
1927 kind: Some(Kind::Repartition(repartition)),
1928 }));
1929 req.wait = ddl_options.wait;
1930 req.timeout = ddl_options.timeout;
1931
1932 info!(
1933 "Submitting repartition task for table {} (table_id={}), from {} to {} partitions, timeout: {:?}, wait: {}",
1934 table_ref,
1935 table_id,
1936 from_partition_exprs_len,
1937 new_partition_exprs_len,
1938 ddl_options.timeout,
1939 ddl_options.wait
1940 );
1941
1942 let response = self
1943 .procedure_executor
1944 .submit_ddl_task(executor_context, req)
1945 .await
1946 .context(error::ExecuteDdlSnafu)?;
1947
1948 if !ddl_options.wait {
1949 return build_procedure_id_output(response.key);
1950 }
1951
1952 let invalidate_keys = vec![
1954 CacheIdent::TableId(table_id),
1955 CacheIdent::TableName(TableName::new(
1956 request.catalog_name,
1957 request.schema_name,
1958 request.table_name,
1959 )),
1960 ];
1961
1962 self.cache_invalidator
1964 .invalidate(&Context::default(), &invalidate_keys)
1965 .await
1966 .context(error::InvalidateTableCacheSnafu)?;
1967
1968 Ok(Output::new_with_affected_rows(0))
1969 }
1970
1971 #[tracing::instrument(skip_all)]
1972 pub async fn alter_table_inner(
1973 &self,
1974 expr: AlterTableExpr,
1975 query_context: QueryContextRef,
1976 trigger_reason: TriggerReason,
1977 ) -> Result<Output> {
1978 ensure_table_definition_writable(&expr.schema_name, &expr.table_name)?;
1979
1980 let catalog_name = if expr.catalog_name.is_empty() {
1981 DEFAULT_CATALOG_NAME.to_string()
1982 } else {
1983 expr.catalog_name.clone()
1984 };
1985
1986 let schema_name = if expr.schema_name.is_empty() {
1987 DEFAULT_SCHEMA_NAME.to_string()
1988 } else {
1989 expr.schema_name.clone()
1990 };
1991
1992 let table_name = expr.table_name.clone();
1993
1994 let table = self
1995 .catalog_manager
1996 .table(
1997 &catalog_name,
1998 &schema_name,
1999 &table_name,
2000 Some(&query_context),
2001 )
2002 .await
2003 .context(CatalogSnafu)?
2004 .with_context(|| TableNotFoundSnafu {
2005 table_name: format_full_table_name(&catalog_name, &schema_name, &table_name),
2006 })?;
2007
2008 let table_id = table.table_info().ident.table_id;
2009 let need_alter = verify_alter(table_id, table.table_info(), expr.clone())?;
2010 if !need_alter {
2011 return Ok(Output::new_with_affected_rows(0));
2012 }
2013 info!(
2014 "Table info before alter is {:?}, expr: {:?}",
2015 table.table_info(),
2016 expr
2017 );
2018
2019 let physical_table_id = self
2020 .table_metadata_manager
2021 .table_route_manager()
2022 .get_physical_table_id(table_id)
2023 .await
2024 .context(TableMetadataManagerSnafu)?;
2025
2026 let executor_context = to_executor_context(query_context, trigger_reason);
2027
2028 let (req, invalidate_keys) = if physical_table_id == table_id {
2029 let req = SubmitDdlTaskRequest::new(DdlTask::new_alter_table(expr));
2031
2032 let invalidate_keys = vec![
2033 CacheIdent::TableId(table_id),
2034 CacheIdent::TableName(TableName::new(catalog_name, schema_name, table_name)),
2035 ];
2036
2037 (req, invalidate_keys)
2038 } else {
2039 let annotation_alter = match expr.kind.as_ref() {
2042 Some(kind) => common_grpc_expr::annotation_alter_family(kind)
2043 .context(AlterExprToRequestSnafu)?
2044 .is_some_and(|family| family.allows_logical_tables()),
2045 None => false,
2046 };
2047 let task = if annotation_alter {
2048 DdlTask::new_alter_table(expr)
2049 } else {
2050 DdlTask::new_alter_logical_tables(vec![expr])
2051 };
2052 let req = SubmitDdlTaskRequest::new(task);
2053
2054 let mut invalidate_keys = vec![
2055 CacheIdent::TableId(physical_table_id),
2056 CacheIdent::TableId(table_id),
2057 CacheIdent::TableName(TableName::new(catalog_name, schema_name, table_name)),
2058 ];
2059
2060 let physical_table = self
2061 .table_metadata_manager
2062 .table_info_manager()
2063 .get(physical_table_id)
2064 .await
2065 .context(TableMetadataManagerSnafu)?
2066 .map(|x| x.into_inner());
2067 if let Some(physical_table) = physical_table {
2068 let physical_table_name = TableName::new(
2069 physical_table.table_info.catalog_name,
2070 physical_table.table_info.schema_name,
2071 physical_table.table_info.name,
2072 );
2073 invalidate_keys.push(CacheIdent::TableName(physical_table_name));
2074 }
2075
2076 (req, invalidate_keys)
2077 };
2078
2079 self.procedure_executor
2080 .submit_ddl_task(executor_context, req)
2081 .await
2082 .context(error::ExecuteDdlSnafu)?;
2083
2084 self.cache_invalidator
2086 .invalidate(&Context::default(), &invalidate_keys)
2087 .await
2088 .context(error::InvalidateTableCacheSnafu)?;
2089
2090 Ok(Output::new_with_affected_rows(0))
2091 }
2092
2093 #[cfg(feature = "enterprise")]
2094 #[tracing::instrument(skip_all)]
2095 pub async fn alter_trigger(
2096 &self,
2097 _alter_expr: AlterTrigger,
2098 _query_context: QueryContextRef,
2099 ) -> Result<Output> {
2100 crate::error::NotSupportedSnafu {
2101 feat: "alter trigger",
2102 }
2103 .fail()
2104 }
2105
2106 #[tracing::instrument(skip_all)]
2107 pub async fn alter_database(
2108 &self,
2109 alter_expr: AlterDatabase,
2110 query_context: QueryContextRef,
2111 ) -> Result<Output> {
2112 let alter_expr = expr_helper::to_alter_database_expr(alter_expr, &query_context)?;
2113 self.alter_database_inner(alter_expr, query_context).await
2114 }
2115
2116 #[tracing::instrument(skip_all)]
2117 pub async fn alter_database_inner(
2118 &self,
2119 alter_expr: AlterDatabaseExpr,
2120 query_context: QueryContextRef,
2121 ) -> Result<Output> {
2122 ensure!(
2123 !is_readonly_schema(&alter_expr.schema_name),
2124 SchemaReadOnlySnafu {
2125 name: query_context.current_schema().clone()
2126 }
2127 );
2128
2129 let exists = self
2130 .catalog_manager
2131 .schema_exists(&alter_expr.catalog_name, &alter_expr.schema_name, None)
2132 .await
2133 .context(CatalogSnafu)?;
2134 ensure!(
2135 exists,
2136 SchemaNotFoundSnafu {
2137 schema_info: alter_expr.schema_name,
2138 }
2139 );
2140
2141 let cache_ident = [CacheIdent::SchemaName(SchemaName {
2142 catalog_name: alter_expr.catalog_name.clone(),
2143 schema_name: alter_expr.schema_name.clone(),
2144 })];
2145
2146 self.alter_database_procedure(alter_expr, query_context)
2147 .await?;
2148
2149 self.cache_invalidator
2151 .invalidate(&Context::default(), &cache_ident)
2152 .await
2153 .context(error::InvalidateTableCacheSnafu)?;
2154
2155 Ok(Output::new_with_affected_rows(0))
2156 }
2157
2158 async fn create_table_procedure(
2159 &self,
2160 create_table: CreateTableExpr,
2161 partitions: Vec<PartitionExpr>,
2162 table_info: TableInfo,
2163 query_context: QueryContextRef,
2164 trigger_reason: TriggerReason,
2165 ) -> Result<SubmitDdlTaskResponse> {
2166 let partitions = partitions
2167 .into_iter()
2168 .map(|expr| expr.as_pb_partition().context(PartitionExprToPbSnafu))
2169 .collect::<Result<Vec<_>>>()?;
2170
2171 let executor_context = to_executor_context_with_origin_frontend(
2172 query_context,
2173 &self.origin_frontend_addr,
2174 trigger_reason,
2175 );
2176 let request = SubmitDdlTaskRequest::new(DdlTask::new_create_table(
2177 create_table,
2178 partitions,
2179 table_info,
2180 ));
2181
2182 self.procedure_executor
2183 .submit_ddl_task(executor_context, request)
2184 .await
2185 .context(error::ExecuteDdlSnafu)
2186 }
2187
2188 async fn create_logical_tables_procedure(
2189 &self,
2190 tables_data: Vec<(CreateTableExpr, TableInfo)>,
2191 query_context: QueryContextRef,
2192 trigger_reason: TriggerReason,
2193 ) -> Result<SubmitDdlTaskResponse> {
2194 let executor_context = to_executor_context_with_origin_frontend(
2195 query_context,
2196 &self.origin_frontend_addr,
2197 trigger_reason,
2198 );
2199 let request = SubmitDdlTaskRequest::new(DdlTask::new_create_logical_tables(tables_data));
2200
2201 self.procedure_executor
2202 .submit_ddl_task(executor_context, request)
2203 .await
2204 .context(error::ExecuteDdlSnafu)
2205 }
2206
2207 async fn alter_logical_tables_procedure(
2208 &self,
2209 tables_data: Vec<AlterTableExpr>,
2210 query_context: QueryContextRef,
2211 trigger_reason: TriggerReason,
2212 ) -> Result<SubmitDdlTaskResponse> {
2213 let executor_context = to_executor_context(query_context, trigger_reason);
2214 let request = SubmitDdlTaskRequest::new(DdlTask::new_alter_logical_tables(tables_data));
2215
2216 self.procedure_executor
2217 .submit_ddl_task(executor_context, request)
2218 .await
2219 .context(error::ExecuteDdlSnafu)
2220 }
2221
2222 async fn drop_table_procedure(
2223 &self,
2224 table_name: &TableName,
2225 table_id: TableId,
2226 drop_if_exists: bool,
2227 query_context: QueryContextRef,
2228 ) -> Result<SubmitDdlTaskResponse> {
2229 let executor_context = to_executor_context(query_context, TriggerReason::Manual);
2230 let request = SubmitDdlTaskRequest::new(DdlTask::new_drop_table(
2231 table_name.catalog_name.clone(),
2232 table_name.schema_name.clone(),
2233 table_name.table_name.clone(),
2234 table_id,
2235 drop_if_exists,
2236 ));
2237
2238 self.procedure_executor
2239 .submit_ddl_task(executor_context, request)
2240 .await
2241 .context(error::ExecuteDdlSnafu)
2242 }
2243
2244 async fn drop_database_procedure(
2245 &self,
2246 catalog: String,
2247 schema: String,
2248 drop_if_exists: bool,
2249 query_context: QueryContextRef,
2250 ) -> Result<SubmitDdlTaskResponse> {
2251 let executor_context = to_executor_context(query_context, TriggerReason::Manual);
2252 let request =
2253 SubmitDdlTaskRequest::new(DdlTask::new_drop_database(catalog, schema, drop_if_exists));
2254
2255 self.procedure_executor
2256 .submit_ddl_task(executor_context, request)
2257 .await
2258 .context(error::ExecuteDdlSnafu)
2259 }
2260
2261 async fn alter_database_procedure(
2262 &self,
2263 alter_expr: AlterDatabaseExpr,
2264 query_context: QueryContextRef,
2265 ) -> Result<SubmitDdlTaskResponse> {
2266 let executor_context = to_executor_context(query_context, TriggerReason::Manual);
2267 let request = SubmitDdlTaskRequest::new(DdlTask::new_alter_database(alter_expr));
2268
2269 self.procedure_executor
2270 .submit_ddl_task(executor_context, request)
2271 .await
2272 .context(error::ExecuteDdlSnafu)
2273 }
2274
2275 async fn truncate_table_procedure(
2276 &self,
2277 table_name: &TableName,
2278 table_id: TableId,
2279 time_ranges: Vec<(Timestamp, Timestamp)>,
2280 query_context: QueryContextRef,
2281 ) -> Result<SubmitDdlTaskResponse> {
2282 let executor_context = to_executor_context(query_context, TriggerReason::Manual);
2283 let request = SubmitDdlTaskRequest::new(DdlTask::new_truncate_table(
2284 table_name.catalog_name.clone(),
2285 table_name.schema_name.clone(),
2286 table_name.table_name.clone(),
2287 table_id,
2288 time_ranges,
2289 ));
2290
2291 self.procedure_executor
2292 .submit_ddl_task(executor_context, request)
2293 .await
2294 .context(error::ExecuteDdlSnafu)
2295 }
2296
2297 #[tracing::instrument(skip_all)]
2298 pub async fn create_database(
2299 &self,
2300 database: &str,
2301 create_if_not_exists: bool,
2302 options: HashMap<String, String>,
2303 query_context: QueryContextRef,
2304 ) -> Result<Output> {
2305 let catalog = query_context.current_catalog();
2306 ensure!(
2307 NAME_PATTERN_REG.is_match(catalog),
2308 error::UnexpectedSnafu {
2309 violated: format!("Invalid catalog name: {}", catalog)
2310 }
2311 );
2312
2313 ensure!(
2314 NAME_PATTERN_REG.is_match(database),
2315 error::UnexpectedSnafu {
2316 violated: format!("Invalid database name: {}", database)
2317 }
2318 );
2319
2320 #[cfg(feature = "enterprise")]
2321 let creator = match &self.create_database_handler {
2322 Some(handler) => handler.creator(&query_context).context(ExternalSnafu)?,
2323 None => None,
2324 };
2325 #[cfg(not(feature = "enterprise"))]
2326 let creator = None;
2327
2328 let output = if !self
2329 .catalog_manager
2330 .schema_exists(catalog, database, None)
2331 .await
2332 .context(CatalogSnafu)?
2333 && !self.catalog_manager.is_reserved_schema_name(database)
2334 {
2335 self.create_database_procedure(
2336 catalog.to_string(),
2337 database.to_string(),
2338 create_if_not_exists,
2339 options,
2340 query_context.clone(),
2341 creator,
2342 )
2343 .await?;
2344
2345 Output::new_with_affected_rows(1)
2346 } else if create_if_not_exists {
2347 Output::new_with_affected_rows(1)
2348 } else {
2349 return error::SchemaExistsSnafu { name: database }.fail();
2350 };
2351
2352 #[cfg(feature = "enterprise")]
2353 if let Some(handler) = &self.create_database_handler {
2354 handler
2355 .refresh_current_user(&query_context)
2356 .await
2357 .context(ExternalSnafu)?;
2358 }
2359
2360 Ok(output)
2361 }
2362
2363 async fn create_database_procedure(
2364 &self,
2365 catalog: String,
2366 database: String,
2367 create_if_not_exists: bool,
2368 options: HashMap<String, String>,
2369 query_context: QueryContextRef,
2370 creator: Option<CreatorGrantIntent>,
2371 ) -> Result<SubmitDdlTaskResponse> {
2372 let executor_context = to_executor_context(query_context, TriggerReason::Manual);
2373 let request = SubmitDdlTaskRequest::new(DdlTask::new_create_database(
2374 catalog,
2375 database,
2376 create_if_not_exists,
2377 options,
2378 creator,
2379 ));
2380
2381 self.procedure_executor
2382 .submit_ddl_task(executor_context, request)
2383 .await
2384 .context(error::ExecuteDdlSnafu)
2385 }
2386}
2387
2388pub fn parse_partitions(
2390 create_table: &CreateTableExpr,
2391 partitions: Option<Partitions>,
2392 query_ctx: &QueryContextRef,
2393) -> Result<(Vec<PartitionExpr>, Vec<String>)> {
2394 let partition_columns = find_partition_columns(&partitions)?;
2397 let partition_exprs =
2398 find_partition_entries(create_table, &partitions, &partition_columns, query_ctx)?;
2399
2400 let exprs = partition_exprs.clone();
2402 MultiDimPartitionRule::try_new(partition_columns.clone(), vec![], exprs, true)
2403 .context(InvalidPartitionSnafu)?;
2404
2405 Ok((partition_exprs, partition_columns))
2406}
2407
2408fn parse_partitions_for_logical_validation(
2409 create_table: &CreateTableExpr,
2410 partitions: &Partitions,
2411 query_ctx: &QueryContextRef,
2412) -> Result<(Vec<String>, Vec<PartitionExpr>)> {
2413 let partition_columns = partitions
2414 .column_list
2415 .iter()
2416 .map(|ident| ident.value.clone())
2417 .collect::<Vec<_>>();
2418
2419 let column_name_and_type = partition_columns
2420 .iter()
2421 .map(|pc| {
2422 let column = create_table
2423 .column_defs
2424 .iter()
2425 .find(|c| &c.name == pc)
2426 .context(ColumnNotFoundSnafu { msg: pc.clone() })?;
2427 let column_name = &column.name;
2428 let data_type = ConcreteDataType::from(
2429 ColumnDataTypeWrapper::try_new(column.data_type, column.datatype_extension.clone())
2430 .context(ColumnDataTypeSnafu)?,
2431 );
2432 Ok((column_name, data_type))
2433 })
2434 .collect::<Result<HashMap<_, _>>>()?;
2435
2436 let mut partition_exprs = Vec::with_capacity(partitions.exprs.len());
2437 for expr in &partitions.exprs {
2438 let partition_expr = convert_one_expr(expr, &column_name_and_type, &query_ctx.timezone())?;
2439 partition_exprs.push(partition_expr);
2440 }
2441
2442 MultiDimPartitionRule::try_new(
2443 partition_columns.clone(),
2444 vec![],
2445 partition_exprs.clone(),
2446 true,
2447 )
2448 .context(InvalidPartitionSnafu)?;
2449
2450 Ok((partition_columns, partition_exprs))
2451}
2452
2453pub fn verify_alter(
2459 table_id: TableId,
2460 table_info: Arc<TableInfo>,
2461 expr: AlterTableExpr,
2462) -> Result<bool> {
2463 let request: AlterTableRequest =
2464 common_grpc_expr::alter_expr_to_request(table_id, expr, Some(&table_info.meta))
2465 .context(AlterExprToRequestSnafu)?;
2466
2467 let AlterTableRequest {
2468 table_name,
2469 alter_kind,
2470 ..
2471 } = &request;
2472
2473 if let AlterKind::RenameTable { new_table_name } = alter_kind {
2474 ensure!(
2475 NAME_PATTERN_REG.is_match(new_table_name),
2476 error::UnexpectedSnafu {
2477 violated: format!("Invalid table name: {}", new_table_name)
2478 }
2479 );
2480 ensure_table_definition_writable(&table_info.schema_name, new_table_name)?;
2483 } else if let AlterKind::AddColumns { columns } = alter_kind {
2484 let column_names: HashSet<_> = table_info
2487 .meta
2488 .schema
2489 .column_schemas()
2490 .iter()
2491 .map(|schema| &schema.name)
2492 .collect();
2493 if columns.iter().all(|column| {
2494 column_names.contains(&column.column_schema.name) && column.add_if_not_exists
2495 }) {
2496 return Ok(false);
2497 }
2498 }
2499
2500 let new_meta = table_info
2501 .meta
2502 .builder_with_alter_kind(table_name, &request.alter_kind)
2503 .context(error::TableSnafu)?
2504 .build()
2505 .context(error::BuildTableMetaSnafu { table_name })?;
2506
2507 validate_json2_columns_append_mode(&new_meta.schema, &new_meta.options)?;
2508
2509 Ok(true)
2510}
2511
2512pub fn create_table_info(
2513 create_table: &CreateTableExpr,
2514 partition_columns: Vec<String>,
2515) -> Result<TableInfo> {
2516 let mut column_schemas = Vec::with_capacity(create_table.column_defs.len());
2517 let mut column_name_to_index_map = HashMap::new();
2518
2519 for (idx, column) in create_table.column_defs.iter().enumerate() {
2520 let schema =
2521 column_def::try_as_column_schema(column).context(error::InvalidColumnDefSnafu {
2522 column: &column.name,
2523 })?;
2524 let schema = schema.with_time_index(column.name == create_table.time_index);
2525
2526 column_schemas.push(schema);
2527 let _ = column_name_to_index_map.insert(column.name.clone(), idx);
2528 }
2529
2530 let next_column_id = column_schemas.len() as u32;
2531 let schema = Arc::new(Schema::try_new(column_schemas).context(ConvertSchemaSnafu)?);
2532
2533 let primary_key_indices = create_table
2534 .primary_keys
2535 .iter()
2536 .map(|name| {
2537 column_name_to_index_map
2538 .get(name)
2539 .cloned()
2540 .context(ColumnNotFoundSnafu { msg: name })
2541 })
2542 .collect::<Result<Vec<_>>>()?;
2543
2544 let partition_key_indices = partition_columns
2545 .into_iter()
2546 .map(|col_name| {
2547 column_name_to_index_map
2548 .get(&col_name)
2549 .cloned()
2550 .context(ColumnNotFoundSnafu { msg: col_name })
2551 })
2552 .collect::<Result<Vec<_>>>()?;
2553
2554 let mut table_options = TableOptions::try_from_iter(&create_table.table_options)
2555 .context(UnrecognizedTableOptionSnafu)?;
2556
2557 validate_json2_columns_append_mode(&schema, &table_options)?;
2558
2559 validate_and_normalize_annotations(&mut table_options, &schema, &partition_key_indices)?;
2560
2561 let meta = TableMeta {
2562 schema,
2563 primary_key_indices,
2564 value_indices: vec![],
2565 engine: create_table.engine.clone(),
2566 next_column_id,
2567 options: table_options,
2568 created_on: Utc::now(),
2569 updated_on: Utc::now(),
2570 partition_key_indices,
2571 column_ids: vec![],
2572 };
2573
2574 let desc = if create_table.desc.is_empty() {
2575 create_table.table_options.get(COMMENT_KEY).cloned()
2576 } else {
2577 Some(create_table.desc.clone())
2578 };
2579
2580 let table_info = TableInfo {
2581 ident: metadata::TableIdent {
2582 table_id: 0,
2584 version: 0,
2585 },
2586 name: create_table.table_name.clone(),
2587 desc,
2588 catalog_name: create_table.catalog_name.clone(),
2589 schema_name: create_table.schema_name.clone(),
2590 meta,
2591 table_type: TableType::Base,
2592 };
2593 Ok(table_info)
2594}
2595
2596fn validate_json2_columns_append_mode(schema: &Schema, table_options: &TableOptions) -> Result<()> {
2597 let append_mode = table_options
2598 .extra_options
2599 .get(APPEND_MODE_KEY)
2600 .is_some_and(|value| value == "true");
2601
2602 for column in schema.column_schemas() {
2603 if column.data_type.is_json2() {
2604 ensure!(
2605 append_mode,
2606 InvalidSqlSnafu {
2607 err_msg: format!(
2608 "JSON2 column `{}` requires {}='true'",
2609 column.name, APPEND_MODE_KEY
2610 ),
2611 }
2612 );
2613 }
2614 }
2615
2616 Ok(())
2617}
2618
2619fn ensure_table_writable(schema: &str, table: &str) -> Result<()> {
2621 ensure!(
2622 !is_readonly_schema(schema),
2623 SchemaReadOnlySnafu {
2624 name: schema.to_string()
2625 }
2626 );
2627 ensure!(
2628 !is_readonly_table(schema, table),
2629 TableReadOnlySnafu {
2630 name: table.to_string()
2631 }
2632 );
2633 Ok(())
2634}
2635
2636fn ensure_table_definition_writable(schema: &str, table: &str) -> Result<()> {
2642 ensure_table_writable(schema, table)?;
2643 ensure!(
2644 !is_ddl_reserved_table(schema, table),
2645 TableDdlReservedSnafu {
2646 name: table.to_string()
2647 }
2648 );
2649 Ok(())
2650}
2651
2652fn validate_and_normalize_annotations(
2656 options: &mut TableOptions,
2657 schema: &Schema,
2658 partition_key_indices: &[usize],
2659) -> Result<()> {
2660 use table::requests::AnnotationValidationError as CheckError;
2661 let cx = AnnotationContext {
2662 schema,
2663 partition_key_indices,
2664 };
2665 validate_and_normalize_annotation_options(options, &cx).map_err(|e| match e {
2666 CheckError::ColumnNotFound { column } => ColumnNotFoundSnafu { msg: column }.build(),
2667 e @ (CheckError::UnknownKey { .. }
2668 | CheckError::InvalidValue { .. }
2669 | CheckError::InvalidPartitionNumHint { .. }
2670 | CheckError::ColumnNotStringForm { .. }) => InvalidSqlSnafu {
2671 err_msg: e.to_string(),
2672 }
2673 .build(),
2674 e @ (CheckError::NotSingleColumn
2675 | CheckError::PartitionMetadataConflict
2676 | CheckError::TimeIndexConflict) => InvalidPartitionRuleSnafu {
2677 reason: e.to_string(),
2678 }
2679 .build(),
2680 })
2681}
2682
2683fn find_partition_columns(partitions: &Option<Partitions>) -> Result<Vec<String>> {
2684 let columns = if let Some(partitions) = partitions {
2685 partitions
2686 .column_list
2687 .iter()
2688 .map(|x| x.value.clone())
2689 .collect::<Vec<_>>()
2690 } else {
2691 vec![]
2692 };
2693 Ok(columns)
2694}
2695
2696fn find_partition_entries(
2700 create_table: &CreateTableExpr,
2701 partitions: &Option<Partitions>,
2702 partition_columns: &[String],
2703 query_ctx: &QueryContextRef,
2704) -> Result<Vec<PartitionExpr>> {
2705 let Some(partitions) = partitions else {
2706 return Ok(vec![]);
2707 };
2708
2709 let column_name_and_type = partition_columns
2711 .iter()
2712 .map(|pc| {
2713 let column = create_table
2714 .column_defs
2715 .iter()
2716 .find(|c| &c.name == pc)
2717 .unwrap();
2719 let column_name = &column.name;
2720 let data_type = ConcreteDataType::from(
2721 ColumnDataTypeWrapper::try_new(column.data_type, column.datatype_extension.clone())
2722 .context(ColumnDataTypeSnafu)?,
2723 );
2724 Ok((column_name, data_type))
2725 })
2726 .collect::<Result<HashMap<_, _>>>()?;
2727
2728 let mut partition_exprs = Vec::with_capacity(partitions.exprs.len());
2730 for partition in &partitions.exprs {
2731 let partition_expr =
2732 convert_one_expr(partition, &column_name_and_type, &query_ctx.timezone())?;
2733 partition_exprs.push(partition_expr);
2734 }
2735
2736 Ok(partition_exprs)
2737}
2738
2739fn column_name_and_type<'a>(
2740 partition_columns: &'a [&'a ColumnSchema],
2741) -> HashMap<&'a String, ConcreteDataType> {
2742 partition_columns
2743 .iter()
2744 .map(|column| (&column.name, column.data_type.clone()))
2745 .collect()
2746}
2747
2748fn validate_and_collect_partition_columns<'a>(
2749 column_names: &[String],
2750 column_schemas: &'a [ColumnSchema],
2751) -> Result<Vec<&'a ColumnSchema>> {
2752 let mut seen = HashSet::with_capacity(column_names.len());
2753 column_names
2754 .iter()
2755 .map(|column_name| {
2756 ensure!(
2757 seen.insert(column_name),
2758 InvalidPartitionRuleSnafu {
2759 reason: format!("duplicate partition column '{}'", column_name)
2760 }
2761 );
2762 column_schemas
2763 .iter()
2764 .find(|column| &column.name == column_name)
2765 .with_context(|| ColumnNotFoundSnafu { msg: column_name })
2766 })
2767 .collect()
2768}
2769
2770fn ensure_partition_expr_columns_in_target(
2771 partition_exprs: &[PartitionExpr],
2772 target_partition_columns: &HashSet<&String>,
2773) -> Result<()> {
2774 for expr in partition_exprs {
2775 ensure_partition_operand_columns_in_target(&expr.lhs, target_partition_columns)?;
2776 ensure_partition_operand_columns_in_target(&expr.rhs, target_partition_columns)?;
2777 }
2778
2779 Ok(())
2780}
2781
2782fn ensure_partition_operand_columns_in_target(
2783 operand: &Operand,
2784 target_partition_columns: &HashSet<&String>,
2785) -> Result<()> {
2786 match operand {
2787 Operand::Column(column) => ensure!(
2788 target_partition_columns.contains(column),
2789 InvalidPartitionRuleSnafu {
2790 reason: format!(
2791 "partition expression references column '{}' that is not in target partition columns",
2792 column
2793 )
2794 }
2795 ),
2796 Operand::Expr(expr) => {
2797 ensure_partition_operand_columns_in_target(&expr.lhs, target_partition_columns)?;
2798 ensure_partition_operand_columns_in_target(&expr.rhs, target_partition_columns)?;
2799 }
2800 Operand::Value(_) => {}
2801 }
2802
2803 Ok(())
2804}
2805
2806fn convert_one_expr(
2807 expr: &Expr,
2808 column_name_and_type: &HashMap<&String, ConcreteDataType>,
2809 timezone: &Timezone,
2810) -> Result<PartitionExpr> {
2811 let Expr::BinaryOp { left, op, right } = expr else {
2812 return InvalidPartitionRuleSnafu {
2813 reason: "partition rule must be a binary expression",
2814 }
2815 .fail();
2816 };
2817
2818 let op =
2819 RestrictedOp::try_from_parser(&op.clone()).with_context(|| InvalidPartitionRuleSnafu {
2820 reason: format!("unsupported operator in partition expr {op}"),
2821 })?;
2822
2823 let (lhs, op, rhs) = match (left.as_ref(), right.as_ref()) {
2825 (Expr::Identifier(ident), Expr::Value(value)) => {
2827 let (column_name, data_type) = convert_identifier(ident, column_name_and_type)?;
2828 let value = convert_value(&value.value, data_type, timezone, None)?;
2829 (Operand::Column(column_name), op, Operand::Value(value))
2830 }
2831 (Expr::Identifier(ident), Expr::UnaryOp { op: unary_op, expr })
2832 if let Expr::Value(v) = &**expr =>
2833 {
2834 let (column_name, data_type) = convert_identifier(ident, column_name_and_type)?;
2835 let value = convert_value(&v.value, data_type, timezone, Some(*unary_op))?;
2836 (Operand::Column(column_name), op, Operand::Value(value))
2837 }
2838 (Expr::Value(value), Expr::Identifier(ident)) => {
2840 let (column_name, data_type) = convert_identifier(ident, column_name_and_type)?;
2841 let value = convert_value(&value.value, data_type, timezone, None)?;
2842 (Operand::Value(value), op, Operand::Column(column_name))
2843 }
2844 (Expr::UnaryOp { op: unary_op, expr }, Expr::Identifier(ident))
2845 if let Expr::Value(v) = &**expr =>
2846 {
2847 let (column_name, data_type) = convert_identifier(ident, column_name_and_type)?;
2848 let value = convert_value(&v.value, data_type, timezone, Some(*unary_op))?;
2849 (Operand::Value(value), op, Operand::Column(column_name))
2850 }
2851 (Expr::BinaryOp { .. }, Expr::BinaryOp { .. }) => {
2852 let lhs = convert_one_expr(left, column_name_and_type, timezone)?;
2854 let rhs = convert_one_expr(right, column_name_and_type, timezone)?;
2855 (Operand::Expr(lhs), op, Operand::Expr(rhs))
2856 }
2857 _ => {
2858 return InvalidPartitionRuleSnafu {
2859 reason: format!("invalid partition expr {expr}"),
2860 }
2861 .fail();
2862 }
2863 };
2864
2865 Ok(PartitionExpr::new(lhs, op, rhs))
2866}
2867
2868fn convert_identifier(
2869 ident: &Ident,
2870 column_name_and_type: &HashMap<&String, ConcreteDataType>,
2871) -> Result<(String, ConcreteDataType)> {
2872 let column_name = ident.value.clone();
2873 let data_type = column_name_and_type
2874 .get(&column_name)
2875 .cloned()
2876 .with_context(|| ColumnNotFoundSnafu { msg: &column_name })?;
2877 Ok((column_name, data_type))
2878}
2879
2880fn convert_value(
2881 value: &ParserValue,
2882 data_type: ConcreteDataType,
2883 timezone: &Timezone,
2884 unary_op: Option<UnaryOperator>,
2885) -> Result<Value> {
2886 sql_value_to_value(
2887 &ColumnSchema::new("<NONAME>", data_type, true),
2888 value,
2889 Some(timezone),
2890 unary_op,
2891 false,
2892 )
2893 .context(error::SqlCommonSnafu)
2894}
2895
2896#[cfg(feature = "enterprise")]
2897async fn execute_undrop_table(
2898 table_metadata_manager: &TableMetadataManagerRef,
2899 procedure_executor: &ProcedureExecutorRef,
2900 cache_invalidator: &CacheInvalidatorRef,
2901 table_name: TableName,
2902 query_context: QueryContextRef,
2903) -> Result<Output> {
2904 ensure_table_definition_writable(&table_name.schema_name, &table_name.table_name)?;
2907
2908 let dropped = table_metadata_manager
2909 .get_dropped_table(&table_name)
2910 .await
2911 .context(TableMetadataManagerSnafu)?
2912 .with_context(|| TableNotFoundSnafu {
2913 table_name: table_name.to_string(),
2914 })?;
2915
2916 let executor_context = to_executor_context(query_context, TriggerReason::Manual);
2917 let request = SubmitDdlTaskRequest::new(DdlTask::new_undrop_table(dropped.table_id));
2918 procedure_executor
2919 .submit_ddl_task(executor_context, request)
2920 .await
2921 .context(error::ExecuteDdlSnafu)?;
2922
2923 if let Err(err) = cache_invalidator
2924 .invalidate(
2925 &Context::default(),
2926 &[
2927 CacheIdent::TableId(dropped.table_id),
2928 CacheIdent::TableName(table_name),
2929 ],
2930 )
2931 .await
2932 {
2933 warn!(
2934 "Failed to invalidate cache after restoring table '{}' (id={}): {}",
2935 dropped.table_name, dropped.table_id, err
2936 );
2937 }
2938
2939 Ok(Output::new_with_affected_rows(0))
2940}
2941
2942#[cfg(test)]
2943mod test {
2944 #[cfg(feature = "enterprise")]
2945 use std::sync::{Arc, Mutex};
2946 use std::time::Duration;
2947
2948 #[cfg(feature = "enterprise")]
2949 use api::v1::meta::{ProcedureDetailResponse, ReconcileRequest, ReconcileResponse};
2950 use arrow::datatypes::{DataType as ArrowDataType, Field as ArrowField};
2951 #[cfg(feature = "enterprise")]
2952 use common_meta::cache_invalidator::{CacheInvalidator, CacheInvalidatorRef};
2953 #[cfg(feature = "enterprise")]
2954 use common_meta::instruction::CacheIdent;
2955 #[cfg(feature = "enterprise")]
2956 use common_meta::key::table_route::TableRouteValue;
2957 #[cfg(feature = "enterprise")]
2958 use common_meta::key::test_utils::new_test_table_info_with_name;
2959 #[cfg(feature = "enterprise")]
2960 use common_meta::key::{TableMetadataManager, TableMetadataManagerRef};
2961 #[cfg(feature = "enterprise")]
2962 use common_meta::kv_backend::memory::MemoryKvBackend;
2963 #[cfg(feature = "enterprise")]
2964 use common_meta::procedure_executor::{
2965 ExecutorContext, ProcedureExecutor, ProcedureExecutorRef,
2966 };
2967 #[cfg(feature = "enterprise")]
2968 use common_meta::rpc::ddl::{DdlTask, SubmitDdlTaskRequest, SubmitDdlTaskResponse};
2969 #[cfg(feature = "enterprise")]
2970 use common_meta::rpc::procedure::{
2971 MigrateRegionRequest, MigrateRegionResponse, ProcedureStateResponse,
2972 };
2973 use datafusion::functions_aggregate::expr_fn::count;
2974 use datafusion::logical_expr::builder::LogicalTableSource;
2975 use datafusion::logical_expr::{LogicalPlanBuilder, col};
2976 use session::context::{QueryContext, QueryContextBuilder};
2977 use sql::dialect::GreptimeDbDialect;
2978 use sql::parser::{ParseOptions, ParserContext};
2979 use sql::statements::statement::Statement;
2980 use sqlparser::parser::Parser;
2981 use table::requests::REPARTITION_COLUMN_HINT_KEY;
2982
2983 use super::*;
2984 use crate::expr_helper;
2985
2986 #[cfg(feature = "enterprise")]
2987 #[derive(Default)]
2988 struct RecordingProcedureExecutor {
2989 requests: Mutex<Vec<SubmitDdlTaskRequest>>,
2990 contexts: Mutex<Vec<ExecutorContext>>,
2991 fail: bool,
2992 }
2993
2994 #[cfg(feature = "enterprise")]
2995 #[async_trait::async_trait]
2996 impl ProcedureExecutor for RecordingProcedureExecutor {
2997 async fn submit_ddl_task(
2998 &self,
2999 ctx: ExecutorContext,
3000 request: SubmitDdlTaskRequest,
3001 ) -> common_meta::error::Result<SubmitDdlTaskResponse> {
3002 self.contexts.lock().unwrap().push(ctx);
3003 self.requests.lock().unwrap().push(request);
3004 if self.fail {
3005 return common_meta::error::UnsupportedSnafu {
3006 operation: "test submit failure",
3007 }
3008 .fail();
3009 }
3010 Ok(SubmitDdlTaskResponse::default())
3011 }
3012
3013 async fn migrate_region(
3014 &self,
3015 _: &ExecutorContext,
3016 _: MigrateRegionRequest,
3017 ) -> common_meta::error::Result<MigrateRegionResponse> {
3018 unimplemented!()
3019 }
3020 async fn reconcile(
3021 &self,
3022 _: &ExecutorContext,
3023 _: ReconcileRequest,
3024 ) -> common_meta::error::Result<ReconcileResponse> {
3025 unimplemented!()
3026 }
3027 async fn query_procedure_state(
3028 &self,
3029 _: &ExecutorContext,
3030 _: &str,
3031 ) -> common_meta::error::Result<ProcedureStateResponse> {
3032 unimplemented!()
3033 }
3034 async fn list_procedures(
3035 &self,
3036 _: &ExecutorContext,
3037 ) -> common_meta::error::Result<ProcedureDetailResponse> {
3038 unimplemented!()
3039 }
3040 }
3041
3042 #[cfg(feature = "enterprise")]
3043 #[derive(Default)]
3044 struct RecordingCacheInvalidator {
3045 invalidations: Mutex<Vec<Vec<CacheIdent>>>,
3046 fail: bool,
3047 }
3048
3049 #[cfg(feature = "enterprise")]
3050 #[async_trait::async_trait]
3051 impl CacheInvalidator for RecordingCacheInvalidator {
3052 async fn invalidate(
3053 &self,
3054 _: &Context,
3055 caches: &[CacheIdent],
3056 ) -> common_meta::error::Result<()> {
3057 self.invalidations.lock().unwrap().push(caches.to_vec());
3058 if self.fail {
3059 return common_meta::error::UnsupportedSnafu {
3060 operation: "test cache failure",
3061 }
3062 .fail();
3063 }
3064 Ok(())
3065 }
3066
3067 fn invalidate_all(&self) -> common_meta::error::Result<()> {
3068 Ok(())
3069 }
3070 }
3071
3072 #[cfg(feature = "enterprise")]
3073 async fn dropped_table_manager(table_id: TableId, name: &TableName) -> TableMetadataManagerRef {
3074 let backend = Arc::new(MemoryKvBackend::default());
3075 let manager = Arc::new(TableMetadataManager::new(backend));
3076 let mut info = new_test_table_info_with_name(table_id, &name.table_name);
3077 info.catalog_name = name.catalog_name.clone();
3078 info.schema_name = name.schema_name.clone();
3079 let route = TableRouteValue::physical(vec![]);
3080 manager
3081 .create_table_metadata(info, route.clone(), HashMap::new())
3082 .await
3083 .unwrap();
3084 manager
3085 .delete_table_metadata(table_id, name, &route, &HashMap::new(), None)
3086 .await
3087 .unwrap();
3088 manager
3089 }
3090
3091 #[cfg(feature = "enterprise")]
3092 #[tokio::test]
3093 async fn test_undrop_table_execution_path() {
3094 let name = TableName::new("greptime", "public", "metrics");
3095 let manager = dropped_table_manager(42, &name).await;
3096 let mut live = new_test_table_info_with_name(99, "metrics");
3098 live.catalog_name = "greptime".to_string();
3099 live.schema_name = "public".to_string();
3100 manager
3101 .create_table_metadata(live, TableRouteValue::physical(vec![]), HashMap::new())
3102 .await
3103 .unwrap();
3104 let procedure = Arc::new(RecordingProcedureExecutor::default());
3105 let cache = Arc::new(RecordingCacheInvalidator::default());
3106
3107 execute_undrop_table(
3108 &manager,
3109 &(procedure.clone() as ProcedureExecutorRef),
3110 &(cache.clone() as CacheInvalidatorRef),
3111 name.clone(),
3112 QueryContext::arc(),
3113 )
3114 .await
3115 .unwrap();
3116
3117 let requests = procedure.requests.lock().unwrap();
3118 assert!(matches!(&requests[0].task, DdlTask::UndropTable(task) if task.table_id == 42));
3119 let contexts = procedure.contexts.lock().unwrap();
3120 assert_eq!(
3121 contexts[0].event_input.as_ref().map(|input| input.reason),
3122 Some(TriggerReason::Manual)
3123 );
3124 assert_eq!(
3125 cache.invalidations.lock().unwrap()[0],
3126 vec![CacheIdent::TableId(42), CacheIdent::TableName(name)]
3127 );
3128 }
3129
3130 #[cfg(feature = "enterprise")]
3131 #[tokio::test]
3132 async fn test_undrop_table_missing_tombstone_submits_nothing() {
3133 let manager = Arc::new(TableMetadataManager::new(Arc::new(
3134 MemoryKvBackend::default(),
3135 )));
3136 let procedure = Arc::new(RecordingProcedureExecutor::default());
3137 let cache = Arc::new(RecordingCacheInvalidator::default());
3138 let err = execute_undrop_table(
3139 &manager,
3140 &(procedure.clone() as ProcedureExecutorRef),
3141 &(cache as CacheInvalidatorRef),
3142 TableName::new("greptime", "public", "missing"),
3143 QueryContext::arc(),
3144 )
3145 .await
3146 .unwrap_err();
3147 assert!(matches!(err, crate::error::Error::TableNotFound { .. }));
3148 assert!(procedure.requests.lock().unwrap().is_empty());
3149 }
3150
3151 #[cfg(feature = "enterprise")]
3152 #[tokio::test]
3153 async fn test_undrop_table_submit_failure_does_not_invalidate() {
3154 let name = TableName::new("greptime", "public", "metrics");
3155 let manager = dropped_table_manager(42, &name).await;
3156 let procedure = Arc::new(RecordingProcedureExecutor {
3157 fail: true,
3158 ..Default::default()
3159 });
3160 let cache = Arc::new(RecordingCacheInvalidator::default());
3161 execute_undrop_table(
3162 &manager,
3163 &(procedure as ProcedureExecutorRef),
3164 &(cache.clone() as CacheInvalidatorRef),
3165 name,
3166 QueryContext::arc(),
3167 )
3168 .await
3169 .unwrap_err();
3170 assert!(cache.invalidations.lock().unwrap().is_empty());
3171 }
3172
3173 #[cfg(feature = "enterprise")]
3174 #[tokio::test]
3175 async fn test_undrop_table_cache_failure_returns_success_after_submit() {
3176 let name = TableName::new("greptime", "public", "metrics");
3177 let manager = dropped_table_manager(42, &name).await;
3178 let procedure = Arc::new(RecordingProcedureExecutor::default());
3179 let cache = Arc::new(RecordingCacheInvalidator {
3180 fail: true,
3181 ..Default::default()
3182 });
3183 execute_undrop_table(
3184 &manager,
3185 &(procedure.clone() as ProcedureExecutorRef),
3186 &(cache.clone() as CacheInvalidatorRef),
3187 name.clone(),
3188 QueryContext::arc(),
3189 )
3190 .await
3191 .unwrap();
3192 assert_eq!(procedure.requests.lock().unwrap().len(), 1);
3193 assert_eq!(
3194 cache.invalidations.lock().unwrap()[0],
3195 vec![CacheIdent::TableId(42), CacheIdent::TableName(name)]
3196 );
3197 }
3198
3199 #[test]
3200 fn test_parse_ddl_options() {
3201 let options = OptionMap::from([
3202 ("timeout".to_string(), "5m".to_string()),
3203 ("wait".to_string(), "false".to_string()),
3204 ]);
3205 let ddl_options = parse_ddl_options(&options).unwrap();
3206 assert!(!ddl_options.wait);
3207 assert_eq!(Duration::from_secs(300), ddl_options.timeout);
3208 }
3209
3210 #[test]
3211 fn test_validate_and_normalize_annotations() {
3212 let schema = Schema::new(vec![
3213 ColumnSchema::new("service_name", ConcreteDataType::string_datatype(), true),
3214 ColumnSchema::new("host_id", ConcreteDataType::string_datatype(), true),
3215 ColumnSchema::new("value", ConcreteDataType::float64_datatype(), true),
3216 ColumnSchema::new("payload", ConcreteDataType::binary_datatype(), true),
3217 ]);
3218 let opts = |pairs: &[(&str, &str)]| TableOptions {
3219 extra_options: pairs
3220 .iter()
3221 .map(|(k, v)| (k.to_string(), v.to_string()))
3222 .collect(),
3223 ..Default::default()
3224 };
3225 let check = |pairs: &[(&str, &str)]| {
3226 let mut options = opts(pairs);
3227 validate_and_normalize_annotations(&mut options, &schema, &[]).map(|()| options)
3228 };
3229
3230 for pairs in [
3232 &[(
3233 "greptime.semantic.entity.process.id",
3234 "service_name,host_id",
3235 )][..],
3236 &[("greptime.semantic.entity.service.id", "value")][..],
3237 ] {
3238 assert!(check(pairs).is_ok());
3239 }
3240
3241 let missing = check(&[("greptime.semantic.entity.service.id", "nope")]).unwrap_err();
3243 assert!(matches!(missing, error::Error::ColumnNotFound { .. }));
3244 assert_eq!(
3245 common_error::status_code::StatusCode::InvalidArguments,
3246 common_error::ext::ErrorExt::status_code(&missing)
3247 );
3248
3249 let binary_id = check(&[("greptime.semantic.entity.service.id", "payload")]).unwrap_err();
3250 assert!(matches!(binary_id, error::Error::InvalidSql { .. }));
3251
3252 let bad_value = check(&[("greptime.semantic.signal_type", "garbage")]).unwrap_err();
3255 assert!(matches!(bad_value, error::Error::InvalidSql { .. }));
3256
3257 let unknown_key = check(&[("greptime.semantic.nonsense", "x")]).unwrap_err();
3258 assert!(matches!(unknown_key, error::Error::InvalidSql { .. }));
3259 }
3260
3261 #[test]
3262 fn test_validate_and_normalize_flow_options_empty() {
3263 assert!(
3264 validate_and_normalize_flow_options(HashMap::new(), None)
3265 .unwrap()
3266 .is_empty()
3267 );
3268 }
3269
3270 #[test]
3271 fn test_validate_and_normalize_flow_options_valid() {
3272 let options = HashMap::from([
3273 (DEFER_ON_MISSING_SOURCE_KEY.to_string(), "TRUE".to_string()),
3274 (
3275 FLOW_EXPERIMENTAL_ENABLE_INCREMENTAL_READ_KEY.to_string(),
3276 "FALSE".to_string(),
3277 ),
3278 ]);
3279
3280 assert_eq!(
3281 validate_and_normalize_flow_options(options, None).unwrap(),
3282 HashMap::from([
3283 (DEFER_ON_MISSING_SOURCE_KEY.to_string(), "true".to_string(),),
3284 (
3285 FLOW_EXPERIMENTAL_ENABLE_INCREMENTAL_READ_KEY.to_string(),
3286 "false".to_string(),
3287 )
3288 ])
3289 );
3290 }
3291
3292 #[cfg(not(feature = "enterprise"))]
3293 #[test]
3294 fn test_validate_and_normalize_flow_options_unknown_option() {
3295 let err = validate_and_normalize_flow_options(
3296 HashMap::from([("foo".to_string(), "bar".to_string())]),
3297 None,
3298 )
3299 .unwrap_err();
3300
3301 assert!(
3302 err.to_string()
3303 .contains("unknown flow option 'foo', supported options:")
3304 );
3305 }
3306
3307 #[test]
3308 fn test_validate_and_normalize_flow_options_reserved_option() {
3309 let err = validate_and_normalize_flow_options(
3310 HashMap::from([(
3311 FlowType::FLOW_TYPE_KEY.to_string(),
3312 FlowType::BATCHING.to_string(),
3313 )]),
3314 None,
3315 )
3316 .unwrap_err();
3317
3318 assert!(
3319 err.to_string()
3320 .contains("flow option 'flow_type' is reserved for internal use")
3321 );
3322 }
3323
3324 #[test]
3325 fn test_validate_and_normalize_flow_options_invalid_bool() {
3326 let err = validate_and_normalize_flow_options(
3327 HashMap::from([(
3328 DEFER_ON_MISSING_SOURCE_KEY.to_string(),
3329 "not-a-bool".to_string(),
3330 )]),
3331 None,
3332 )
3333 .unwrap_err();
3334
3335 assert!(
3336 err.to_string()
3337 .contains("invalid flow option 'defer_on_missing_source': 'not-a-bool'")
3338 );
3339 }
3340
3341 #[test]
3342 fn test_parser_flow_option_access_key_id() {
3343 let sql = r"
3344CREATE FLOW task_6
3345SINK TO schema_1.table_1
3346WITH (access_key_id = [true])
3347AS
3348SELECT max(c1), min(c2) FROM schema_2.table_2;";
3349 let stmt =
3350 ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
3351 .unwrap()
3352 .pop()
3353 .unwrap();
3354
3355 let Statement::CreateFlow(create_flow) = stmt else {
3356 unreachable!()
3357 };
3358 let expr =
3359 expr_helper::to_create_flow_task_expr(create_flow, &QueryContext::arc()).unwrap();
3360
3361 #[cfg(not(feature = "enterprise"))]
3362 {
3363 let err = validate_and_normalize_flow_options(expr.flow_options, None).unwrap_err();
3364 assert!(
3365 err.to_string()
3366 .contains("unknown flow option 'access_key_id'")
3367 );
3368 }
3369 #[cfg(feature = "enterprise")]
3370 assert_eq!(
3371 validate_and_normalize_flow_options(expr.flow_options, None).unwrap(),
3372 HashMap::from([("access_key_id".to_string(), "['true']".to_string())])
3373 );
3374 }
3375
3376 fn stateless_test_scan(name: &str) -> LogicalPlan {
3377 let schema = arrow::datatypes::Schema::new(vec![ArrowField::new(
3378 "value",
3379 ArrowDataType::Int32,
3380 true,
3381 )]);
3382 LogicalPlanBuilder::scan(
3383 name,
3384 Arc::new(LogicalTableSource::new(Arc::new(schema))),
3385 None,
3386 )
3387 .unwrap()
3388 .build()
3389 .unwrap()
3390 }
3391
3392 #[test]
3393 fn test_is_stateless_flow_plan_accepts_scan_projection_and_filter() {
3394 let scan = stateless_test_scan("source");
3395 assert!(is_stateless_flow_plan(&scan));
3396
3397 let projection = LogicalPlanBuilder::from(scan.clone())
3398 .project(vec![col("value")])
3399 .unwrap()
3400 .build()
3401 .unwrap();
3402 assert!(is_stateless_flow_plan(&projection));
3403
3404 let filter = LogicalPlanBuilder::from(projection)
3405 .filter(col("value").gt(datafusion_expr::lit(0)))
3406 .unwrap()
3407 .build()
3408 .unwrap();
3409 assert!(is_stateless_flow_plan(&filter));
3410 }
3411
3412 #[test]
3413 fn test_is_stateless_flow_plan_rejects_aggregate_and_multiple_scans() {
3414 let scan = stateless_test_scan("source");
3415 let aggregate = LogicalPlanBuilder::from(scan.clone())
3416 .aggregate(
3417 Vec::<datafusion_expr::Expr>::new(),
3418 vec![count(col("value"))],
3419 )
3420 .unwrap()
3421 .build()
3422 .unwrap();
3423 assert!(!is_stateless_flow_plan(&aggregate));
3424
3425 let distinct = LogicalPlanBuilder::from(scan.clone())
3426 .distinct()
3427 .unwrap()
3428 .build()
3429 .unwrap();
3430 assert!(is_stateless_flow_plan(&distinct));
3431
3432 let distinct_on_aggregate = LogicalPlanBuilder::from(aggregate)
3433 .distinct()
3434 .unwrap()
3435 .build()
3436 .unwrap();
3437 assert!(!is_stateless_flow_plan(&distinct_on_aggregate));
3438
3439 let multiple_scans = LogicalPlanBuilder::from(stateless_test_scan("left"))
3440 .cross_join(stateless_test_scan("right"))
3441 .unwrap()
3442 .build()
3443 .unwrap();
3444 assert!(!is_stateless_flow_plan(&multiple_scans));
3445 }
3446
3447 #[test]
3450 fn test_eval_interval_rejected_non_positive() {
3451 let err = validate_and_normalize_flow_options(HashMap::new(), Some(0)).unwrap_err();
3453 assert!(err.to_string().contains("EVAL INTERVAL must be positive"));
3454
3455 let err = validate_and_normalize_flow_options(HashMap::new(), Some(-5)).unwrap_err();
3457 assert!(err.to_string().contains("EVAL INTERVAL must be positive"));
3458
3459 let result = validate_and_normalize_flow_options(HashMap::new(), Some(300));
3461 assert!(result.is_ok());
3462 }
3463
3464 #[test]
3465 fn test_internal_schedule_key_rejected() {
3466 let err = validate_and_normalize_flow_options(
3467 HashMap::from([(INTERNAL_EVAL_SCHEDULE_KEY.to_string(), "value".to_string())]),
3468 Some(300),
3469 )
3470 .unwrap_err();
3471 assert!(err.to_string().contains("reserved for internal use"));
3472 }
3473
3474 #[cfg(not(feature = "enterprise"))]
3475 #[test]
3476 fn test_schedule_keys_rejected_as_unknown_options() {
3477 for key in [
3478 "eval_interval_anchor",
3479 "eval_interval_start",
3480 "eval_interval_missed_tick_policy",
3481 "eval_interval_catchup_max_runs",
3482 "eval_interval_catchup_max_lag",
3483 ] {
3484 let err = validate_and_normalize_flow_options(
3485 HashMap::from([(key.to_string(), "value".to_string())]),
3486 Some(300),
3487 )
3488 .unwrap_err();
3489 assert!(
3490 err.to_string()
3491 .contains(&format!("unknown flow option '{key}'"))
3492 );
3493 }
3494 }
3495
3496 #[cfg(feature = "enterprise")]
3497 #[test]
3498 fn test_schedule_keys_preserved() {
3499 let options = HashMap::from([
3500 ("eval_interval_anchor".to_string(), "anchor".to_string()),
3501 ("eval_interval_start".to_string(), "start".to_string()),
3502 (
3503 "eval_interval_missed_tick_policy".to_string(),
3504 "missed".to_string(),
3505 ),
3506 (
3507 "eval_interval_catchup_max_runs".to_string(),
3508 "runs".to_string(),
3509 ),
3510 (
3511 "eval_interval_catchup_max_lag".to_string(),
3512 "lag".to_string(),
3513 ),
3514 ]);
3515 assert_eq!(
3516 validate_and_normalize_flow_options(options.clone(), Some(300)).unwrap(),
3517 options
3518 );
3519 }
3520
3521 #[test]
3522 fn test_internal_transport_keys_rejected_as_reserved() {
3523 for key in [
3524 "__greptime_internal_eval_schedule",
3525 "__greptime_internal_eval_offset_secs",
3526 ] {
3527 let err = validate_and_normalize_flow_options(
3528 HashMap::from([(key.to_string(), "120".to_string())]),
3529 Some(300),
3530 )
3531 .unwrap_err();
3532
3533 assert!(
3534 err.to_string()
3535 .contains(&format!("flow option '{key}' is reserved for internal use")),
3536 "unexpected error for {key}: {err}"
3537 );
3538 }
3539 }
3540
3541 #[test]
3542 fn test_determine_flow_type_for_source_state_missing_sources_require_opt_in() {
3543 let err =
3544 determine_flow_type_for_source_state("my_flow", &HashMap::new(), true, false, false)
3545 .unwrap_err();
3546
3547 assert!(err.to_string().contains(
3548 "missing source tables for flow 'my_flow'; use WITH (defer_on_missing_source = true) to create a pending flow"
3549 ));
3550 }
3551
3552 #[test]
3553 fn test_determine_flow_type_for_source_state_missing_sources_prefer_batching() {
3554 let flow_options =
3555 HashMap::from([(DEFER_ON_MISSING_SOURCE_KEY.to_string(), "true".to_string())]);
3556
3557 assert_eq!(
3558 determine_flow_type_for_source_state("my_flow", &flow_options, true, true, false)
3559 .unwrap(),
3560 Some(FlowType::Batching)
3561 );
3562 let err = determine_flow_type_for_source_state("my_flow", &flow_options, true, true, true)
3563 .unwrap_err();
3564 assert!(err.to_string().contains(
3565 "flow 'my_flow' with EVAL INTERVAL requires the batching scheduler, but source tables with ttl=instant are not supported under batching mode"
3566 ));
3567 }
3568
3569 #[test]
3570 fn test_determine_flow_type_for_source_state_existing_sources_require_plan() {
3571 assert_eq!(
3572 determine_flow_type_for_source_state("my_flow", &HashMap::new(), false, true, false)
3573 .unwrap(),
3574 None
3575 );
3576 }
3577
3578 #[test]
3579 fn test_determine_flow_type_for_source_state_instant_ttl_with_eval_interval_is_rejected() {
3580 let err =
3585 determine_flow_type_for_source_state("my_flow", &HashMap::new(), false, true, true)
3586 .unwrap_err();
3587 assert!(
3588 err.to_string().contains(
3589 "flow 'my_flow' with EVAL INTERVAL requires the batching scheduler, but source tables with ttl=instant are not supported under batching mode"
3590 ),
3591 "unexpected error: {err}"
3592 );
3593 }
3594
3595 #[test]
3596 fn test_name_is_match() {
3597 assert!(!NAME_PATTERN_REG.is_match("/adaf"));
3598 assert!(!NAME_PATTERN_REG.is_match("🈲"));
3599 assert!(NAME_PATTERN_REG.is_match("hello"));
3600 assert!(NAME_PATTERN_REG.is_match("test@"));
3601 assert!(!NAME_PATTERN_REG.is_match("@test"));
3602 assert!(NAME_PATTERN_REG.is_match("test#"));
3603 assert!(!NAME_PATTERN_REG.is_match("#test"));
3604 assert!(!NAME_PATTERN_REG.is_match("@"));
3605 assert!(!NAME_PATTERN_REG.is_match("#"));
3606 }
3607
3608 #[test]
3609 fn test_partition_expr_equivalence_with_swapped_operands() {
3610 let column_name = "device_id".to_string();
3611 let column_name_and_type =
3612 HashMap::from([(&column_name, ConcreteDataType::int32_datatype())]);
3613 let timezone = Timezone::from_tz_string("UTC").unwrap();
3614 let dialect = GreptimeDbDialect {};
3615
3616 let mut parser = Parser::new(&dialect)
3617 .try_with_sql("device_id < 100")
3618 .unwrap();
3619 let expr_left = parser.parse_expr().unwrap();
3620
3621 let mut parser = Parser::new(&dialect)
3622 .try_with_sql("100 > device_id")
3623 .unwrap();
3624 let expr_right = parser.parse_expr().unwrap();
3625
3626 let partition_left =
3627 convert_one_expr(&expr_left, &column_name_and_type, &timezone).unwrap();
3628 let partition_right =
3629 convert_one_expr(&expr_right, &column_name_and_type, &timezone).unwrap();
3630
3631 assert_eq!(partition_left, partition_right);
3632 assert!([partition_left.clone()].contains(&partition_right));
3633
3634 let mut physical_partition_exprs = vec![partition_left];
3635 let mut logical_partition_exprs = vec![partition_right];
3636 physical_partition_exprs.sort_unstable();
3637 logical_partition_exprs.sort_unstable();
3638 assert_eq!(physical_partition_exprs, logical_partition_exprs);
3639 }
3640
3641 #[test]
3642 fn test_repartition_target_partition_columns_are_overwrite_context() {
3643 let device_id = ColumnSchema::new("device_id", ConcreteDataType::int32_datatype(), true);
3644 let area = ColumnSchema::new("area", ConcreteDataType::string_datatype(), true);
3645 let existing_partition_columns = vec![&device_id];
3646 let target_partition_columns = vec![&device_id, &area];
3647 let existing_column_name_and_type = column_name_and_type(&existing_partition_columns);
3648 let target_column_name_and_type = column_name_and_type(&target_partition_columns);
3649 let timezone = Timezone::from_tz_string("UTC").unwrap();
3650 let dialect = GreptimeDbDialect {};
3651
3652 let mut parser = Parser::new(&dialect)
3653 .try_with_sql("device_id < 100 AND area < 'South'")
3654 .unwrap();
3655 let expr = parser.parse_expr().unwrap();
3656
3657 let err = convert_one_expr(&expr, &existing_column_name_and_type, &timezone).unwrap_err();
3658 assert!(err.to_string().contains("area"));
3659
3660 let partition_expr = convert_one_expr(&expr, &target_column_name_and_type, &timezone)
3661 .expect("target columns should overwrite the conversion context");
3662 let partition_expr = partition_expr.to_string();
3663 assert!(partition_expr.contains("device_id"));
3664 assert!(partition_expr.contains("area"));
3665 assert!(partition_expr.contains("South"));
3666 }
3667
3668 #[test]
3669 fn test_repartition_rejects_remaining_expr_outside_target_columns() {
3670 let device_id = "device_id".to_string();
3671 let area = "area".to_string();
3672 let timezone = Timezone::from_tz_string("UTC").unwrap();
3673 let column_name_and_type = HashMap::from([
3674 (&device_id, ConcreteDataType::int32_datatype()),
3675 (&area, ConcreteDataType::string_datatype()),
3676 ]);
3677 let dialect = GreptimeDbDialect {};
3678 let mut parser = Parser::new(&dialect)
3679 .try_with_sql("device_id >= 100")
3680 .unwrap();
3681 let remaining_old_expr = convert_one_expr(
3682 &parser.parse_expr().unwrap(),
3683 &column_name_and_type,
3684 &timezone,
3685 )
3686 .unwrap();
3687 let target_partition_columns = HashSet::from([&area]);
3688
3689 let err = ensure_partition_expr_columns_in_target(
3690 &[remaining_old_expr],
3691 &target_partition_columns,
3692 )
3693 .unwrap_err();
3694
3695 assert!(err.to_string().contains("device_id"));
3696 assert!(err.to_string().contains("target partition columns"));
3697 }
3698
3699 #[test]
3700 fn test_repartition_rejects_duplicate_target_partition_columns() {
3701 let device_id = ColumnSchema::new("device_id", ConcreteDataType::int32_datatype(), true);
3702 let column_schemas = vec![device_id];
3703 let target_partition_columns = vec!["device_id".to_string(), "device_id".to_string()];
3704
3705 let err =
3706 validate_and_collect_partition_columns(&target_partition_columns, &column_schemas)
3707 .unwrap_err();
3708
3709 assert!(err.to_string().contains("duplicate partition column"));
3710 assert!(err.to_string().contains("device_id"));
3711 }
3712
3713 fn create_expr_from_sql(sql: &str) -> CreateTableExpr {
3714 let result =
3715 ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
3716 .unwrap();
3717
3718 match &result[0] {
3719 Statement::CreateTable(create) => {
3720 expr_helper::create_to_expr(create, &QueryContext::arc()).unwrap()
3721 }
3722 _ => unreachable!(),
3723 }
3724 }
3725
3726 #[test]
3727 fn test_create_table_with_repartition_partition_num_hint() {
3728 let key = table::requests::REPARTITION_PARTITION_NUM_HINT_KEY;
3729 for (value, expected) in [
3730 (" 8 ", Some("8")),
3731 ("0", None),
3732 ("-1", None),
3733 ("4294967296", None),
3734 ] {
3735 let expr = create_expr_from_sql(&format!(
3736 "CREATE TABLE metrics (host STRING, ts TIMESTAMP TIME INDEX, PRIMARY KEY(host)) WITH ('{key}' = '{value}')"
3737 ));
3738 let result = create_table_info(&expr, vec![]);
3739 if let Some(expected) = expected {
3740 let info = result.unwrap();
3741 assert_eq!(
3742 info.meta.options.extra_options.get(key).map(String::as_str),
3743 Some(expected)
3744 );
3745 } else {
3746 assert!(
3747 result
3748 .unwrap_err()
3749 .to_string()
3750 .contains("expects a positive integer")
3751 );
3752 }
3753 }
3754 }
3755
3756 #[test]
3757 fn test_create_table_with_repartition_column_hint() {
3758 let expr = create_expr_from_sql(
3759 r"
3760CREATE TABLE metrics (
3761 host STRING,
3762 ts TIMESTAMP TIME INDEX,
3763 cpu DOUBLE,
3764 PRIMARY KEY(host)
3765)
3766WITH ('repartition.column.hint' = ' host ')",
3767 );
3768
3769 let table_info = create_table_info(&expr, vec![]).unwrap();
3770 assert_eq!(
3771 table_info
3772 .meta
3773 .options
3774 .extra_options
3775 .get(REPARTITION_COLUMN_HINT_KEY),
3776 Some(&"host".to_string())
3777 );
3778 }
3779
3780 #[test]
3781 fn test_create_table_info_rejects_non_timestamp_time_index() {
3782 let expr = CreateTableExpr {
3783 catalog_name: "greptime".to_string(),
3784 schema_name: "public".to_string(),
3785 table_name: "demo".to_string(),
3786 desc: String::new(),
3787 column_defs: vec![api::v1::ColumnDef {
3788 name: "host".to_string(),
3789 data_type: api::v1::ColumnDataType::String as i32,
3790 is_nullable: true,
3791 default_constraint: vec![],
3792 semantic_type: 0,
3793 comment: String::new(),
3794 datatype_extension: None,
3795 options: None,
3796 }],
3797 time_index: "host".to_string(),
3798 primary_keys: vec![],
3799 create_if_not_exists: false,
3800 table_options: HashMap::new(),
3801 table_id: None,
3802 engine: "mito".to_string(),
3803 };
3804
3805 let err = create_table_info(&expr, vec![]).unwrap_err();
3806 assert_eq!(
3807 common_error::ext::ErrorExt::status_code(&err),
3808 common_error::status_code::StatusCode::InvalidArguments
3809 );
3810 }
3811
3812 #[test]
3813 fn test_verify_alter_guards_entity_column_types() {
3814 let expr = create_expr_from_sql(
3815 "CREATE TABLE t (svc STRING, ts TIMESTAMP TIME INDEX) \
3816 WITH ('greptime.semantic.entity.service.id'='svc');",
3817 );
3818 let table_info = Arc::new(create_table_info(&expr, vec![]).unwrap());
3819 let alter = |kind| AlterTableExpr {
3820 catalog_name: "greptime".to_string(),
3821 schema_name: "public".to_string(),
3822 table_name: "t".to_string(),
3823 kind: Some(kind),
3824 };
3825
3826 let modify = alter(Kind::ModifyColumnTypes(api::v1::ModifyColumnTypes {
3827 modify_column_types: vec![api::v1::ModifyColumnType {
3828 column_name: "svc".to_string(),
3829 target_type: api::v1::ColumnDataType::Binary as i32,
3830 target_type_extension: None,
3831 }],
3832 }));
3833 let err = verify_alter(1, table_info.clone(), modify).unwrap_err();
3834 let msg = common_error::ext::ErrorExt::output_msg(&err);
3835 assert!(
3836 msg.contains("must keep a type that renders as a string"),
3837 "{msg}"
3838 );
3839
3840 let drop = alter(Kind::DropColumns(api::v1::DropColumns {
3843 drop_columns: vec![api::v1::DropColumn {
3844 name: "svc".to_string(),
3845 }],
3846 }));
3847 assert!(verify_alter(1, table_info, drop).unwrap());
3848 }
3849
3850 #[test]
3851 fn test_json2_requires_append_mode() {
3852 let cases = [
3853 "CREATE TABLE monitor (payload JSON2, ts TIMESTAMP TIME INDEX);",
3854 "CREATE TABLE monitor (payload JSON2, ts TIMESTAMP TIME INDEX) WITH (append_mode='false');",
3855 ];
3856
3857 for sql in cases {
3858 let expr = create_expr_from_sql(sql);
3859 let err = create_table_info(&expr, vec![]).unwrap_err();
3860 assert!(
3861 err.to_string()
3862 .contains("JSON2 column `payload` requires append_mode='true'"),
3863 "{err}"
3864 );
3865 }
3866
3867 let expr = create_expr_from_sql(
3868 "CREATE TABLE monitor (payload JSON2, ts TIMESTAMP TIME INDEX) WITH (append_mode='true');",
3869 );
3870 create_table_info(&expr, vec![]).unwrap();
3871 }
3872
3873 #[test]
3874 fn test_create_table_with_empty_repartition_column_hint() {
3875 let expr = create_expr_from_sql(
3876 r"
3877CREATE TABLE metrics (
3878 host STRING,
3879 ts TIMESTAMP TIME INDEX,
3880 cpu DOUBLE,
3881 PRIMARY KEY(host)
3882)
3883WITH ('repartition.column.hint' = '')",
3884 );
3885
3886 let err = create_table_info(&expr, vec![]).unwrap_err();
3887 assert!(
3888 err.to_string()
3889 .contains("repartition.column.hint expects exactly one column name")
3890 );
3891 }
3892
3893 #[test]
3894 fn test_create_table_with_multiple_repartition_column_hints() {
3895 let expr = create_expr_from_sql(
3896 r"
3897CREATE TABLE metrics (
3898 host STRING,
3899 region_id STRING,
3900 ts TIMESTAMP TIME INDEX,
3901 cpu DOUBLE,
3902 PRIMARY KEY(host)
3903)
3904WITH ('repartition.column.hint' = 'host,region_id')",
3905 );
3906
3907 let err = create_table_info(&expr, vec![]).unwrap_err();
3908 assert!(
3909 err.to_string()
3910 .contains("repartition.column.hint expects exactly one column name")
3911 );
3912 }
3913
3914 #[test]
3915 fn test_create_table_with_missing_repartition_column_hint() {
3916 let expr = create_expr_from_sql(
3917 r"
3918CREATE TABLE metrics (
3919 host STRING,
3920 ts TIMESTAMP TIME INDEX,
3921 cpu DOUBLE,
3922 PRIMARY KEY(host)
3923)
3924WITH ('repartition.column.hint' = 'region_id')",
3925 );
3926
3927 let err = create_table_info(&expr, vec![]).unwrap_err();
3928 assert!(
3929 err.to_string()
3930 .contains("Cannot find column by name: region")
3931 );
3932 }
3933
3934 #[test]
3935 fn test_create_table_with_time_index_repartition_column_hint() {
3936 let expr = create_expr_from_sql(
3937 r"
3938CREATE TABLE metrics (
3939 host STRING,
3940 ts TIMESTAMP TIME INDEX,
3941 cpu DOUBLE,
3942 PRIMARY KEY(host)
3943)
3944WITH ('repartition.column.hint' = 'ts')",
3945 );
3946
3947 let err = create_table_info(&expr, vec![]).unwrap_err();
3948 assert!(
3949 err.to_string()
3950 .contains("cannot set repartition.column.hint to the time index column")
3951 );
3952 }
3953
3954 #[test]
3955 fn test_create_partitioned_table_with_repartition_column_hint() {
3956 let expr = create_expr_from_sql(
3957 r"
3958CREATE TABLE metrics (
3959 host STRING,
3960 ts TIMESTAMP TIME INDEX,
3961 cpu DOUBLE,
3962 PRIMARY KEY(host)
3963)
3964WITH ('repartition.column.hint' = 'host')",
3965 );
3966
3967 let err = create_table_info(&expr, vec!["host".to_string()]).unwrap_err();
3968 assert!(
3969 err.to_string()
3970 .contains("cannot set repartition.column.hint on a table with partition metadata")
3971 );
3972 }
3973
3974 #[tokio::test]
3975 #[ignore = "TODO(ruihang): WIP new partition rule"]
3976 async fn test_parse_partitions() {
3977 common_telemetry::init_default_ut_logging();
3978 let cases = [
3979 (
3980 r"
3981CREATE TABLE rcx ( a INT, b STRING, c TIMESTAMP, TIME INDEX (c) )
3982PARTITION ON COLUMNS (b) (
3983 b < 'hz',
3984 b >= 'hz' AND b < 'sh',
3985 b >= 'sh'
3986)
3987ENGINE=mito",
3988 r#"[{"column_list":["b"],"value_list":["{\"Value\":{\"String\":\"hz\"}}"]},{"column_list":["b"],"value_list":["{\"Value\":{\"String\":\"sh\"}}"]},{"column_list":["b"],"value_list":["\"MaxValue\""]}]"#,
3989 ),
3990 (
3991 r"
3992CREATE TABLE rcx ( a INT, b STRING, c TIMESTAMP, TIME INDEX (c) )
3993PARTITION BY RANGE COLUMNS (b, a) (
3994 PARTITION r0 VALUES LESS THAN ('hz', 10),
3995 b < 'hz' AND a < 10,
3996 b >= 'hz' AND b < 'sh' AND a >= 10 AND a < 20,
3997 b >= 'sh' AND a >= 20
3998)
3999ENGINE=mito",
4000 r#"[{"column_list":["b","a"],"value_list":["{\"Value\":{\"String\":\"hz\"}}","{\"Value\":{\"Int32\":10}}"]},{"column_list":["b","a"],"value_list":["{\"Value\":{\"String\":\"sh\"}}","{\"Value\":{\"Int32\":20}}"]},{"column_list":["b","a"],"value_list":["\"MaxValue\"","\"MaxValue\""]}]"#,
4001 ),
4002 ];
4003 let ctx = QueryContextBuilder::default().build().into();
4004 for (sql, expected) in cases {
4005 let result = ParserContext::create_with_dialect(
4006 sql,
4007 &GreptimeDbDialect {},
4008 ParseOptions::default(),
4009 )
4010 .unwrap();
4011 match &result[0] {
4012 Statement::CreateTable(c) => {
4013 let expr = expr_helper::create_to_expr(c, &QueryContext::arc()).unwrap();
4014 let (partitions, _) =
4015 parse_partitions(&expr, c.partitions.clone(), &ctx).unwrap();
4016 let json = serde_json::to_string(&partitions).unwrap();
4017 assert_eq!(json, expected);
4018 }
4019 _ => unreachable!(),
4020 }
4021 }
4022 }
4023}