Skip to main content

common_meta/
ddl_manager.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::sync::Arc;
16use std::time::Duration;
17
18use api::v1::alter_table_expr::Kind;
19use api::v1::repartition::Source as PbRepartitionSource;
20use api::v1::{PartitionExprs, Repartition};
21use common_error::ext::BoxedError;
22use common_procedure::{
23    BoxedProcedure, BoxedProcedureLoader, Output, ProcedureId, ProcedureManagerRef,
24    ProcedureWithId, watcher,
25};
26use common_telemetry::tracing_context::{FutureExt, TracingContext};
27use common_telemetry::{debug, info, tracing};
28use derive_builder::Builder;
29use snafu::{OptionExt, ResultExt, ensure};
30use store_api::storage::TableId;
31use table::table_name::TableName;
32
33use crate::ddl::alter_database::AlterDatabaseProcedure;
34use crate::ddl::alter_logical_tables::AlterLogicalTablesProcedure;
35use crate::ddl::alter_table::AlterTableProcedure;
36use crate::ddl::comment_on::CommentOnProcedure;
37use crate::ddl::create_database::CreateDatabaseProcedure;
38use crate::ddl::create_flow::CreateFlowProcedure;
39use crate::ddl::create_logical_tables::CreateLogicalTablesProcedure;
40use crate::ddl::create_table::CreateTableProcedure;
41use crate::ddl::create_view::CreateViewProcedure;
42use crate::ddl::drop_database::DropDatabaseProcedure;
43use crate::ddl::drop_flow::DropFlowProcedure;
44use crate::ddl::drop_table::DropTableProcedure;
45use crate::ddl::drop_view::DropViewProcedure;
46use crate::ddl::truncate_table::TruncateTableProcedure;
47use crate::ddl::{DdlContext, utils};
48use crate::error::{
49    self, CreateRepartitionProcedureSnafu, EmptyDdlTasksSnafu, ProcedureOutputSnafu,
50    RegisterProcedureLoaderSnafu, RegisterRepartitionProcedureLoaderSnafu, Result,
51    SubmitProcedureSnafu, TableInfoNotFoundSnafu, TableNotFoundSnafu, TableRouteNotFoundSnafu,
52    UnexpectedLogicalRouteTableSnafu, WaitProcedureSnafu,
53};
54use crate::key::table_info::TableInfoValue;
55use crate::key::table_name::TableNameKey;
56use crate::key::{DeserializedValueWithBytes, TableMetadataManagerRef};
57use crate::procedure_executor::ExecutorContext;
58#[cfg(feature = "enterprise")]
59use crate::rpc::ddl::DdlTask::CreateTrigger;
60#[cfg(feature = "enterprise")]
61use crate::rpc::ddl::DdlTask::DropTrigger;
62use crate::rpc::ddl::DdlTask::{
63    AlterDatabase, AlterLogicalTables, AlterTable, CommentOn, CreateDatabase, CreateFlow,
64    CreateLogicalTables, CreateTable, CreateView, DropDatabase, DropFlow, DropLogicalTables,
65    DropTable, DropView, TruncateTable,
66};
67#[cfg(feature = "enterprise")]
68use crate::rpc::ddl::trigger::CreateTriggerTask;
69#[cfg(feature = "enterprise")]
70use crate::rpc::ddl::trigger::DropTriggerTask;
71use crate::rpc::ddl::{
72    AlterDatabaseTask, AlterTableTask, CommentOnTask, CreateDatabaseTask, CreateFlowTask,
73    CreateTableTask, CreateViewTask, DropDatabaseTask, DropFlowTask, DropTableTask, DropViewTask,
74    QueryContext, SubmitDdlTaskRequest, SubmitDdlTaskResponse, TruncateTableTask,
75};
76
77/// A configurator that customizes or enhances a [`DdlManager`].
78#[async_trait::async_trait]
79pub trait DdlManagerConfigurator<C>: Send + Sync {
80    /// Configures the given [`DdlManager`] using the provided [`DdlManagerConfigureContext`].
81    async fn configure(
82        &self,
83        ddl_manager: DdlManager,
84        ctx: C,
85    ) -> std::result::Result<DdlManager, BoxedError>;
86}
87
88pub type DdlManagerConfiguratorRef<C> = Arc<dyn DdlManagerConfigurator<C>>;
89
90pub type DdlManagerRef = Arc<DdlManager>;
91
92pub type BoxedProcedureLoaderFactory = dyn Fn(DdlContext) -> BoxedProcedureLoader;
93
94/// The [DdlManager] provides the ability to execute Ddl.
95#[derive(Builder)]
96pub struct DdlManager {
97    ddl_context: DdlContext,
98    procedure_manager: ProcedureManagerRef,
99    repartition_procedure_factory: RepartitionProcedureFactoryRef,
100    #[cfg(feature = "enterprise")]
101    trigger_ddl_manager: Option<TriggerDdlManagerRef>,
102}
103
104/// This trait is responsible for handling DDL tasks about triggers. e.g.,
105/// create trigger, drop trigger, etc.
106#[cfg(feature = "enterprise")]
107#[async_trait::async_trait]
108pub trait TriggerDdlManager: Send + Sync {
109    async fn create_trigger(
110        &self,
111        create_trigger_task: CreateTriggerTask,
112        procedure_manager: ProcedureManagerRef,
113        ddl_context: DdlContext,
114        query_context: QueryContext,
115    ) -> Result<SubmitDdlTaskResponse>;
116
117    async fn drop_trigger(
118        &self,
119        drop_trigger_task: DropTriggerTask,
120        procedure_manager: ProcedureManagerRef,
121        ddl_context: DdlContext,
122        query_context: QueryContext,
123    ) -> Result<SubmitDdlTaskResponse>;
124
125    fn as_any(&self) -> &dyn std::any::Any;
126}
127
128#[cfg(feature = "enterprise")]
129pub type TriggerDdlManagerRef = Arc<dyn TriggerDdlManager>;
130
131macro_rules! procedure_loader_entry {
132    ($procedure:ident) => {
133        (
134            $procedure::TYPE_NAME,
135            &|context: DdlContext| -> BoxedProcedureLoader {
136                Box::new(move |json: &str| {
137                    let context = context.clone();
138                    $procedure::from_json(json, context).map(|p| Box::new(p) as _)
139                })
140            },
141        )
142    };
143}
144
145macro_rules! procedure_loader {
146    ($($procedure:ident),*) => {
147        vec![
148            $(procedure_loader_entry!($procedure)),*
149        ]
150    };
151}
152
153pub type RepartitionProcedureFactoryRef = Arc<dyn RepartitionProcedureFactory>;
154
155pub enum RepartitionSource {
156    Partitioned { exprs: Vec<String> },
157    Unpartitioned { partition_columns: Vec<String> },
158}
159
160pub trait RepartitionProcedureFactory: Send + Sync {
161    fn create(
162        &self,
163        ddl_ctx: &DdlContext,
164        table_name: TableName,
165        table_id: TableId,
166        source: RepartitionSource,
167        to_exprs: Vec<String>,
168        timeout: Option<Duration>,
169    ) -> std::result::Result<BoxedProcedure, BoxedError>;
170
171    fn register_loaders(
172        &self,
173        ddl_ctx: &DdlContext,
174        procedure_manager: &ProcedureManagerRef,
175    ) -> std::result::Result<(), BoxedError>;
176}
177
178/// The options for DDL tasks.
179///
180/// Note: These options may not be utilized by all procedures.
181/// At present, they are specifically applied in `RepartitionProcedure`.
182#[derive(Debug, Clone, Copy)]
183pub struct DdlOptions {
184    /// The timeout will be passed to the procedure.
185    ///
186    /// Note: Each procedure may implement its own timeout handling mechanism.
187    pub timeout: Duration,
188    /// The flag that controls whether to wait for the procedure to complete.
189    ///
190    /// If wait is `true`, the procedure will wait for completion(success or failure) and the result will be returned.
191    /// Otherwise, the procedure will be submitted and return the [ProcedureId](common_procedure::ProcedureId) immediately.
192    ///
193    /// Note: The value of `wait` is independent of the `timeout` option. If a procedure ignores the `timeout` and `wait` is set to true, the operation returns until the procedure completes.
194    pub wait: bool,
195}
196
197impl DdlManager {
198    /// Returns a new [DdlManager] with all Ddl [BoxedProcedureLoader](common_procedure::procedure::BoxedProcedureLoader)s registered.
199    pub fn try_new(
200        ddl_context: DdlContext,
201        procedure_manager: ProcedureManagerRef,
202        repartition_procedure_factory: RepartitionProcedureFactoryRef,
203        register_loaders: bool,
204    ) -> Result<Self> {
205        let manager = Self {
206            ddl_context,
207            procedure_manager,
208            repartition_procedure_factory,
209            #[cfg(feature = "enterprise")]
210            trigger_ddl_manager: None,
211        };
212        if register_loaders {
213            manager.register_loaders()?;
214        }
215        Ok(manager)
216    }
217
218    #[cfg(feature = "enterprise")]
219    pub fn with_trigger_ddl_manager(mut self, trigger_ddl_manager: TriggerDdlManagerRef) -> Self {
220        self.trigger_ddl_manager = Some(trigger_ddl_manager);
221        self
222    }
223
224    /// Returns the [TableMetadataManagerRef].
225    pub fn table_metadata_manager(&self) -> &TableMetadataManagerRef {
226        &self.ddl_context.table_metadata_manager
227    }
228
229    /// Returns the [DdlContext]
230    pub fn create_context(&self) -> DdlContext {
231        self.ddl_context.clone()
232    }
233
234    /// Registers all Ddl loaders.
235    pub fn register_loaders(&self) -> Result<()> {
236        let loaders: Vec<(&str, &BoxedProcedureLoaderFactory)> = procedure_loader!(
237            CreateTableProcedure,
238            CreateLogicalTablesProcedure,
239            CreateViewProcedure,
240            CreateFlowProcedure,
241            AlterTableProcedure,
242            AlterLogicalTablesProcedure,
243            AlterDatabaseProcedure,
244            DropTableProcedure,
245            DropFlowProcedure,
246            TruncateTableProcedure,
247            CreateDatabaseProcedure,
248            DropDatabaseProcedure,
249            DropViewProcedure,
250            CommentOnProcedure
251        );
252
253        for (type_name, loader_factory) in loaders {
254            let context = self.create_context();
255            self.procedure_manager
256                .register_loader(type_name, loader_factory(context))
257                .context(RegisterProcedureLoaderSnafu { type_name })?;
258        }
259
260        self.repartition_procedure_factory
261            .register_loaders(&self.ddl_context, &self.procedure_manager)
262            .context(RegisterRepartitionProcedureLoaderSnafu)?;
263
264        Ok(())
265    }
266
267    /// Submits a repartition procedure for the specified table.
268    ///
269    /// This creates a repartition procedure using the provided `table_id`,
270    /// `table_name`, and `Repartition` configuration, and then either executes it
271    /// to completion or just submits it for asynchronous execution.
272    ///
273    /// The `Repartition` argument contains the original (`from_partition_exprs`)
274    /// and target (`into_partition_exprs`) partition expressions that define how
275    /// the table should be repartitioned.
276    ///
277    /// The `wait` flag controls whether this method waits for the repartition
278    /// procedure to finish:
279    /// - If `wait` is `true`, the procedure is executed and this method awaits
280    ///   its completion, returning both the generated `ProcedureId` and the
281    ///   final `Output` of the procedure.
282    /// - If `wait` is `false`, the procedure is only submitted to the procedure
283    ///   manager for asynchronous execution, and this method returns the
284    ///   `ProcedureId` along with `None` as the output.
285    async fn submit_repartition_task(
286        &self,
287        table_id: TableId,
288        table_name: TableName,
289        repartition: Repartition,
290        wait: bool,
291        timeout: Duration,
292    ) -> Result<(ProcedureId, Option<Output>)> {
293        let context = self.create_context();
294
295        let into_partition_exprs = repartition.into_partition_exprs;
296        let source = repartition.source;
297
298        let source = match source {
299            Some(PbRepartitionSource::PartitionExprs(PartitionExprs { exprs })) => {
300                RepartitionSource::Partitioned { exprs }
301            }
302            Some(PbRepartitionSource::Unpartitioned(source)) => RepartitionSource::Unpartitioned {
303                partition_columns: source.partition_columns,
304            },
305            None => {
306                // Reads the deprecated field for backward compatibility with old persisted DDL tasks.
307                #[allow(deprecated)]
308                RepartitionSource::Partitioned {
309                    exprs: repartition.from_partition_exprs,
310                }
311            }
312        };
313
314        let procedure = self
315            .repartition_procedure_factory
316            .create(
317                &context,
318                table_name,
319                table_id,
320                source,
321                into_partition_exprs,
322                Some(timeout),
323            )
324            .context(CreateRepartitionProcedureSnafu)?;
325        let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure));
326        if wait {
327            self.execute_procedure_and_wait(procedure_with_id).await
328        } else {
329            self.submit_procedure(procedure_with_id)
330                .await
331                .map(|p| (p, None))
332        }
333    }
334
335    /// Submits and executes an alter table task.
336    #[tracing::instrument(skip_all)]
337    pub async fn submit_alter_table_task(
338        &self,
339        table_id: TableId,
340        alter_table_task: AlterTableTask,
341        ddl_options: DdlOptions,
342    ) -> Result<(ProcedureId, Option<Output>)> {
343        // make alter_table_task mutable so we can call .take() on its field
344        let mut alter_table_task = alter_table_task;
345        if let Some(Kind::Repartition(_)) = alter_table_task.alter_table.kind.as_ref()
346            && let Kind::Repartition(repartition) =
347                alter_table_task.alter_table.kind.take().unwrap()
348        {
349            let table_name = TableName::new(
350                alter_table_task.alter_table.catalog_name,
351                alter_table_task.alter_table.schema_name,
352                alter_table_task.alter_table.table_name,
353            );
354            return self
355                .submit_repartition_task(
356                    table_id,
357                    table_name,
358                    repartition,
359                    ddl_options.wait,
360                    ddl_options.timeout,
361                )
362                .await;
363        }
364
365        let context = self.create_context();
366        let procedure = AlterTableProcedure::new(table_id, alter_table_task, context)?;
367
368        let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure));
369
370        self.execute_procedure_and_wait(procedure_with_id).await
371    }
372
373    /// Submits and executes a create table task.
374    #[tracing::instrument(skip_all)]
375    pub async fn submit_create_table_task(
376        &self,
377        create_table_task: CreateTableTask,
378        query_context: QueryContext,
379    ) -> Result<(ProcedureId, Option<Output>)> {
380        let context = self.create_context();
381
382        let procedure = CreateTableProcedure::new_with_query_context(
383            create_table_task,
384            query_context,
385            context,
386        )?;
387
388        let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure));
389
390        self.execute_procedure_and_wait(procedure_with_id).await
391    }
392
393    /// Submits and executes a `[CreateViewTask]`.
394    #[tracing::instrument(skip_all)]
395    pub async fn submit_create_view_task(
396        &self,
397        create_view_task: CreateViewTask,
398    ) -> Result<(ProcedureId, Option<Output>)> {
399        let context = self.create_context();
400
401        let procedure = CreateViewProcedure::new(create_view_task, context);
402
403        let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure));
404
405        self.execute_procedure_and_wait(procedure_with_id).await
406    }
407
408    /// Submits and executes a create multiple logical table tasks.
409    #[tracing::instrument(skip_all)]
410    pub async fn submit_create_logical_table_tasks(
411        &self,
412        create_table_tasks: Vec<CreateTableTask>,
413        physical_table_id: TableId,
414    ) -> Result<(ProcedureId, Option<Output>)> {
415        let context = self.create_context();
416
417        let procedure =
418            CreateLogicalTablesProcedure::new(create_table_tasks, physical_table_id, context);
419
420        let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure));
421
422        self.execute_procedure_and_wait(procedure_with_id).await
423    }
424
425    /// Submits and executes alter multiple table tasks.
426    #[tracing::instrument(skip_all)]
427    pub async fn submit_alter_logical_table_tasks(
428        &self,
429        alter_table_tasks: Vec<AlterTableTask>,
430        physical_table_id: TableId,
431    ) -> Result<(ProcedureId, Option<Output>)> {
432        let context = self.create_context();
433
434        let procedure =
435            AlterLogicalTablesProcedure::new(alter_table_tasks, physical_table_id, context);
436
437        let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure));
438
439        self.execute_procedure_and_wait(procedure_with_id).await
440    }
441
442    /// Submits and executes a drop table task.
443    #[tracing::instrument(skip_all)]
444    pub async fn submit_drop_table_task(
445        &self,
446        drop_table_task: DropTableTask,
447    ) -> Result<(ProcedureId, Option<Output>)> {
448        let context = self.create_context();
449
450        let procedure = DropTableProcedure::new(drop_table_task, context);
451
452        let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure));
453
454        self.execute_procedure_and_wait(procedure_with_id).await
455    }
456
457    /// Submits and executes a create database task.
458    #[tracing::instrument(skip_all)]
459    pub async fn submit_create_database(
460        &self,
461        CreateDatabaseTask {
462            catalog,
463            schema,
464            create_if_not_exists,
465            options,
466        }: CreateDatabaseTask,
467    ) -> Result<(ProcedureId, Option<Output>)> {
468        let context = self.create_context();
469        let procedure =
470            CreateDatabaseProcedure::new(catalog, schema, create_if_not_exists, options, context);
471        let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure));
472
473        self.execute_procedure_and_wait(procedure_with_id).await
474    }
475
476    /// Submits and executes a drop table task.
477    #[tracing::instrument(skip_all)]
478    pub async fn submit_drop_database(
479        &self,
480        DropDatabaseTask {
481            catalog,
482            schema,
483            drop_if_exists,
484        }: DropDatabaseTask,
485    ) -> Result<(ProcedureId, Option<Output>)> {
486        let context = self.create_context();
487        let procedure = DropDatabaseProcedure::new(catalog, schema, drop_if_exists, context);
488        let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure));
489
490        self.execute_procedure_and_wait(procedure_with_id).await
491    }
492
493    pub async fn submit_alter_database(
494        &self,
495        alter_database_task: AlterDatabaseTask,
496    ) -> Result<(ProcedureId, Option<Output>)> {
497        let context = self.create_context();
498        let procedure = AlterDatabaseProcedure::new(alter_database_task, context)?;
499        let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure));
500
501        self.execute_procedure_and_wait(procedure_with_id).await
502    }
503
504    /// Submits and executes a create flow task.
505    #[tracing::instrument(skip_all)]
506    pub async fn submit_create_flow_task(
507        &self,
508        create_flow: CreateFlowTask,
509        query_context: QueryContext,
510    ) -> Result<(ProcedureId, Option<Output>)> {
511        let context = self.create_context();
512        let procedure = CreateFlowProcedure::new(create_flow, query_context, context);
513        let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure));
514
515        self.execute_procedure_and_wait(procedure_with_id).await
516    }
517
518    /// Submits and executes a drop flow task.
519    #[tracing::instrument(skip_all)]
520    pub async fn submit_drop_flow_task(
521        &self,
522        drop_flow: DropFlowTask,
523    ) -> Result<(ProcedureId, Option<Output>)> {
524        let context = self.create_context();
525        let procedure = DropFlowProcedure::new(drop_flow, context);
526        let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure));
527
528        self.execute_procedure_and_wait(procedure_with_id).await
529    }
530
531    /// Submits and executes a drop view task.
532    #[tracing::instrument(skip_all)]
533    pub async fn submit_drop_view_task(
534        &self,
535        drop_view: DropViewTask,
536    ) -> Result<(ProcedureId, Option<Output>)> {
537        let context = self.create_context();
538        let procedure = DropViewProcedure::new(drop_view, context);
539        let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure));
540
541        self.execute_procedure_and_wait(procedure_with_id).await
542    }
543
544    /// Submits and executes a truncate table task.
545    #[tracing::instrument(skip_all)]
546    pub async fn submit_truncate_table_task(
547        &self,
548        truncate_table_task: TruncateTableTask,
549        table_info_value: DeserializedValueWithBytes<TableInfoValue>,
550    ) -> Result<(ProcedureId, Option<Output>)> {
551        let context = self.create_context();
552        let procedure = TruncateTableProcedure::new(truncate_table_task, table_info_value, context);
553
554        let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure));
555
556        self.execute_procedure_and_wait(procedure_with_id).await
557    }
558
559    /// Submits and executes a comment on task.
560    #[tracing::instrument(skip_all)]
561    pub async fn submit_comment_on_task(
562        &self,
563        comment_on_task: CommentOnTask,
564    ) -> Result<(ProcedureId, Option<Output>)> {
565        let context = self.create_context();
566        let procedure = CommentOnProcedure::new(comment_on_task, context);
567        let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure));
568
569        self.execute_procedure_and_wait(procedure_with_id).await
570    }
571
572    /// Executes a procedure and waits for the result.
573    async fn execute_procedure_and_wait(
574        &self,
575        procedure_with_id: ProcedureWithId,
576    ) -> Result<(ProcedureId, Option<Output>)> {
577        let procedure_id = procedure_with_id.id;
578
579        let mut watcher = self
580            .procedure_manager
581            .submit(procedure_with_id)
582            .await
583            .context(SubmitProcedureSnafu)?;
584
585        let output = watcher::wait(&mut watcher)
586            .await
587            .context(WaitProcedureSnafu)?;
588
589        Ok((procedure_id, output))
590    }
591
592    /// Submits a procedure and returns the procedure id.
593    async fn submit_procedure(&self, procedure_with_id: ProcedureWithId) -> Result<ProcedureId> {
594        let procedure_id = procedure_with_id.id;
595        let _ = self
596            .procedure_manager
597            .submit(procedure_with_id)
598            .await
599            .context(SubmitProcedureSnafu)?;
600
601        Ok(procedure_id)
602    }
603
604    pub async fn submit_ddl_task(
605        &self,
606        ctx: &ExecutorContext,
607        request: SubmitDdlTaskRequest,
608    ) -> Result<SubmitDdlTaskResponse> {
609        let span = ctx
610            .tracing_context
611            .as_ref()
612            .map(TracingContext::from_w3c)
613            .unwrap_or_else(TracingContext::from_current_span)
614            .attach(tracing::info_span!("DdlManager::submit_ddl_task"));
615        let ddl_options = DdlOptions {
616            wait: request.wait,
617            timeout: request.timeout,
618        };
619        async move {
620            debug!("Submitting Ddl task: {:?}", request.task);
621            match request.task {
622                CreateTable(create_table_task) => {
623                    handle_create_table_task(self, create_table_task, request.query_context).await
624                }
625                DropTable(drop_table_task) => handle_drop_table_task(self, drop_table_task).await,
626                AlterTable(alter_table_task) => {
627                    handle_alter_table_task(self, alter_table_task, ddl_options).await
628                }
629                TruncateTable(truncate_table_task) => {
630                    handle_truncate_table_task(self, truncate_table_task).await
631                }
632                CreateLogicalTables(create_table_tasks) => {
633                    handle_create_logical_table_tasks(self, create_table_tasks).await
634                }
635                AlterLogicalTables(alter_table_tasks) => {
636                    handle_alter_logical_table_tasks(self, alter_table_tasks).await
637                }
638                DropLogicalTables(_) => todo!(),
639                CreateDatabase(create_database_task) => {
640                    handle_create_database_task(self, create_database_task).await
641                }
642                DropDatabase(drop_database_task) => {
643                    handle_drop_database_task(self, drop_database_task).await
644                }
645                AlterDatabase(alter_database_task) => {
646                    handle_alter_database_task(self, alter_database_task).await
647                }
648                CreateFlow(create_flow_task) => {
649                    handle_create_flow_task(self, create_flow_task, request.query_context).await
650                }
651                DropFlow(drop_flow_task) => handle_drop_flow_task(self, drop_flow_task).await,
652                CreateView(create_view_task) => {
653                    handle_create_view_task(self, create_view_task).await
654                }
655                DropView(drop_view_task) => handle_drop_view_task(self, drop_view_task).await,
656                CommentOn(comment_on_task) => handle_comment_on_task(self, comment_on_task).await,
657                #[cfg(feature = "enterprise")]
658                CreateTrigger(create_trigger_task) => {
659                    handle_create_trigger_task(self, create_trigger_task, request.query_context)
660                        .await
661                }
662                #[cfg(feature = "enterprise")]
663                DropTrigger(drop_trigger_task) => {
664                    handle_drop_trigger_task(self, drop_trigger_task, request.query_context).await
665                }
666            }
667        }
668        .trace(span)
669        .await
670    }
671}
672
673async fn handle_truncate_table_task(
674    ddl_manager: &DdlManager,
675    truncate_table_task: TruncateTableTask,
676) -> Result<SubmitDdlTaskResponse> {
677    let table_id = truncate_table_task.table_id;
678    let table_metadata_manager = &ddl_manager.table_metadata_manager();
679    let table_ref = truncate_table_task.table_ref();
680
681    let table_info_value = table_metadata_manager
682        .table_info_manager()
683        .get(table_id)
684        .await?
685        .with_context(|| TableInfoNotFoundSnafu {
686            table: table_ref.to_string(),
687        })?;
688    let physical_table_id = table_metadata_manager
689        .table_route_manager()
690        .get_physical_table_id(table_id)
691        .await?;
692    ensure!(
693        physical_table_id == table_id,
694        error::UnexpectedSnafu {
695            err_msg: "Truncate table is only supported for physical tables."
696        }
697    );
698
699    let (id, _) = ddl_manager
700        .submit_truncate_table_task(truncate_table_task, table_info_value)
701        .await?;
702
703    info!("Table: {table_id} is truncated via procedure_id {id:?}");
704
705    Ok(SubmitDdlTaskResponse {
706        key: id.to_string().into(),
707        ..Default::default()
708    })
709}
710
711async fn handle_alter_table_task(
712    ddl_manager: &DdlManager,
713    alter_table_task: AlterTableTask,
714    ddl_options: DdlOptions,
715) -> Result<SubmitDdlTaskResponse> {
716    let table_ref = alter_table_task.table_ref();
717
718    let table_id = ddl_manager
719        .table_metadata_manager()
720        .table_name_manager()
721        .get(TableNameKey::new(
722            table_ref.catalog,
723            table_ref.schema,
724            table_ref.table,
725        ))
726        .await?
727        .with_context(|| TableNotFoundSnafu {
728            table_name: table_ref.to_string(),
729        })?
730        .table_id();
731
732    let table_route_value = ddl_manager
733        .table_metadata_manager()
734        .table_route_manager()
735        .table_route_storage()
736        .get(table_id)
737        .await?
738        .context(TableRouteNotFoundSnafu { table_id })?;
739    ensure!(
740        table_route_value.is_physical(),
741        UnexpectedLogicalRouteTableSnafu {
742            err_msg: format!("{:?} is a non-physical TableRouteValue.", table_ref),
743        }
744    );
745
746    let (id, _) = ddl_manager
747        .submit_alter_table_task(table_id, alter_table_task, ddl_options)
748        .await?;
749
750    info!("Table: {table_id} is altered via procedure_id {id:?}");
751
752    Ok(SubmitDdlTaskResponse {
753        key: id.to_string().into(),
754        ..Default::default()
755    })
756}
757
758async fn handle_drop_table_task(
759    ddl_manager: &DdlManager,
760    drop_table_task: DropTableTask,
761) -> Result<SubmitDdlTaskResponse> {
762    let table_id = drop_table_task.table_id;
763    let (id, _) = ddl_manager.submit_drop_table_task(drop_table_task).await?;
764
765    info!("Table: {table_id} is dropped via procedure_id {id:?}");
766
767    Ok(SubmitDdlTaskResponse {
768        key: id.to_string().into(),
769        ..Default::default()
770    })
771}
772
773async fn handle_create_table_task(
774    ddl_manager: &DdlManager,
775    create_table_task: CreateTableTask,
776    query_context: QueryContext,
777) -> Result<SubmitDdlTaskResponse> {
778    let (id, output) = ddl_manager
779        .submit_create_table_task(create_table_task, query_context)
780        .await?;
781
782    let procedure_id = id.to_string();
783    let output = output.context(ProcedureOutputSnafu {
784        procedure_id: &procedure_id,
785        err_msg: "empty output",
786    })?;
787    let table_id = *(output.downcast_ref::<u32>().context(ProcedureOutputSnafu {
788        procedure_id: &procedure_id,
789        err_msg: "downcast to `u32`",
790    })?);
791    info!("Table: {table_id} is created via procedure_id {id:?}");
792
793    Ok(SubmitDdlTaskResponse {
794        key: procedure_id.into(),
795        table_ids: vec![table_id],
796    })
797}
798
799async fn handle_create_logical_table_tasks(
800    ddl_manager: &DdlManager,
801    create_table_tasks: Vec<CreateTableTask>,
802) -> Result<SubmitDdlTaskResponse> {
803    ensure!(
804        !create_table_tasks.is_empty(),
805        EmptyDdlTasksSnafu {
806            name: "create logical tables"
807        }
808    );
809    let physical_table_id = utils::check_and_get_physical_table_id(
810        ddl_manager.table_metadata_manager(),
811        &create_table_tasks,
812    )
813    .await?;
814    let num_logical_tables = create_table_tasks.len();
815
816    let (id, output) = ddl_manager
817        .submit_create_logical_table_tasks(create_table_tasks, physical_table_id)
818        .await?;
819
820    info!(
821        "{num_logical_tables} logical tables on physical table: {physical_table_id:?} is created via procedure_id {id:?}"
822    );
823
824    let procedure_id = id.to_string();
825    let output = output.context(ProcedureOutputSnafu {
826        procedure_id: &procedure_id,
827        err_msg: "empty output",
828    })?;
829    let table_ids = output
830        .downcast_ref::<Vec<TableId>>()
831        .context(ProcedureOutputSnafu {
832            procedure_id: &procedure_id,
833            err_msg: "downcast to `Vec<TableId>`",
834        })?
835        .clone();
836
837    Ok(SubmitDdlTaskResponse {
838        key: procedure_id.into(),
839        table_ids,
840    })
841}
842
843async fn handle_create_database_task(
844    ddl_manager: &DdlManager,
845    create_database_task: CreateDatabaseTask,
846) -> Result<SubmitDdlTaskResponse> {
847    let (id, _) = ddl_manager
848        .submit_create_database(create_database_task.clone())
849        .await?;
850
851    let procedure_id = id.to_string();
852    info!(
853        "Database {}.{} is created via procedure_id {id:?}",
854        create_database_task.catalog, create_database_task.schema
855    );
856
857    Ok(SubmitDdlTaskResponse {
858        key: procedure_id.into(),
859        ..Default::default()
860    })
861}
862
863async fn handle_drop_database_task(
864    ddl_manager: &DdlManager,
865    drop_database_task: DropDatabaseTask,
866) -> Result<SubmitDdlTaskResponse> {
867    let (id, _) = ddl_manager
868        .submit_drop_database(drop_database_task.clone())
869        .await?;
870
871    let procedure_id = id.to_string();
872    info!(
873        "Database {}.{} is dropped via procedure_id {id:?}",
874        drop_database_task.catalog, drop_database_task.schema
875    );
876
877    Ok(SubmitDdlTaskResponse {
878        key: procedure_id.into(),
879        ..Default::default()
880    })
881}
882
883async fn handle_alter_database_task(
884    ddl_manager: &DdlManager,
885    alter_database_task: AlterDatabaseTask,
886) -> Result<SubmitDdlTaskResponse> {
887    let (id, _) = ddl_manager
888        .submit_alter_database(alter_database_task.clone())
889        .await?;
890
891    let procedure_id = id.to_string();
892    info!(
893        "Database {}.{} is altered via procedure_id {id:?}",
894        alter_database_task.catalog(),
895        alter_database_task.schema()
896    );
897
898    Ok(SubmitDdlTaskResponse {
899        key: procedure_id.into(),
900        ..Default::default()
901    })
902}
903
904async fn handle_drop_flow_task(
905    ddl_manager: &DdlManager,
906    drop_flow_task: DropFlowTask,
907) -> Result<SubmitDdlTaskResponse> {
908    let (id, _) = ddl_manager
909        .submit_drop_flow_task(drop_flow_task.clone())
910        .await?;
911
912    let procedure_id = id.to_string();
913    info!(
914        "Flow {}.{}({}) is dropped via procedure_id {id:?}",
915        drop_flow_task.catalog_name, drop_flow_task.flow_name, drop_flow_task.flow_id,
916    );
917
918    Ok(SubmitDdlTaskResponse {
919        key: procedure_id.into(),
920        ..Default::default()
921    })
922}
923
924#[cfg(feature = "enterprise")]
925async fn handle_drop_trigger_task(
926    ddl_manager: &DdlManager,
927    drop_trigger_task: DropTriggerTask,
928    query_context: QueryContext,
929) -> Result<SubmitDdlTaskResponse> {
930    let Some(m) = ddl_manager.trigger_ddl_manager.as_ref() else {
931        use crate::error::UnsupportedSnafu;
932
933        return UnsupportedSnafu {
934            operation: "drop trigger",
935        }
936        .fail();
937    };
938
939    m.drop_trigger(
940        drop_trigger_task,
941        ddl_manager.procedure_manager.clone(),
942        ddl_manager.ddl_context.clone(),
943        query_context,
944    )
945    .await
946}
947
948async fn handle_drop_view_task(
949    ddl_manager: &DdlManager,
950    drop_view_task: DropViewTask,
951) -> Result<SubmitDdlTaskResponse> {
952    let (id, _) = ddl_manager
953        .submit_drop_view_task(drop_view_task.clone())
954        .await?;
955
956    let procedure_id = id.to_string();
957    info!(
958        "View {}({}) is dropped via procedure_id {id:?}",
959        drop_view_task.table_ref(),
960        drop_view_task.view_id,
961    );
962
963    Ok(SubmitDdlTaskResponse {
964        key: procedure_id.into(),
965        ..Default::default()
966    })
967}
968
969async fn handle_create_flow_task(
970    ddl_manager: &DdlManager,
971    create_flow_task: CreateFlowTask,
972    query_context: QueryContext,
973) -> Result<SubmitDdlTaskResponse> {
974    let (id, output) = ddl_manager
975        .submit_create_flow_task(create_flow_task.clone(), query_context)
976        .await?;
977
978    let procedure_id = id.to_string();
979    let output = output.context(ProcedureOutputSnafu {
980        procedure_id: &procedure_id,
981        err_msg: "empty output",
982    })?;
983    let flow_id = *(output.downcast_ref::<u32>().context(ProcedureOutputSnafu {
984        procedure_id: &procedure_id,
985        err_msg: "downcast to `u32`",
986    })?);
987    if !create_flow_task.or_replace {
988        info!(
989            "Flow {}.{}({flow_id}) is created via procedure_id {id:?}",
990            create_flow_task.catalog_name, create_flow_task.flow_name,
991        );
992    } else {
993        info!(
994            "Flow {}.{}({flow_id}) is replaced via procedure_id {id:?}",
995            create_flow_task.catalog_name, create_flow_task.flow_name,
996        );
997    }
998
999    Ok(SubmitDdlTaskResponse {
1000        key: procedure_id.into(),
1001        ..Default::default()
1002    })
1003}
1004
1005#[cfg(feature = "enterprise")]
1006async fn handle_create_trigger_task(
1007    ddl_manager: &DdlManager,
1008    create_trigger_task: CreateTriggerTask,
1009    query_context: QueryContext,
1010) -> Result<SubmitDdlTaskResponse> {
1011    let Some(m) = ddl_manager.trigger_ddl_manager.as_ref() else {
1012        use crate::error::UnsupportedSnafu;
1013
1014        return UnsupportedSnafu {
1015            operation: "create trigger",
1016        }
1017        .fail();
1018    };
1019
1020    m.create_trigger(
1021        create_trigger_task,
1022        ddl_manager.procedure_manager.clone(),
1023        ddl_manager.ddl_context.clone(),
1024        query_context,
1025    )
1026    .await
1027}
1028
1029async fn handle_alter_logical_table_tasks(
1030    ddl_manager: &DdlManager,
1031    alter_table_tasks: Vec<AlterTableTask>,
1032) -> Result<SubmitDdlTaskResponse> {
1033    ensure!(
1034        !alter_table_tasks.is_empty(),
1035        EmptyDdlTasksSnafu {
1036            name: "alter logical tables"
1037        }
1038    );
1039
1040    // Use the physical table id in the first logical table, then it will be checked in the procedure.
1041    let first_table = TableNameKey {
1042        catalog: &alter_table_tasks[0].alter_table.catalog_name,
1043        schema: &alter_table_tasks[0].alter_table.schema_name,
1044        table: &alter_table_tasks[0].alter_table.table_name,
1045    };
1046    let physical_table_id =
1047        utils::get_physical_table_id(ddl_manager.table_metadata_manager(), first_table).await?;
1048    let num_logical_tables = alter_table_tasks.len();
1049
1050    let (id, _) = ddl_manager
1051        .submit_alter_logical_table_tasks(alter_table_tasks, physical_table_id)
1052        .await?;
1053
1054    info!(
1055        "{num_logical_tables} logical tables on physical table: {physical_table_id:?} is altered via procedure_id {id:?}"
1056    );
1057
1058    let procedure_id = id.to_string();
1059
1060    Ok(SubmitDdlTaskResponse {
1061        key: procedure_id.into(),
1062        ..Default::default()
1063    })
1064}
1065
1066/// Handle the `[CreateViewTask]` and returns the DDL response when success.
1067async fn handle_create_view_task(
1068    ddl_manager: &DdlManager,
1069    create_view_task: CreateViewTask,
1070) -> Result<SubmitDdlTaskResponse> {
1071    let (id, output) = ddl_manager
1072        .submit_create_view_task(create_view_task)
1073        .await?;
1074
1075    let procedure_id = id.to_string();
1076    let output = output.context(ProcedureOutputSnafu {
1077        procedure_id: &procedure_id,
1078        err_msg: "empty output",
1079    })?;
1080    let view_id = *(output.downcast_ref::<u32>().context(ProcedureOutputSnafu {
1081        procedure_id: &procedure_id,
1082        err_msg: "downcast to `u32`",
1083    })?);
1084    info!("View: {view_id} is created via procedure_id {id:?}");
1085
1086    Ok(SubmitDdlTaskResponse {
1087        key: procedure_id.into(),
1088        table_ids: vec![view_id],
1089    })
1090}
1091
1092async fn handle_comment_on_task(
1093    ddl_manager: &DdlManager,
1094    comment_on_task: CommentOnTask,
1095) -> Result<SubmitDdlTaskResponse> {
1096    let (id, _) = ddl_manager
1097        .submit_comment_on_task(comment_on_task.clone())
1098        .await?;
1099
1100    let procedure_id = id.to_string();
1101    info!(
1102        "Comment on {}.{}.{} is updated via procedure_id {id:?}",
1103        comment_on_task.catalog_name, comment_on_task.schema_name, comment_on_task.object_name
1104    );
1105
1106    Ok(SubmitDdlTaskResponse {
1107        key: procedure_id.into(),
1108        ..Default::default()
1109    })
1110}
1111
1112#[cfg(test)]
1113mod tests {
1114    use std::sync::Arc;
1115    use std::time::Duration;
1116
1117    use common_error::ext::BoxedError;
1118    use common_procedure::local::LocalManager;
1119    use common_procedure::test_util::InMemoryPoisonStore;
1120    use common_procedure::{BoxedProcedure, ProcedureManagerRef};
1121    use store_api::storage::TableId;
1122    use table::table_name::TableName;
1123
1124    use super::DdlManager;
1125    use crate::cache_invalidator::DummyCacheInvalidator;
1126    use crate::ddl::alter_table::AlterTableProcedure;
1127    use crate::ddl::create_table::CreateTableProcedure;
1128    use crate::ddl::drop_table::DropTableProcedure;
1129    use crate::ddl::flow_meta::FlowMetadataAllocator;
1130    use crate::ddl::table_meta::TableMetadataAllocator;
1131    use crate::ddl::truncate_table::TruncateTableProcedure;
1132    use crate::ddl::{DdlContext, NoopRegionFailureDetectorControl};
1133    use crate::ddl_manager::{RepartitionProcedureFactory, RepartitionSource};
1134    use crate::key::TableMetadataManager;
1135    use crate::key::flow::FlowMetadataManager;
1136    use crate::kv_backend::memory::MemoryKvBackend;
1137    use crate::node_manager::{DatanodeManager, DatanodeRef, FlownodeManager, FlownodeRef};
1138    use crate::peer::Peer;
1139    use crate::region_keeper::MemoryRegionKeeper;
1140    use crate::region_registry::LeaderRegionRegistry;
1141    use crate::sequence::SequenceBuilder;
1142    use crate::state_store::KvStateStore;
1143    use crate::wal_provider::WalProvider;
1144
1145    /// A dummy implemented [NodeManager].
1146    pub struct DummyDatanodeManager;
1147
1148    #[async_trait::async_trait]
1149    impl DatanodeManager for DummyDatanodeManager {
1150        async fn datanode(&self, _datanode: &Peer) -> DatanodeRef {
1151            unimplemented!()
1152        }
1153    }
1154
1155    #[async_trait::async_trait]
1156    impl FlownodeManager for DummyDatanodeManager {
1157        async fn flownode(&self, _node: &Peer) -> FlownodeRef {
1158            unimplemented!()
1159        }
1160    }
1161
1162    struct DummyRepartitionProcedureFactory;
1163
1164    #[async_trait::async_trait]
1165    impl RepartitionProcedureFactory for DummyRepartitionProcedureFactory {
1166        fn create(
1167            &self,
1168            _ddl_ctx: &DdlContext,
1169            _table_name: TableName,
1170            _table_id: TableId,
1171            _source: RepartitionSource,
1172            _to_exprs: Vec<String>,
1173            _timeout: Option<Duration>,
1174        ) -> std::result::Result<BoxedProcedure, BoxedError> {
1175            unimplemented!()
1176        }
1177
1178        fn register_loaders(
1179            &self,
1180            _ddl_ctx: &DdlContext,
1181            _procedure_manager: &ProcedureManagerRef,
1182        ) -> std::result::Result<(), BoxedError> {
1183            Ok(())
1184        }
1185    }
1186
1187    #[test]
1188    fn test_try_new() {
1189        let kv_backend = Arc::new(MemoryKvBackend::new());
1190        let table_metadata_manager = Arc::new(TableMetadataManager::new(kv_backend.clone()));
1191        let table_metadata_allocator = Arc::new(TableMetadataAllocator::new(
1192            Arc::new(SequenceBuilder::new("test", kv_backend.clone()).build()),
1193            Arc::new(WalProvider::default()),
1194        ));
1195        let flow_metadata_manager = Arc::new(FlowMetadataManager::new(kv_backend.clone()));
1196        let flow_metadata_allocator = Arc::new(FlowMetadataAllocator::with_noop_peer_allocator(
1197            Arc::new(SequenceBuilder::new("flow-test", kv_backend.clone()).build()),
1198        ));
1199
1200        let state_store = Arc::new(KvStateStore::new(kv_backend.clone()));
1201        let poison_manager = Arc::new(InMemoryPoisonStore::default());
1202        let procedure_manager = Arc::new(LocalManager::new(
1203            Default::default(),
1204            state_store,
1205            poison_manager,
1206            None,
1207            None,
1208        ));
1209
1210        let _ = DdlManager::try_new(
1211            DdlContext {
1212                node_manager: Arc::new(DummyDatanodeManager),
1213                cache_invalidator: Arc::new(DummyCacheInvalidator),
1214                table_metadata_manager,
1215                table_metadata_allocator,
1216                flow_metadata_manager,
1217                flow_metadata_allocator,
1218                memory_region_keeper: Arc::new(MemoryRegionKeeper::default()),
1219                leader_region_registry: Arc::new(LeaderRegionRegistry::default()),
1220                region_failure_detector_controller: Arc::new(NoopRegionFailureDetectorControl),
1221            },
1222            procedure_manager.clone(),
1223            Arc::new(DummyRepartitionProcedureFactory),
1224            true,
1225        );
1226
1227        let expected_loaders = vec![
1228            CreateTableProcedure::TYPE_NAME,
1229            AlterTableProcedure::TYPE_NAME,
1230            DropTableProcedure::TYPE_NAME,
1231            TruncateTableProcedure::TYPE_NAME,
1232        ];
1233
1234        for loader in expected_loaders {
1235            assert!(procedure_manager.contains_loader(loader));
1236        }
1237    }
1238}