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