diff --git a/src/datanode/src/sql/alter.rs b/src/datanode/src/sql/alter.rs index caae6535d7b..6ec8669d464 100644 --- a/src/datanode/src/sql/alter.rs +++ b/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. diff --git a/src/datanode/src/sql/drop_table.rs b/src/datanode/src/sql/drop_table.rs index 5ad684a2452..0a24d4cb9a6 100644 --- a/src/datanode/src/sql/drop_table.rs +++ b/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 { + 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 { + 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))); + } } diff --git a/src/table-procedure/src/alter.rs b/src/table-procedure/src/alter.rs index 81cf202edd7..f6b772a357b 100644 --- a/src/table-procedure/src/alter.rs +++ b/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, - /// 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. diff --git a/src/table-procedure/src/drop.rs b/src/table-procedure/src/drop.rs new file mode 100644 index 00000000000..252a136a427 --- /dev/null +++ b/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 { + 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 { + 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 { + let data: DropTableData = serde_json::from_str(json).context(DeserializeProcedureSnafu)?; + + Ok(DropTableProcedure { + data, + catalog_manager, + engine_procedure, + }) + } + + async fn on_prepare(&mut self) -> Result { + 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 { + 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 { + // 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, +} + +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, + } + )); + } +} diff --git a/src/table-procedure/src/lib.rs b/src/table-procedure/src/lib.rs index 7c1b7eaeacf..ae6713b0b00 100644 --- a/src/table-procedure/src/lib.rs +++ b/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); }