From 0db3ce18254f66aa386fa646c896177101792b54 Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" Date: Sun, 3 May 2026 00:59:08 +0800 Subject: [PATCH] feat: add undrop and purge table procedures Add soft-drop recovery and cleanup procedures, wire their DDL task handling, and update \`greptime-proto\` so the new tasks can round-trip through protobuf. Files: - \`Cargo.toml\` - \`Cargo.lock\` - \`src/common/meta/src/ddl.rs\` - \`src/common/meta/src/ddl/undrop_table.rs\` - \`src/common/meta/src/ddl/purge_dropped_table.rs\` - \`src/common/meta/src/ddl_manager.rs\` - \`src/common/meta/src/rpc/ddl.rs\` - \`src/common/meta/src/key.rs\` - \`src/common/meta/src/ddl/tests/drop_table.rs\` - \`src/mito2/src/engine/open_test.rs\` Signed-off-by: Lei, HUANG --- src/common/meta/src/ddl.rs | 2 + .../meta/src/ddl/purge_dropped_table.rs | 202 +++++++++++ src/common/meta/src/ddl/tests/drop_table.rs | 330 +++++++++++++++++- src/common/meta/src/ddl/undrop_table.rs | 298 ++++++++++++++++ src/common/meta/src/ddl_manager.rs | 74 +++- src/common/meta/src/key.rs | 5 +- src/common/meta/src/rpc/ddl.rs | 186 +++++++++- src/mito2/src/engine/open_test.rs | 64 ++++ 8 files changed, 1156 insertions(+), 5 deletions(-) create mode 100644 src/common/meta/src/ddl/purge_dropped_table.rs create mode 100644 src/common/meta/src/ddl/undrop_table.rs diff --git a/src/common/meta/src/ddl.rs b/src/common/meta/src/ddl.rs index 80130be5d6..71f9500c86 100644 --- a/src/common/meta/src/ddl.rs +++ b/src/common/meta/src/ddl.rs @@ -44,12 +44,14 @@ pub mod drop_flow; pub mod drop_table; pub mod drop_view; pub mod flow_meta; +pub mod purge_dropped_table; pub mod table_meta; #[cfg(any(test, feature = "testing"))] pub mod test_util; #[cfg(test)] pub(crate) mod tests; pub mod truncate_table; +pub mod undrop_table; pub mod utils; /// Metadata allocated to a table. diff --git a/src/common/meta/src/ddl/purge_dropped_table.rs b/src/common/meta/src/ddl/purge_dropped_table.rs new file mode 100644 index 0000000000..2ddd6bb094 --- /dev/null +++ b/src/common/meta/src/ddl/purge_dropped_table.rs @@ -0,0 +1,202 @@ +// 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. + +use std::collections::HashMap; + +use async_trait::async_trait; +use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; +use common_procedure::{ + Context as ProcedureContext, LockKey, Procedure, Result as ProcedureResult, Status, +}; +use common_wal::options::WalOptions; +use serde::{Deserialize, Serialize}; +use snafu::ResultExt; +use store_api::storage::{RegionNumber, TableId}; +use strum::AsRefStr; +use table::table_name::TableName; + +use crate::ddl::DdlContext; +use crate::ddl::drop_table::executor::DropTableExecutor; +use crate::ddl::utils::map_to_procedure_error; +use crate::error::Result; +use crate::key::table_route::TableRouteValue; +use crate::lock_key::{CatalogLock, SchemaLock, TableLock, TableNameLock}; +use crate::rpc::ddl::PurgeDroppedTableTask; +use crate::rpc::router::RegionRoute; + +pub struct PurgeDroppedTableProcedure { + context: DdlContext, + data: PurgeDroppedTableData, +} + +impl PurgeDroppedTableProcedure { + pub const TYPE_NAME: &'static str = "metasrv-procedure::PurgeDroppedTable"; + + pub fn new(task: PurgeDroppedTableTask, context: DdlContext) -> Self { + Self { + context, + data: PurgeDroppedTableData::new(task), + } + } + + pub fn from_json(json: &str, context: DdlContext) -> ProcedureResult { + let data: PurgeDroppedTableData = serde_json::from_str(json).context(FromJsonSnafu)?; + Ok(Self { context, data }) + } + + async fn on_prepare(&mut self) -> Result { + let dropped_table = if let Some(table_id) = self.data.task.table_id { + self.context + .table_metadata_manager + .get_dropped_table_by_id(table_id) + .await? + } else { + self.context + .table_metadata_manager + .get_dropped_table(&self.data.task.table_name()) + .await? + }; + + let Some(dropped_table) = dropped_table else { + return Ok(Status::done()); + }; + self.data.table_id = Some(dropped_table.table_id); + self.data.table_name = Some(dropped_table.table_name); + self.data.table_route_value = Some(dropped_table.table_route_value); + self.data.region_wal_options = dropped_table.region_wal_options; + self.data.state = PurgeDroppedTableState::DropRegions; + Ok(Status::executing(true)) + } + + async fn on_drop_regions(&mut self) -> Result { + if let Some(region_routes) = self.data.physical_region_routes() { + self.executor() + .on_drop_regions( + &self.context.node_manager, + &self.context.leader_region_registry, + region_routes, + false, + false, + false, + ) + .await?; + } + self.data.state = PurgeDroppedTableState::DeleteTombstone; + Ok(Status::executing(true)) + } + + async fn on_delete_tombstone(&mut self) -> Result { + self.context + .table_metadata_manager + .delete_table_metadata_tombstone( + self.data.table_id(), + self.data.table_name(), + self.data.table_route_value(), + &self.data.region_wal_options, + ) + .await?; + Ok(Status::done()) + } + + fn executor(&self) -> DropTableExecutor { + DropTableExecutor::new(self.data.table_name().clone(), self.data.table_id(), false) + } +} + +#[async_trait] +impl Procedure for PurgeDroppedTableProcedure { + fn type_name(&self) -> &str { + Self::TYPE_NAME + } + + fn recover(&mut self) -> ProcedureResult<()> { + Ok(()) + } + + async fn execute(&mut self, _: &ProcedureContext) -> ProcedureResult { + match self.data.state { + PurgeDroppedTableState::Prepare => self.on_prepare().await, + PurgeDroppedTableState::DropRegions => self.on_drop_regions().await, + PurgeDroppedTableState::DeleteTombstone => self.on_delete_tombstone().await, + } + .map_err(map_to_procedure_error) + } + + fn dump(&self) -> ProcedureResult { + serde_json::to_string(&self.data).context(ToJsonSnafu) + } + + fn lock_key(&self) -> LockKey { + let table_ref = self.data.task.table_ref(); + let mut keys = vec![ + CatalogLock::Read(table_ref.catalog).into(), + SchemaLock::read(table_ref.catalog, table_ref.schema).into(), + TableNameLock::new(table_ref.catalog, table_ref.schema, table_ref.table).into(), + ]; + if let Some(table_id) = self.data.task.table_id { + keys.push(TableLock::Write(table_id).into()); + } + LockKey::new(keys) + } +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct PurgeDroppedTableData { + state: PurgeDroppedTableState, + task: PurgeDroppedTableTask, + table_id: Option, + table_name: Option, + table_route_value: Option, + #[serde(default)] + region_wal_options: HashMap, +} + +impl PurgeDroppedTableData { + fn new(task: PurgeDroppedTableTask) -> Self { + Self { + state: PurgeDroppedTableState::Prepare, + task, + table_id: None, + table_name: None, + table_route_value: None, + region_wal_options: HashMap::new(), + } + } + + fn table_id(&self) -> TableId { + self.table_id.unwrap() + } + + fn table_name(&self) -> &TableName { + self.table_name.as_ref().unwrap() + } + + fn table_route_value(&self) -> &TableRouteValue { + self.table_route_value.as_ref().unwrap() + } + + fn physical_region_routes(&self) -> Option<&[RegionRoute]> { + match self.table_route_value() { + TableRouteValue::Physical(route) => Some(&route.region_routes), + TableRouteValue::Logical(_) => None, + } + } +} + +#[derive(Debug, Serialize, Deserialize, AsRefStr, PartialEq)] +enum PurgeDroppedTableState { + Prepare, + DropRegions, + DeleteTombstone, +} diff --git a/src/common/meta/src/ddl/tests/drop_table.rs b/src/common/meta/src/ddl/tests/drop_table.rs index c5bdbaa5c7..0c9873613b 100644 --- a/src/common/meta/src/ddl/tests/drop_table.rs +++ b/src/common/meta/src/ddl/tests/drop_table.rs @@ -33,18 +33,20 @@ use crate::ddl::TableMetadata; use crate::ddl::create_logical_tables::CreateLogicalTablesProcedure; use crate::ddl::create_table::CreateTableProcedure; use crate::ddl::drop_table::{DropTableProcedure, DropTableState}; +use crate::ddl::purge_dropped_table::PurgeDroppedTableProcedure; use crate::ddl::test_util::create_table::test_create_table_task; use crate::ddl::test_util::datanode_handler::{DatanodeWatcher, NaiveDatanodeHandler}; use crate::ddl::test_util::{ create_logical_table, create_physical_table, create_physical_table_metadata, put_datanode_address, test_create_logical_table_task, test_create_physical_table_task, }; +use crate::ddl::undrop_table::UndropTableProcedure; use crate::error::Error; use crate::key::table_name::TableNameKey; use crate::key::table_route::TableRouteValue; use crate::kv_backend::memory::MemoryKvBackend; use crate::peer::Peer; -use crate::rpc::ddl::DropTableTask; +use crate::rpc::ddl::{DropTableTask, PurgeDroppedTableTask, UndropTableTask}; use crate::rpc::router::{Region, RegionRoute}; use crate::test_util::{MockDatanodeManager, new_ddl_context, new_ddl_context_with_kv_backend}; @@ -495,6 +497,311 @@ async fn test_hard_drop_recreated_table_ignores_previous_orphan_tombstone() { procedure.on_prepare().await.unwrap(); } +#[tokio::test] +async fn test_undrop_table_restores_metadata_and_reopens_regions() { + let (tx, mut rx) = mpsc::channel(8); + let datanode_handler = DatanodeWatcher::new(tx); + let node_manager = Arc::new(MockDatanodeManager::new(datanode_handler)); + let mut ddl_context = new_ddl_context(node_manager); + ddl_context.soft_drop_enabled = true; + let table_id = 1024; + let table_name = "foo"; + let task = test_create_table_task(table_name, table_id); + ddl_context + .table_metadata_manager + .create_table_metadata( + task.table_info.clone(), + TableRouteValue::physical(vec![RegionRoute { + region: Region::new_test(RegionId::new(table_id, 1)), + leader_peer: Some(Peer::empty(1)), + follower_peers: vec![Peer::empty(2)], + leader_state: None, + leader_down_since: None, + write_route_policy: None, + }]), + HashMap::new(), + ) + .await + .unwrap(); + + let mut drop_procedure = DropTableProcedure::new( + new_drop_table_task(table_name, table_id, false), + ddl_context.clone(), + ); + execute_procedure_until_done(&mut drop_procedure).await; + + while rx.try_recv().is_ok() {} + + let mut procedure = UndropTableProcedure::new( + new_undrop_table_task(table_name, table_id), + ddl_context.clone(), + ); + execute_procedure_until_done(&mut procedure).await; + + let live_table = ddl_context + .table_metadata_manager + .table_name_manager() + .get(TableNameKey::new( + DEFAULT_CATALOG_NAME, + DEFAULT_SCHEMA_NAME, + table_name, + )) + .await + .unwrap() + .unwrap(); + assert_eq!(live_table.table_id(), table_id); + assert!( + ddl_context + .table_metadata_manager + .get_dropped_table(&drop_procedure.data.task.table_name()) + .await + .unwrap() + .is_none() + ); + + let (peer, request) = rx.try_recv().unwrap(); + assert_eq!(peer.id, 1); + let Some(region_request::Body::Open(req)) = request.body else { + unreachable!(); + }; + assert_eq!(req.region_id, RegionId::new(table_id, 1).as_u64()); + assert!(rx.try_recv().is_err()); +} + +#[tokio::test] +async fn test_undrop_logical_table_skips_datanode_open() { + let (tx, mut rx) = mpsc::channel(8); + let datanode_handler = DatanodeWatcher::new(tx); + let node_manager = Arc::new(MockDatanodeManager::new(datanode_handler)); + let mut ddl_context = new_ddl_context(node_manager); + ddl_context.soft_drop_enabled = true; + let physical_table_id = 1024; + let logical_table_id = 1025; + let table_name = "foo"; + ddl_context + .table_metadata_manager + .create_table_metadata( + test_create_table_task("phy", physical_table_id).table_info, + TableRouteValue::physical(vec![RegionRoute { + region: Region::new_test(RegionId::new(physical_table_id, 1)), + leader_peer: Some(Peer::empty(1)), + follower_peers: vec![], + leader_state: None, + leader_down_since: None, + write_route_policy: None, + }]), + HashMap::new(), + ) + .await + .unwrap(); + let task = test_create_table_task(table_name, logical_table_id); + ddl_context + .table_metadata_manager + .create_table_metadata( + task.table_info.clone(), + TableRouteValue::logical(physical_table_id), + HashMap::new(), + ) + .await + .unwrap(); + + let mut drop_procedure = DropTableProcedure::new( + new_drop_table_task(table_name, logical_table_id, false), + ddl_context.clone(), + ); + execute_procedure_until_done(&mut drop_procedure).await; + + while rx.try_recv().is_ok() {} + + let mut procedure = UndropTableProcedure::new( + new_undrop_table_task(table_name, logical_table_id), + ddl_context.clone(), + ); + execute_procedure_until_done(&mut procedure).await; + + let live_table = ddl_context + .table_metadata_manager + .table_name_manager() + .get(TableNameKey::new( + DEFAULT_CATALOG_NAME, + DEFAULT_SCHEMA_NAME, + table_name, + )) + .await + .unwrap() + .unwrap(); + assert_eq!(live_table.table_id(), logical_table_id); + assert!(rx.try_recv().is_err()); +} + +#[tokio::test] +async fn test_undrop_table_fails_when_live_name_exists() { + let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)); + let mut ddl_context = new_ddl_context(node_manager); + ddl_context.soft_drop_enabled = true; + let dropped_table_id = 1024; + let live_table_id = 1025; + let table_name = "foo"; + let task = test_create_table_task(table_name, dropped_table_id); + ddl_context + .table_metadata_manager + .create_table_metadata( + task.table_info.clone(), + TableRouteValue::physical(vec![]), + HashMap::new(), + ) + .await + .unwrap(); + let mut drop_procedure = DropTableProcedure::new( + new_drop_table_task(table_name, dropped_table_id, false), + ddl_context.clone(), + ); + execute_procedure_until_done(&mut drop_procedure).await; + ddl_context + .table_metadata_manager + .create_table_metadata( + test_create_table_task(table_name, live_table_id).table_info, + TableRouteValue::physical(vec![]), + HashMap::new(), + ) + .await + .unwrap(); + + let mut procedure = UndropTableProcedure::new( + new_undrop_table_task(table_name, dropped_table_id), + ddl_context, + ); + let err = procedure.on_prepare().await.unwrap_err(); + + assert_matches!(err, Error::TableAlreadyExists { .. }); +} + +#[tokio::test] +async fn test_purge_dropped_table_drops_regions_and_deletes_tombstone() { + let (tx, mut rx) = mpsc::channel(8); + let datanode_handler = DatanodeWatcher::new(tx); + let node_manager = Arc::new(MockDatanodeManager::new(datanode_handler)); + let mut ddl_context = new_ddl_context(node_manager); + ddl_context.soft_drop_enabled = true; + let table_id = 1024; + let table_name = "foo"; + let task = test_create_table_task(table_name, table_id); + ddl_context + .table_metadata_manager + .create_table_metadata( + task.table_info.clone(), + TableRouteValue::physical(vec![RegionRoute { + region: Region::new_test(RegionId::new(table_id, 1)), + leader_peer: Some(Peer::empty(1)), + follower_peers: vec![Peer::empty(2)], + leader_state: None, + leader_down_since: None, + write_route_policy: None, + }]), + HashMap::new(), + ) + .await + .unwrap(); + let mut drop_procedure = DropTableProcedure::new( + new_drop_table_task(table_name, table_id, false), + ddl_context.clone(), + ); + execute_procedure_until_done(&mut drop_procedure).await; + while rx.try_recv().is_ok() {} + + let mut procedure = PurgeDroppedTableProcedure::new( + new_purge_dropped_table_task(table_name, Some(table_id)), + ddl_context.clone(), + ); + execute_procedure_until_done(&mut procedure).await; + + let mut requests = Vec::new(); + for _ in 0..2 { + let (peer, request) = rx.try_recv().unwrap(); + requests.push((peer.id, request.body.unwrap())); + } + requests.sort_unstable_by_key(|(peer_id, _)| *peer_id); + assert_matches!(requests[0].1, region_request::Body::Drop(_)); + assert_matches!(requests[1].1, region_request::Body::Close(_)); + assert!( + ddl_context + .table_metadata_manager + .get_dropped_table(&drop_procedure.data.task.table_name()) + .await + .unwrap() + .is_none() + ); +} + +#[tokio::test] +async fn test_purge_dropped_table_by_name_selects_tombstone_when_live_table_exists() { + let (tx, mut rx) = mpsc::channel(8); + let datanode_handler = DatanodeWatcher::new(tx); + let node_manager = Arc::new(MockDatanodeManager::new(datanode_handler)); + let mut ddl_context = new_ddl_context(node_manager); + ddl_context.soft_drop_enabled = true; + let dropped_table_id = 1024; + let live_table_id = 1025; + let table_name = "foo"; + ddl_context + .table_metadata_manager + .create_table_metadata( + test_create_table_task(table_name, dropped_table_id).table_info, + TableRouteValue::physical(vec![RegionRoute { + region: Region::new_test(RegionId::new(dropped_table_id, 1)), + leader_peer: Some(Peer::empty(1)), + follower_peers: vec![], + leader_state: None, + leader_down_since: None, + write_route_policy: None, + }]), + HashMap::new(), + ) + .await + .unwrap(); + let mut drop_procedure = DropTableProcedure::new( + new_drop_table_task(table_name, dropped_table_id, false), + ddl_context.clone(), + ); + execute_procedure_until_done(&mut drop_procedure).await; + ddl_context + .table_metadata_manager + .create_table_metadata( + test_create_table_task(table_name, live_table_id).table_info, + TableRouteValue::physical(vec![]), + HashMap::new(), + ) + .await + .unwrap(); + while rx.try_recv().is_ok() {} + + let mut procedure = PurgeDroppedTableProcedure::new( + new_purge_dropped_table_task(table_name, None), + ddl_context.clone(), + ); + execute_procedure_until_done(&mut procedure).await; + + let live_table = ddl_context + .table_metadata_manager + .table_name_manager() + .get(TableNameKey::new( + DEFAULT_CATALOG_NAME, + DEFAULT_SCHEMA_NAME, + table_name, + )) + .await + .unwrap() + .unwrap(); + assert_eq!(live_table.table_id(), live_table_id); + + let (_, request) = rx.try_recv().unwrap(); + let Some(region_request::Body::Drop(req)) = request.body else { + unreachable!(); + }; + assert_eq!(req.region_id, RegionId::new(dropped_table_id, 1).as_u64()); + assert!(rx.try_recv().is_err()); +} + #[tokio::test] async fn test_on_rollback() { let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)); @@ -566,6 +873,27 @@ fn new_drop_table_task(table_name: &str, table_id: TableId, drop_if_exists: bool } } +fn new_undrop_table_task(table_name: &str, table_id: TableId) -> UndropTableTask { + UndropTableTask { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table: table_name.to_string(), + table_id, + } +} + +fn new_purge_dropped_table_task( + table_name: &str, + table_id: Option, +) -> PurgeDroppedTableTask { + PurgeDroppedTableTask { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table: table_name.to_string(), + table_id, + } +} + #[tokio::test] async fn test_memory_region_keeper_guard_dropped_on_procedure_done() { let node_manager = Arc::new(MockDatanodeManager::new(NaiveDatanodeHandler)); diff --git a/src/common/meta/src/ddl/undrop_table.rs b/src/common/meta/src/ddl/undrop_table.rs new file mode 100644 index 0000000000..9cbfce63c4 --- /dev/null +++ b/src/common/meta/src/ddl/undrop_table.rs @@ -0,0 +1,298 @@ +// 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. + +use std::collections::HashMap; + +use api::v1::region::{ + OpenRequest as PbOpenRegionRequest, RegionRequest, RegionRequestHeader, region_request, +}; +use async_trait::async_trait; +use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; +use common_procedure::{ + Context as ProcedureContext, LockKey, Procedure, Result as ProcedureResult, Status, +}; +use common_telemetry::tracing_context::TracingContext; +use common_wal::options::WalOptions; +use futures::future::join_all; +use serde::{Deserialize, Serialize}; +use snafu::{OptionExt, ResultExt, ensure}; +use store_api::storage::{RegionId, RegionNumber}; +use strum::AsRefStr; +use table::metadata::TableId; +use table::table_name::TableName; +use table::table_reference::TableReference; + +use crate::ddl::utils::{add_peer_context_if_needed, map_to_procedure_error, region_storage_path}; +use crate::ddl::{CreateRequestBuilder, DdlContext, build_template_from_raw_table_info}; +use crate::error::{self, Result}; +use crate::instruction::CacheIdent; +use crate::key::table_name::TableNameKey; +use crate::key::table_route::TableRouteValue; +use crate::lock_key::{CatalogLock, SchemaLock, TableLock, TableNameLock}; +use crate::rpc::ddl::UndropTableTask; +use crate::rpc::router::{RegionRoute, find_leader_regions, find_leaders}; + +pub struct UndropTableProcedure { + context: DdlContext, + data: UndropTableData, +} + +impl UndropTableProcedure { + pub const TYPE_NAME: &'static str = "metasrv-procedure::UndropTable"; + + pub fn new(task: UndropTableTask, context: DdlContext) -> Self { + Self { + context, + data: UndropTableData::new(task), + } + } + + pub fn from_json(json: &str, context: DdlContext) -> ProcedureResult { + let data: UndropTableData = serde_json::from_str(json).context(FromJsonSnafu)?; + Ok(Self { context, data }) + } + + pub(crate) async fn on_prepare(&mut self) -> Result { + let table_ref = self.data.table_ref(); + ensure!( + !self + .context + .table_metadata_manager + .table_name_manager() + .exists(TableNameKey::new( + table_ref.catalog, + table_ref.schema, + table_ref.table + )) + .await?, + error::TableAlreadyExistsSnafu { + table_name: table_ref.to_string() + } + ); + + let dropped_table = self + .context + .table_metadata_manager + .get_dropped_table_by_id(self.data.task.table_id) + .await? + .with_context(|| error::TableNotFoundSnafu { + table_name: table_ref.to_string(), + })?; + self.data.table_name = Some(dropped_table.table_name.clone()); + self.data.table_route_value = Some(dropped_table.table_route_value.clone()); + self.data.region_wal_options = dropped_table.region_wal_options; + self.data.table_info = Some(dropped_table.table_info_value.table_info); + self.data.state = UndropTableState::RestoreMetadata; + Ok(Status::executing(true)) + } + + async fn on_restore_metadata(&mut self) -> Result { + let table_route_value = self.data.table_route_value(); + self.context + .table_metadata_manager + .restore_table_metadata( + self.data.task.table_id, + self.data.table_name(), + table_route_value, + &self.data.region_wal_options, + ) + .await?; + self.data.state = UndropTableState::OpenRegions; + Ok(Status::executing(true)) + } + + async fn on_open_regions(&mut self) -> Result { + let TableRouteValue::Physical(route) = self.data.table_route_value() else { + self.data.state = UndropTableState::InvalidateTableCache; + return Ok(Status::executing(true)); + }; + + open_regions( + &self.context, + self.data.task.table_id, + self.data.table_name(), + self.data.table_info(), + &route.region_routes, + &self.data.region_wal_options, + ) + .await?; + self.data.state = UndropTableState::InvalidateTableCache; + Ok(Status::executing(true)) + } + + async fn on_broadcast(&mut self) -> Result { + let ctx = crate::cache_invalidator::Context { + subject: Some(format!( + "Invalidate table cache by undropping table {}, table_id: {}", + self.data.table_name().table_ref(), + self.data.task.table_id, + )), + }; + self.context + .cache_invalidator + .invalidate( + &ctx, + &[ + CacheIdent::TableName(self.data.table_name().table_ref().into()), + CacheIdent::TableId(self.data.task.table_id), + ], + ) + .await?; + Ok(Status::done()) + } +} + +#[async_trait] +impl Procedure for UndropTableProcedure { + fn type_name(&self) -> &str { + Self::TYPE_NAME + } + + fn recover(&mut self) -> ProcedureResult<()> { + Ok(()) + } + + async fn execute(&mut self, _: &ProcedureContext) -> ProcedureResult { + match self.data.state { + UndropTableState::Prepare => self.on_prepare().await, + UndropTableState::RestoreMetadata => self.on_restore_metadata().await, + UndropTableState::OpenRegions => self.on_open_regions().await, + UndropTableState::InvalidateTableCache => self.on_broadcast().await, + } + .map_err(map_to_procedure_error) + } + + fn dump(&self) -> ProcedureResult { + serde_json::to_string(&self.data).context(ToJsonSnafu) + } + + fn lock_key(&self) -> LockKey { + let table_ref = self.data.table_ref(); + LockKey::new(vec![ + CatalogLock::Read(table_ref.catalog).into(), + SchemaLock::read(table_ref.catalog, table_ref.schema).into(), + TableNameLock::new(table_ref.catalog, table_ref.schema, table_ref.table).into(), + TableLock::Write(self.data.task.table_id).into(), + ]) + } +} + +async fn open_regions( + context: &DdlContext, + table_id: TableId, + table_name: &TableName, + table_info: &table::metadata::TableInfo, + region_routes: &[RegionRoute], + region_wal_options: &HashMap, +) -> Result<()> { + let template = build_template_from_raw_table_info(table_info)?; + let builder = CreateRequestBuilder::new(template, None); + let storage_path = region_storage_path(&table_name.catalog_name, &table_name.schema_name); + let wal_options = region_wal_options + .iter() + .map(|(region_number, wal_options)| { + serde_json::to_string(wal_options) + .map(|wal_options| (*region_number, wal_options)) + .context(error::SerdeJsonSnafu) + }) + .collect::>>()?; + let leaders = find_leaders(region_routes); + let mut tasks = Vec::with_capacity(leaders.len()); + for datanode in leaders { + let requester = context.node_manager.datanode(&datanode).await; + for region_number in find_leader_regions(region_routes, &datanode) { + let region_id = RegionId::new(table_id, region_number); + let create_request = builder.build_one( + region_id, + storage_path.clone(), + &wal_options, + &HashMap::new(), + ); + let request = RegionRequest { + header: Some(RegionRequestHeader { + tracing_context: TracingContext::from_current_span().to_w3c(), + ..Default::default() + }), + body: Some(region_request::Body::Open(PbOpenRegionRequest { + region_id: create_request.region_id, + engine: create_request.engine, + path: create_request.path, + options: create_request.options, + })), + }; + let datanode = datanode.clone(); + let requester = requester.clone(); + tasks.push(async move { + requester + .handle(request) + .await + .map_err(add_peer_context_if_needed(datanode)) + }); + } + } + + join_all(tasks) + .await + .into_iter() + .collect::>>()?; + Ok(()) +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct UndropTableData { + state: UndropTableState, + task: UndropTableTask, + table_name: Option, + table_info: Option, + table_route_value: Option, + #[serde(default)] + region_wal_options: HashMap, +} + +impl UndropTableData { + fn new(task: UndropTableTask) -> Self { + Self { + state: UndropTableState::Prepare, + task, + table_name: None, + table_info: None, + table_route_value: None, + region_wal_options: HashMap::new(), + } + } + + fn table_ref(&self) -> TableReference<'_> { + self.task.table_ref() + } + + fn table_name(&self) -> &TableName { + self.table_name.as_ref().unwrap() + } + + fn table_info(&self) -> &table::metadata::TableInfo { + self.table_info.as_ref().unwrap() + } + + fn table_route_value(&self) -> &TableRouteValue { + self.table_route_value.as_ref().unwrap() + } +} + +#[derive(Debug, Serialize, Deserialize, AsRefStr, PartialEq)] +enum UndropTableState { + Prepare, + RestoreMetadata, + OpenRegions, + InvalidateTableCache, +} diff --git a/src/common/meta/src/ddl_manager.rs b/src/common/meta/src/ddl_manager.rs index bba40c1a7e..bfd858a196 100644 --- a/src/common/meta/src/ddl_manager.rs +++ b/src/common/meta/src/ddl_manager.rs @@ -43,7 +43,9 @@ use crate::ddl::drop_database::DropDatabaseProcedure; use crate::ddl::drop_flow::DropFlowProcedure; use crate::ddl::drop_table::DropTableProcedure; use crate::ddl::drop_view::DropViewProcedure; +use crate::ddl::purge_dropped_table::PurgeDroppedTableProcedure; use crate::ddl::truncate_table::TruncateTableProcedure; +use crate::ddl::undrop_table::UndropTableProcedure; use crate::ddl::{DdlContext, utils}; use crate::error::{ self, CreateRepartitionProcedureSnafu, EmptyDdlTasksSnafu, ProcedureOutputSnafu, @@ -62,7 +64,7 @@ use crate::rpc::ddl::DdlTask::DropTrigger; use crate::rpc::ddl::DdlTask::{ AlterDatabase, AlterLogicalTables, AlterTable, CommentOn, CreateDatabase, CreateFlow, CreateLogicalTables, CreateTable, CreateView, DropDatabase, DropFlow, DropLogicalTables, - DropTable, DropView, TruncateTable, + DropTable, DropView, PurgeDroppedTable, TruncateTable, UndropTable, }; #[cfg(feature = "enterprise")] use crate::rpc::ddl::trigger::CreateTriggerTask; @@ -71,7 +73,8 @@ use crate::rpc::ddl::trigger::DropTriggerTask; use crate::rpc::ddl::{ AlterDatabaseTask, AlterTableTask, CommentOnTask, CreateDatabaseTask, CreateFlowTask, CreateTableTask, CreateViewTask, DropDatabaseTask, DropFlowTask, DropTableTask, DropViewTask, - QueryContext, SubmitDdlTaskRequest, SubmitDdlTaskResponse, TruncateTableTask, + PurgeDroppedTableTask, QueryContext, SubmitDdlTaskRequest, SubmitDdlTaskResponse, + TruncateTableTask, UndropTableTask, }; /// A configurator that customizes or enhances a [`DdlManager`]. @@ -252,6 +255,8 @@ impl DdlManager { AlterLogicalTablesProcedure, AlterDatabaseProcedure, DropTableProcedure, + UndropTableProcedure, + PurgeDroppedTableProcedure, DropFlowProcedure, TruncateTableProcedure, CreateDatabaseProcedure, @@ -468,6 +473,32 @@ impl DdlManager { self.execute_procedure_and_wait(procedure_with_id).await } + /// Submits and executes an undrop table task. + #[tracing::instrument(skip_all)] + pub async fn submit_undrop_table_task( + &self, + undrop_table_task: UndropTableTask, + ) -> Result<(ProcedureId, Option)> { + let context = self.create_context(); + let procedure = UndropTableProcedure::new(undrop_table_task, context); + let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + + self.execute_procedure_and_wait(procedure_with_id).await + } + + /// Submits and executes a purge dropped table task. + #[tracing::instrument(skip_all)] + pub async fn submit_purge_dropped_table_task( + &self, + purge_dropped_table_task: PurgeDroppedTableTask, + ) -> Result<(ProcedureId, Option)> { + let context = self.create_context(); + let procedure = PurgeDroppedTableProcedure::new(purge_dropped_table_task, context); + let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + + self.execute_procedure_and_wait(procedure_with_id).await + } + /// Submits and executes a create database task. #[tracing::instrument(skip_all)] pub async fn submit_create_database( @@ -643,6 +674,12 @@ impl DdlManager { handle_create_table_task(self, create_table_task, request.query_context).await } DropTable(drop_table_task) => handle_drop_table_task(self, drop_table_task).await, + UndropTable(undrop_table_task) => { + handle_undrop_table_task(self, undrop_table_task).await + } + PurgeDroppedTable(purge_dropped_table_task) => { + handle_purge_dropped_table_task(self, purge_dropped_table_task).await + } AlterTable(alter_table_task) => { handle_alter_table_task(self, alter_table_task, ddl_options).await } @@ -790,6 +827,39 @@ async fn handle_drop_table_task( }) } +async fn handle_undrop_table_task( + ddl_manager: &DdlManager, + undrop_table_task: UndropTableTask, +) -> Result { + let table_id = undrop_table_task.table_id; + let (id, _) = ddl_manager + .submit_undrop_table_task(undrop_table_task) + .await?; + + info!("Table: {table_id} is undropped via procedure_id {id:?}"); + + Ok(SubmitDdlTaskResponse { + key: id.to_string().into(), + ..Default::default() + }) +} + +async fn handle_purge_dropped_table_task( + ddl_manager: &DdlManager, + purge_dropped_table_task: PurgeDroppedTableTask, +) -> Result { + let (id, _) = ddl_manager + .submit_purge_dropped_table_task(purge_dropped_table_task) + .await?; + + info!("Dropped table is purged via procedure_id {id:?}"); + + Ok(SubmitDdlTaskResponse { + key: id.to_string().into(), + ..Default::default() + }) +} + async fn handle_create_table_task( ddl_manager: &DdlManager, create_table_task: CreateTableTask, diff --git a/src/common/meta/src/key.rs b/src/common/meta/src/key.rs index 0be3ee1a45..4f52822d8c 100644 --- a/src/common/meta/src/key.rs +++ b/src/common/meta/src/key.rs @@ -1155,7 +1155,10 @@ impl TableMetadataManager { table_route_value: &TableRouteValue, ) -> Result> { let mut region_wal_options = HashMap::new(); - let datanode_table_keys = region_distribution(table_route_value.region_routes()?) + let Some(region_routes) = table_route_value.region_routes().ok() else { + return Ok(region_wal_options); + }; + let datanode_table_keys = region_distribution(region_routes) .into_keys() .map(|datanode_id| DatanodeTableKey::new(datanode_id, table_id)) .collect::>(); diff --git a/src/common/meta/src/rpc/ddl.rs b/src/common/meta/src/rpc/ddl.rs index fd11e4c310..d78c72b31a 100644 --- a/src/common/meta/src/rpc/ddl.rs +++ b/src/common/meta/src/rpc/ddl.rs @@ -31,7 +31,8 @@ use api::v1::meta::{ DdlTaskResponse as PbDdlTaskResponse, DropDatabaseTask as PbDropDatabaseTask, DropFlowTask as PbDropFlowTask, DropTableTask as PbDropTableTask, DropTableTasks as PbDropTableTasks, DropViewTask as PbDropViewTask, Partition, ProcedureId, - TruncateTableTask as PbTruncateTableTask, + PurgeDroppedTableTask as PbPurgeDroppedTableTask, TruncateTableTask as PbTruncateTableTask, + UndropTableTask as PbUndropTableTask, }; use api::v1::{ AlterDatabaseExpr, AlterTableExpr, CommentObjectType as PbCommentObjectType, CommentOnExpr, @@ -71,6 +72,8 @@ pub const ORIGIN_FRONTEND_ADDR_EXTENSION_KEY: &str = "__greptime_origin_frontend pub enum DdlTask { CreateTable(CreateTableTask), DropTable(DropTableTask), + UndropTable(UndropTableTask), + PurgeDroppedTable(PurgeDroppedTableTask), AlterTable(AlterTableTask), TruncateTable(TruncateTableTask), CreateLogicalTables(Vec), @@ -152,6 +155,36 @@ impl DdlTask { }) } + /// Creates a [`DdlTask`] to undrop a table. + pub fn new_undrop_table( + catalog: String, + schema: String, + table: String, + table_id: TableId, + ) -> Self { + DdlTask::UndropTable(UndropTableTask { + catalog, + schema, + table, + table_id, + }) + } + + /// Creates a [`DdlTask`] to purge a dropped table. + pub fn new_purge_dropped_table( + catalog: String, + schema: String, + table: String, + table_id: Option, + ) -> Self { + DdlTask::PurgeDroppedTable(PurgeDroppedTableTask { + catalog, + schema, + table, + table_id, + }) + } + /// Creates a [`DdlTask`] to create a database. pub fn new_create_database( catalog: String, @@ -225,6 +258,12 @@ impl TryFrom for DdlTask { Ok(DdlTask::CreateTable(create_table.try_into()?)) } Task::DropTableTask(drop_table) => Ok(DdlTask::DropTable(drop_table.try_into()?)), + Task::UndropTableTask(undrop_table) => { + Ok(DdlTask::UndropTable(undrop_table.try_into()?)) + } + Task::PurgeDroppedTableTask(purge_dropped_table) => { + Ok(DdlTask::PurgeDroppedTable(purge_dropped_table.into())) + } Task::AlterTableTask(alter_table) => Ok(DdlTask::AlterTable(alter_table.try_into()?)), Task::TruncateTableTask(truncate_table) => { Ok(DdlTask::TruncateTable(truncate_table.try_into()?)) @@ -335,6 +374,8 @@ impl TryFrom for PbDdlTaskRequest { let task = match request.task { DdlTask::CreateTable(task) => Task::CreateTableTask(task.try_into()?), DdlTask::DropTable(task) => Task::DropTableTask(task.into()), + DdlTask::UndropTable(task) => Task::UndropTableTask(task.into()), + DdlTask::PurgeDroppedTable(task) => Task::PurgeDroppedTableTask(task.into()), DdlTask::AlterTable(task) => Task::AlterTableTask(task.try_into()?), DdlTask::TruncateTable(task) => Task::TruncateTableTask(task.try_into()?), DdlTask::CreateLogicalTables(tasks) => { @@ -596,6 +637,109 @@ pub struct DropTableTask { pub drop_if_exists: bool, } +#[derive(Debug, PartialEq, Serialize, Deserialize, Clone)] +pub struct UndropTableTask { + pub catalog: String, + pub schema: String, + pub table: String, + pub table_id: TableId, +} + +impl UndropTableTask { + pub fn table_ref(&self) -> TableReference<'_> { + TableReference { + catalog: &self.catalog, + schema: &self.schema, + table: &self.table, + } + } + + pub fn table_name(&self) -> TableName { + TableName { + catalog_name: self.catalog.clone(), + schema_name: self.schema.clone(), + table_name: self.table.clone(), + } + } +} + +#[derive(Debug, PartialEq, Serialize, Deserialize, Clone)] +pub struct PurgeDroppedTableTask { + pub catalog: String, + pub schema: String, + pub table: String, + pub table_id: Option, +} + +impl PurgeDroppedTableTask { + pub fn table_ref(&self) -> TableReference<'_> { + TableReference { + catalog: &self.catalog, + schema: &self.schema, + table: &self.table, + } + } + + pub fn table_name(&self) -> TableName { + TableName { + catalog_name: self.catalog.clone(), + schema_name: self.schema.clone(), + table_name: self.table.clone(), + } + } +} + +impl TryFrom for UndropTableTask { + type Error = error::Error; + + fn try_from(pb: PbUndropTableTask) -> Result { + Ok(Self { + catalog: pb.catalog_name, + schema: pb.schema_name, + table: pb.table_name, + table_id: pb + .table_id + .context(error::InvalidProtoMsgSnafu { + err_msg: "expected table_id", + })? + .id, + }) + } +} + +impl From for PbUndropTableTask { + fn from(task: UndropTableTask) -> Self { + Self { + catalog_name: task.catalog, + schema_name: task.schema, + table_name: task.table, + table_id: Some(api::v1::TableId { id: task.table_id }), + } + } +} + +impl From for PurgeDroppedTableTask { + fn from(pb: PbPurgeDroppedTableTask) -> Self { + Self { + catalog: pb.catalog_name, + schema: pb.schema_name, + table: pb.table_name, + table_id: pb.table_id.map(|table_id| table_id.id), + } + } +} + +impl From for PbPurgeDroppedTableTask { + fn from(task: PurgeDroppedTableTask) -> Self { + Self { + catalog_name: task.catalog, + schema_name: task.schema, + table_name: task.table, + table_id: task.table_id.map(|id| api::v1::TableId { id }), + } + } +} + impl DropTableTask { pub fn table_ref(&self) -> TableReference<'_> { TableReference { @@ -1764,6 +1908,46 @@ mod tests { assert_eq!(task, de); } + #[test] + fn test_undrop_table_task_pb_roundtrip() { + let expected = UndropTableTask { + catalog: "greptime".to_string(), + schema: "public".to_string(), + table: "foo".to_string(), + table_id: 1024, + }; + let request = SubmitDdlTaskRequest::new( + QueryContext::default(), + DdlTask::UndropTable(expected.clone()), + ); + + let pb = PbDdlTaskRequest::try_from(request).unwrap(); + let pb_task = pb.task.unwrap(); + let de = DdlTask::try_from(pb_task).unwrap(); + + assert!(matches!(de, DdlTask::UndropTable(task) if task == expected)); + } + + #[test] + fn test_purge_dropped_table_task_pb_roundtrip() { + let expected = PurgeDroppedTableTask { + catalog: "greptime".to_string(), + schema: "public".to_string(), + table: "foo".to_string(), + table_id: Some(1024), + }; + let request = SubmitDdlTaskRequest::new( + QueryContext::default(), + DdlTask::PurgeDroppedTable(expected.clone()), + ); + + let pb = PbDdlTaskRequest::try_from(request).unwrap(); + let pb_task = pb.task.unwrap(); + let de = DdlTask::try_from(pb_task).unwrap(); + + assert!(matches!(de, DdlTask::PurgeDroppedTable(task) if task == expected)); + } + #[test] fn test_sort_columns() { // construct RawSchema diff --git a/src/mito2/src/engine/open_test.rs b/src/mito2/src/engine/open_test.rs index b22e17dd01..1f0053aab5 100644 --- a/src/mito2/src/engine/open_test.rs +++ b/src/mito2/src/engine/open_test.rs @@ -125,6 +125,70 @@ async fn test_engine_reopen_region() { test_engine_reopen_region_with_format(true).await; } +#[tokio::test] +async fn test_engine_reopens_closed_soft_dropped_region() { + test_engine_reopens_closed_soft_dropped_region_with_format(false).await; + test_engine_reopens_closed_soft_dropped_region_with_format(true).await; +} + +async fn test_engine_reopens_closed_soft_dropped_region_with_format(flat_format: bool) { + let mut env = TestEnv::with_prefix("reopen-soft-dropped-region").await; + let engine = env + .create_engine(MitoConfig { + default_flat_format: flat_format, + ..Default::default() + }) + .await; + + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + let table_dir = request.table_dir.clone(); + let column_schemas = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + put_rows( + &engine, + region_id, + Rows { + schema: column_schemas.clone(), + rows: build_rows(0, 2), + }, + ) + .await; + + engine + .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .await + .unwrap(); + assert!(!engine.is_region_exists(region_id)); + + engine + .handle_request( + region_id, + RegionRequest::Open(RegionOpenRequest { + engine: String::new(), + table_dir, + path_type: PathType::Bare, + options: HashMap::default(), + skip_wal_replay: false, + checkpoint: None, + }), + ) + .await + .unwrap(); + assert!(engine.is_region_exists(region_id)); + + let scanner = engine + .scanner(region_id, ScanRequest::default()) + .await + .unwrap(); + let stream = scanner.scan().await.unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + assert_eq!(2, batches.iter().map(|b| b.num_rows()).sum::()); +} + async fn test_engine_reopen_region_with_format(flat_format: bool) { let mut env = TestEnv::with_prefix("reopen-region").await; let engine = env