feat: Drop table by procedure (#1401)

* feat: Add drop table procedure

* feat: support dropping table by procedure on datanode

* test: Add test for DropTableProcedure

* test: Test drop table by procedure

* chore: update comments

* fix: Make on_remove_from_catalog idempotent

src/datanode/src/sql/alter.rs

@@ -221,7 +221,7 @@ mod tests {
     }
 
     #[tokio::test(flavor = "multi_thread")]
-    async fn alter_table_by_procedure() {
+    async fn test_alter_table_by_procedure() {
         let instance = MockInstance::with_procedure_enabled("alter_table_by_procedure").await;
 
         // Create table first.

src/datanode/src/sql/drop_table.rs

@@ -15,17 +15,23 @@
 use catalog::error::TableNotExistSnafu;
 use catalog::DeregisterTableRequest;
 use common_error::prelude::BoxedError;
+use common_procedure::{watcher, ProcedureManagerRef, ProcedureWithId};
 use common_query::Output;
 use common_telemetry::info;
 use snafu::{OptionExt, ResultExt};
 use table::engine::{EngineContext, TableReference};
 use table::requests::DropTableRequest;
+use table_procedure::DropTableProcedure;
 
 use crate::error::{self, Result};
 use crate::sql::SqlHandler;
 
 impl SqlHandler {
     pub async fn drop_table(&self, req: DropTableRequest) -> Result<Output> {
+        if let Some(procedure_manager) = &self.procedure_manager {
+            return self.drop_table_by_procedure(procedure_manager, req).await;
+        }
+
         let deregister_table_req = DeregisterTableRequest {
             catalog: req.catalog_name.clone(),
             schema: req.schema_name.clone(),
@@ -76,4 +82,80 @@ impl SqlHandler {
 
         Ok(Output::AffectedRows(1))
     }
+
+    pub(crate) async fn drop_table_by_procedure(
+        &self,
+        procedure_manager: &ProcedureManagerRef,
+        req: DropTableRequest,
+    ) -> Result<Output> {
+        let table_name = req.table_name.clone();
+        let procedure = DropTableProcedure::new(
+            req,
+            self.catalog_manager.clone(),
+            self.engine_procedure.clone(),
+        );
+
+        let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure));
+        let procedure_id = procedure_with_id.id;
+
+        info!("Drop table {} by procedure {}", table_name, procedure_id);
+
+        let mut watcher = procedure_manager
+            .submit(procedure_with_id)
+            .await
+            .context(error::SubmitProcedureSnafu { procedure_id })?;
+
+        watcher::wait(&mut watcher)
+            .await
+            .context(error::WaitProcedureSnafu { procedure_id })?;
+
+        Ok(Output::AffectedRows(1))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use query::parser::{QueryLanguageParser, QueryStatement};
+    use query::query_engine::SqlStatementExecutor;
+    use session::context::QueryContext;
+
+    use super::*;
+    use crate::tests::test_util::MockInstance;
+
+    #[tokio::test(flavor = "multi_thread")]
+    async fn test_drop_table_by_procedure() {
+        let instance = MockInstance::with_procedure_enabled("alter_table_by_procedure").await;
+
+        // Create table first.
+        let sql = r#"create table test_drop(
+                            host string,
+                            ts timestamp,
+                            cpu double default 0,
+                            TIME INDEX (ts),
+                            PRIMARY KEY(host)
+                        ) engine=mito with(regions=1);"#;
+        let stmt = match QueryLanguageParser::parse_sql(sql).unwrap() {
+            QueryStatement::Sql(sql) => sql,
+            _ => unreachable!(),
+        };
+        let output = instance
+            .inner()
+            .execute_sql(stmt, QueryContext::arc())
+            .await
+            .unwrap();
+        assert!(matches!(output, Output::AffectedRows(0)));
+
+        // Drop table.
+        let sql = r#"drop table test_drop"#;
+        let stmt = match QueryLanguageParser::parse_sql(sql).unwrap() {
+            QueryStatement::Sql(sql) => sql,
+            _ => unreachable!(),
+        };
+        let output = instance
+            .inner()
+            .execute_sql(stmt, QueryContext::arc())
+            .await
+            .unwrap();
+        assert!(matches!(output, Output::AffectedRows(1)));
+    }
 }

src/table-procedure/src/alter.rs

@@ -32,6 +32,7 @@ use crate::error::{
     SerializeProcedureSnafu, SubprocedureFailedSnafu, TableExistsSnafu, TableNotFoundSnafu,
 };
 
+/// Procedure to alter a table.
 pub struct AlterTableProcedure {
     data: AlterTableData,
     catalog_manager: CatalogManagerRef,
@@ -88,8 +89,8 @@ impl AlterTableProcedure {
             data: AlterTableData {
                 state: AlterTableState::Prepare,
                 request,
-                subprocedure_id: None,
                 table_id: None,
+                subprocedure_id: None,
             },
             catalog_manager,
             engine_procedure,
@@ -137,13 +138,13 @@ impl AlterTableProcedure {
             .catalog_manager
             .catalog(&self.data.request.catalog_name)
             .context(AccessCatalogSnafu)?
-            .with_context(|| CatalogNotFoundSnafu {
+            .context(CatalogNotFoundSnafu {
                 name: &self.data.request.catalog_name,
             })?;
         let schema = catalog
             .schema(&self.data.request.schema_name)
             .context(AccessCatalogSnafu)?
-            .with_context(|| SchemaNotFoundSnafu {
+            .context(SchemaNotFoundSnafu {
                 name: &self.data.request.schema_name,
             })?;
 
@@ -276,7 +277,7 @@ struct AlterTableData {
     ///
     /// Available after [AlterTableState::Prepare] state.
     table_id: Option<TableId>,
-    /// Id of the subprocedure to create this table in the engine.
+    /// Id of the subprocedure to alter this table in the engine.
     ///
     /// This id is `Some` while the procedure is in [AlterTableState::EngineAlterTable]
     /// state.

src/table-procedure/src/drop.rs

@@ -0,0 +1,309 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Procedure to drop a table.
+
+use async_trait::async_trait;
+use catalog::{CatalogManagerRef, DeregisterTableRequest};
+use common_procedure::{
+    Context, Error, LockKey, Procedure, ProcedureId, ProcedureManager, ProcedureState,
+    ProcedureWithId, Result, Status,
+};
+use common_telemetry::logging;
+use serde::{Deserialize, Serialize};
+use snafu::{OptionExt, ResultExt};
+use table::engine::{EngineContext, TableEngineProcedureRef, TableReference};
+use table::requests::DropTableRequest;
+
+use crate::error::{
+    AccessCatalogSnafu, DeserializeProcedureSnafu, SerializeProcedureSnafu,
+    SubprocedureFailedSnafu, TableNotFoundSnafu,
+};
+
+/// Procedure to drop a table.
+pub struct DropTableProcedure {
+    data: DropTableData,
+    catalog_manager: CatalogManagerRef,
+    engine_procedure: TableEngineProcedureRef,
+}
+
+#[async_trait]
+impl Procedure for DropTableProcedure {
+    fn type_name(&self) -> &str {
+        Self::TYPE_NAME
+    }
+
+    async fn execute(&mut self, ctx: &Context) -> Result<Status> {
+        match self.data.state {
+            DropTableState::Prepare => self.on_prepare().await,
+            DropTableState::RemoveFromCatalog => self.on_remove_from_catalog().await,
+            DropTableState::EngineDropTable => self.on_engine_drop_table(ctx).await,
+        }
+    }
+
+    fn dump(&self) -> Result<String> {
+        let json = serde_json::to_string(&self.data).context(SerializeProcedureSnafu)?;
+        Ok(json)
+    }
+
+    fn lock_key(&self) -> LockKey {
+        // We lock the whole table.
+        let table_name = self.data.table_ref().to_string();
+        LockKey::single(table_name)
+    }
+}
+
+impl DropTableProcedure {
+    const TYPE_NAME: &str = "table-procedure::DropTableProcedure";
+
+    /// Returns a new [DropTableProcedure].
+    pub fn new(
+        request: DropTableRequest,
+        catalog_manager: CatalogManagerRef,
+        engine_procedure: TableEngineProcedureRef,
+    ) -> DropTableProcedure {
+        DropTableProcedure {
+            data: DropTableData {
+                state: DropTableState::Prepare,
+                request,
+                subprocedure_id: None,
+            },
+            catalog_manager,
+            engine_procedure,
+        }
+    }
+
+    /// Register the loader of this procedure to the `procedure_manager`.
+    ///
+    /// # Panics
+    /// Panics on error.
+    pub fn register_loader(
+        catalog_manager: CatalogManagerRef,
+        engine_procedure: TableEngineProcedureRef,
+        procedure_manager: &dyn ProcedureManager,
+    ) {
+        procedure_manager
+            .register_loader(
+                Self::TYPE_NAME,
+                Box::new(move |data| {
+                    Self::from_json(data, catalog_manager.clone(), engine_procedure.clone())
+                        .map(|p| Box::new(p) as _)
+                }),
+            )
+            .unwrap()
+    }
+
+    /// Recover the procedure from json.
+    fn from_json(
+        json: &str,
+        catalog_manager: CatalogManagerRef,
+        engine_procedure: TableEngineProcedureRef,
+    ) -> Result<Self> {
+        let data: DropTableData = serde_json::from_str(json).context(DeserializeProcedureSnafu)?;
+
+        Ok(DropTableProcedure {
+            data,
+            catalog_manager,
+            engine_procedure,
+        })
+    }
+
+    async fn on_prepare(&mut self) -> Result<Status> {
+        let request = &self.data.request;
+        // Ensure the table exists.
+        self.catalog_manager
+            .table(
+                &request.catalog_name,
+                &request.schema_name,
+                &request.table_name,
+            )
+            .await
+            .context(AccessCatalogSnafu)?
+            .context(TableNotFoundSnafu {
+                name: &request.table_name,
+            })?;
+
+        self.data.state = DropTableState::RemoveFromCatalog;
+
+        Ok(Status::executing(true))
+    }
+
+    async fn on_remove_from_catalog(&mut self) -> Result<Status> {
+        let request = &self.data.request;
+        let has_table = self
+            .catalog_manager
+            .table(
+                &request.catalog_name,
+                &request.schema_name,
+                &request.table_name,
+            )
+            .await
+            .context(AccessCatalogSnafu)?
+            .is_some();
+        if has_table {
+            // The table is still in the catalog.
+            let deregister_table_req = DeregisterTableRequest {
+                catalog: self.data.request.catalog_name.clone(),
+                schema: self.data.request.schema_name.clone(),
+                table_name: self.data.request.table_name.clone(),
+            };
+            self.catalog_manager
+                .deregister_table(deregister_table_req)
+                .await
+                .context(AccessCatalogSnafu)?;
+        }
+
+        self.data.state = DropTableState::EngineDropTable;
+        // Assign procedure id to the subprocedure.
+        self.data.subprocedure_id = Some(ProcedureId::random());
+
+        Ok(Status::executing(true))
+    }
+
+    async fn on_engine_drop_table(&mut self, ctx: &Context) -> Result<Status> {
+        // Safety: subprocedure id is always set in this state.
+        let sub_id = self.data.subprocedure_id.unwrap();
+
+        // Query subprocedure state.
+        let Some(sub_state) = ctx.provider.procedure_state(sub_id).await? else {
+            logging::info!(
+                "On engine drop table {}, subprocedure not found, sub_id: {}",
+                self.data.request.table_name,
+                sub_id
+            );
+
+            // If the subprocedure is not found, we create a new subprocedure with the same id.
+            let engine_ctx = EngineContext::default();
+            let procedure = self
+                .engine_procedure
+                .drop_table_procedure(&engine_ctx, self.data.request.clone())
+                .map_err(Error::from_error_ext)?;
+            return Ok(Status::Suspended {
+                subprocedures: vec![ProcedureWithId {
+                    id: sub_id,
+                    procedure,
+                }],
+                persist: true,
+            });
+        };
+
+        match sub_state {
+            ProcedureState::Running | ProcedureState::Retrying { .. } => Ok(Status::Suspended {
+                subprocedures: Vec::new(),
+                persist: false,
+            }),
+            ProcedureState::Done => {
+                logging::info!(
+                    "On engine drop table {}, done, sub_id: {}",
+                    self.data.request.table_name,
+                    sub_id
+                );
+
+                Ok(Status::Done)
+            }
+            ProcedureState::Failed { .. } => {
+                // Return error if the subprocedure is failed.
+                SubprocedureFailedSnafu {
+                    subprocedure_id: sub_id,
+                }
+                .fail()?
+            }
+        }
+    }
+}
+
+/// Represents each step while dropping a table in the datanode.
+#[derive(Debug, Serialize, Deserialize)]
+enum DropTableState {
+    /// Validate request and prepare to drop table.
+    Prepare,
+    /// Remove the table from the catalog.
+    RemoveFromCatalog,
+    /// Drop table in the table engine.
+    EngineDropTable,
+}
+
+/// Serializable data of [DropTableProcedure].
+#[derive(Debug, Serialize, Deserialize)]
+struct DropTableData {
+    /// Current state.
+    state: DropTableState,
+    /// Request to drop this table.
+    request: DropTableRequest,
+    /// Id of the subprocedure to drop this table from the engine.
+    ///
+    /// This id is `Some` while the procedure is in [DropTableState::EngineDropTable]
+    /// state.
+    subprocedure_id: Option<ProcedureId>,
+}
+
+impl DropTableData {
+    fn table_ref(&self) -> TableReference {
+        TableReference {
+            catalog: &self.request.catalog_name,
+            schema: &self.request.schema_name,
+            table: &self.request.table_name,
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
+    use table::engine::TableEngine;
+
+    use super::*;
+    use crate::test_util::TestEnv;
+
+    #[tokio::test]
+    async fn test_drop_table_procedure() {
+        let env = TestEnv::new("drop");
+        let table_name = "test_drop";
+        env.create_table(table_name).await;
+
+        let request = DropTableRequest {
+            catalog_name: DEFAULT_CATALOG_NAME.to_string(),
+            schema_name: DEFAULT_SCHEMA_NAME.to_string(),
+            table_name: table_name.to_string(),
+        };
+        let TestEnv {
+            dir: _dir,
+            table_engine,
+            procedure_manager,
+            catalog_manager,
+        } = env;
+        let procedure =
+            DropTableProcedure::new(request, catalog_manager.clone(), table_engine.clone());
+        let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure));
+
+        let mut watcher = procedure_manager.submit(procedure_with_id).await.unwrap();
+        watcher.changed().await.unwrap();
+
+        let catalog = catalog_manager
+            .catalog(DEFAULT_CATALOG_NAME)
+            .unwrap()
+            .unwrap();
+        let schema = catalog.schema(DEFAULT_SCHEMA_NAME).unwrap().unwrap();
+        assert!(schema.table(table_name).await.unwrap().is_none());
+        let ctx = EngineContext::default();
+        assert!(!table_engine.table_exists(
+            &ctx,
+            &TableReference {
+                catalog: DEFAULT_CATALOG_NAME,
+                schema: DEFAULT_SCHEMA_NAME,
+                table: table_name,
+            }
+        ));
+    }
+}

src/table-procedure/src/lib.rs

@@ -16,6 +16,7 @@
 
 mod alter;
 mod create;
+mod drop;
 pub mod error;
 #[cfg(test)]
 mod test_util;
@@ -24,6 +25,7 @@ pub use alter::AlterTableProcedure;
 use catalog::CatalogManagerRef;
 use common_procedure::ProcedureManager;
 pub use create::CreateTableProcedure;
+pub use drop::DropTableProcedure;
 use table::engine::{TableEngineProcedureRef, TableEngineRef};
 
 /// Register all procedure loaders to the procedure manager.
@@ -42,5 +44,10 @@ pub fn register_procedure_loaders(
         table_engine,
         procedure_manager,
     );
-    AlterTableProcedure::register_loader(catalog_manager, engine_procedure, procedure_manager);
+    AlterTableProcedure::register_loader(
+        catalog_manager.clone(),
+        engine_procedure.clone(),
+        procedure_manager,
+    );
+    DropTableProcedure::register_loader(catalog_manager, engine_procedure, procedure_manager);
 }