common_meta/ddl/drop_table/

executor.rs

// 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;
#[cfg(feature = "enterprise")]
use std::collections::HashSet;

use api::v1::region::{
    CleanUpRequest as PbCleanUpRequest, CloseRequest as PbCloseRegionRequest,
    DropRequest as PbDropRegionRequest, RegionRequest, RegionRequestHeader, region_request,
};
use common_error::ext::ErrorExt;
use common_error::status_code::StatusCode;
use common_telemetry::tracing_context::TracingContext;
use common_telemetry::{debug, error};
use common_wal::options::WalOptions;
use futures::future::join_all;
use snafu::ensure;
use store_api::storage::{RegionId, RegionNumber};
use table::metadata::{TableId, TableInfo};
use table::table_name::TableName;

use crate::cache_invalidator::Context;
use crate::ddl::utils::{
    add_peer_context_if_needed, convert_region_routes_to_detecting_regions, region_storage_path,
};
use crate::ddl::{CreateRequestBuilder, DdlContext, build_template_from_raw_table_info};
use crate::error::{self, Result};
use crate::instruction::CacheIdent;
#[cfg(feature = "enterprise")]
use crate::key::DroppedTableLifecycle;
use crate::key::table_name::TableNameKey;
use crate::key::table_route::TableRouteValue;
use crate::node_manager::NodeManagerRef;
use crate::region_registry::LeaderRegionRegistryRef;
use crate::rpc::router::{
    RegionRoute, find_follower_regions, find_followers, find_leader_regions, find_leaders,
    operating_leader_regions,
};

/// [Control] indicated to the caller whether to go to the next step.
#[derive(Debug)]
pub enum Control<T> {
    Continue(T),
    Stop,
}

impl<T> Control<T> {
    /// Returns true if it's [Control::Stop].
    pub fn stop(&self) -> bool {
        matches!(self, Control::Stop)
    }
}

impl DropTableExecutor {
    /// Returns the [DropTableExecutor].
    pub fn new(table: TableName, table_id: TableId, drop_if_exists: bool) -> Self {
        Self {
            table,
            table_id,
            drop_if_exists,
        }
    }
}

/// [DropTableExecutor] performs:
/// - Drops the metadata of the table.
/// - Invalidates the cache on the Frontend nodes.
/// - Drops the regions on the Datanode nodes.
pub struct DropTableExecutor {
    table: TableName,
    table_id: TableId,
    drop_if_exists: bool,
}

impl DropTableExecutor {
    /// Checks whether table exists.
    /// - Early returns if table not exists and `drop_if_exists` is `true`.
    /// - Throws an error if table not exists and `drop_if_exists` is `false`.
    pub async fn on_prepare(&self, ctx: &DdlContext) -> Result<Control<()>> {
        let table_ref = self.table.table_ref();

        let exist = ctx
            .table_metadata_manager
            .table_name_manager()
            .exists(TableNameKey::new(
                table_ref.catalog,
                table_ref.schema,
                table_ref.table,
            ))
            .await?;

        if !exist && self.drop_if_exists {
            return Ok(Control::Stop);
        }

        ensure!(
            exist,
            error::TableNotFoundSnafu {
                table_name: table_ref.to_string()
            }
        );

        Ok(Control::Continue(()))
    }

    /// Rejects dropping a recreated live table while an older tombstone still owns the same
    /// fully qualified name.
    #[cfg(feature = "enterprise")]
    pub async fn check_tombstone_conflict(
        &self,
        ctx: &DdlContext,
        soft_drop_enabled: bool,
    ) -> Result<()> {
        let table_ref = self.table.table_ref();
        if let Some(dropped_table) = ctx
            .table_metadata_manager
            .get_dropped_table(&self.table)
            .await?
            && dropped_table.table_id != self.table_id
            && (soft_drop_enabled || dropped_table.dropped_at.is_some())
        {
            return error::TableNameTombstoneConflictSnafu {
                table_name: table_ref.to_string(),
                existing_table_id: dropped_table.table_id,
                dropping_table_id: self.table_id,
            }
            .fail();
        }

        Ok(())
    }

    /// Deletes the table metadata **logically**.
    pub async fn on_delete_metadata(
        &self,
        ctx: &DdlContext,
        table_route_value: &TableRouteValue,
        region_wal_options: &HashMap<RegionNumber, WalOptions>,
    ) -> Result<()> {
        ctx.table_metadata_manager
            .delete_table_metadata(
                self.table_id,
                &self.table,
                table_route_value,
                region_wal_options,
                None,
            )
            .await
    }

    /// Soft-deletes table metadata and retains its tombstone lifecycle markers.
    #[cfg(feature = "enterprise")]
    pub async fn on_soft_delete_metadata(
        &self,
        ctx: &DdlContext,
        table_route_value: &TableRouteValue,
        region_wal_options: &HashMap<RegionNumber, WalOptions>,
        dropped_at: Option<i64>,
        retention_expires_at: Option<i64>,
        drop_generation: Option<&str>,
    ) -> Result<()> {
        ctx.table_metadata_manager
            .delete_table_metadata_with_retention_and_generation(
                self.table_id,
                &self.table,
                table_route_value,
                region_wal_options,
                DroppedTableLifecycle {
                    dropped_at,
                    retention_expires_at,
                    drop_generation,
                },
            )
            .await
    }

    /// Deletes the table metadata tombstone **permanently**.
    pub async fn on_delete_metadata_tombstone(
        &self,
        ctx: &DdlContext,
        table_route_value: &TableRouteValue,
        region_wal_options: &HashMap<u32, WalOptions>,
    ) -> Result<()> {
        ctx.table_metadata_manager
            .delete_table_metadata_tombstone(
                self.table_id,
                &self.table,
                table_route_value,
                region_wal_options,
            )
            .await
    }

    /// Deletes metadata for table **permanently**.
    pub async fn on_destroy_metadata(
        &self,
        ctx: &DdlContext,
        table_route_value: &TableRouteValue,
        region_wal_options: &HashMap<u32, WalOptions>,
    ) -> Result<()> {
        ctx.table_metadata_manager
            .destroy_table_metadata(
                self.table_id,
                &self.table,
                table_route_value,
                region_wal_options,
            )
            .await?;

        let detecting_regions = if table_route_value.is_physical() {
            // Safety: checked.
            let regions = table_route_value.region_routes().unwrap();
            convert_region_routes_to_detecting_regions(regions)
        } else {
            vec![]
        };
        ctx.deregister_failure_detectors(detecting_regions).await;
        Ok(())
    }

    /// Restores the table metadata.
    pub async fn on_restore_metadata(
        &self,
        ctx: &DdlContext,
        table_route_value: &TableRouteValue,
        region_wal_options: &HashMap<u32, WalOptions>,
    ) -> Result<()> {
        ctx.table_metadata_manager
            .restore_table_metadata(
                self.table_id,
                &self.table,
                table_route_value,
                region_wal_options,
            )
            .await
    }

    /// Invalidates caches for the table.
    pub async fn invalidate_table_cache(&self, ctx: &DdlContext) -> Result<()> {
        let cache_invalidator = &ctx.cache_invalidator;
        let ctx = Context {
            subject: Some(format!(
                "Invalidate table cache by dropping table {}, table_id: {}",
                self.table.table_ref(),
                self.table_id,
            )),
        };

        cache_invalidator
            .invalidate(
                &ctx,
                &[
                    CacheIdent::TableName(self.table.table_ref().into()),
                    CacheIdent::TableId(self.table_id),
                ],
            )
            .await?;

        Ok(())
    }

    /// Drops regions on datanodes.
    ///
    /// Arguments:
    /// - `node_manager`: resolves datanode clients from peers in `region_routes`.
    /// - `leader_region_registry`: tracks in-flight leader region operations.
    /// - `region_routes`: table region placement; leaders receive drop requests and followers
    ///   receive close requests.
    /// - `fast_path`: forwards to datanode drop requests to skip extra cleanup when safe.
    /// - `force`: forwards to datanode drop requests to allow forced region removal.
    /// - `partial_drop`: forwards to datanode drop requests for partial table/region drops.
    #[allow(clippy::too_many_arguments)]
    pub async fn on_drop_regions(
        &self,
        node_manager: &NodeManagerRef,
        leader_region_registry: &LeaderRegionRegistryRef,
        region_routes: &[RegionRoute],
        fast_path: bool,
        force: bool,
        partial_drop: bool,
    ) -> Result<()> {
        // Drops leader regions on datanodes.
        let leaders = find_leaders(region_routes);
        let mut drop_region_tasks = Vec::with_capacity(leaders.len());
        let table_id = self.table_id;
        for datanode in leaders {
            let requester = node_manager.datanode(&datanode).await;
            let regions = find_leader_regions(region_routes, &datanode);
            let region_ids = regions
                .iter()
                .map(|region_number| RegionId::new(table_id, *region_number))
                .collect::<Vec<_>>();

            for region_id in region_ids {
                debug!("Dropping region {region_id} on Datanode {datanode:?}");
                let request = RegionRequest {
                    header: Some(RegionRequestHeader {
                        tracing_context: TracingContext::from_current_span().to_w3c(),
                        ..Default::default()
                    }),
                    body: Some(region_request::Body::Drop(PbDropRegionRequest {
                        region_id: region_id.as_u64(),
                        fast_path,
                        force,
                        partial_drop,
                        soft_drop: false,
                    })),
                };
                let datanode = datanode.clone();
                let requester = requester.clone();
                drop_region_tasks.push(async move {
                    if let Err(err) = requester.handle(request).await
                        && err.status_code() != StatusCode::RegionNotFound
                    {
                        return Err(add_peer_context_if_needed(datanode)(err));
                    }
                    Ok(())
                });
            }
        }

        join_all(drop_region_tasks)
            .await
            .into_iter()
            .collect::<Result<Vec<_>>>()?;

        // Drops follower regions on datanodes.
        let followers = find_followers(region_routes);
        let mut close_region_tasks = Vec::with_capacity(followers.len());
        for datanode in followers {
            let requester = node_manager.datanode(&datanode).await;
            let regions = find_follower_regions(region_routes, &datanode);
            let region_ids = regions
                .iter()
                .map(|region_number| RegionId::new(table_id, *region_number))
                .collect::<Vec<_>>();

            for region_id in region_ids {
                debug!("Closing region {region_id} on Datanode {datanode:?}");
                let request = RegionRequest {
                    header: Some(RegionRequestHeader {
                        tracing_context: TracingContext::from_current_span().to_w3c(),
                        ..Default::default()
                    }),
                    body: Some(region_request::Body::Close(PbCloseRegionRequest {
                        region_id: region_id.as_u64(),
                        flush_on_close: false,
                    })),
                };

                let datanode = datanode.clone();
                let requester = requester.clone();
                close_region_tasks.push(async move {
                    if let Err(err) = requester.handle(request).await
                        && err.status_code() != StatusCode::RegionNotFound
                    {
                        return Err(add_peer_context_if_needed(datanode)(err));
                    }
                    Ok(())
                });
            }
        }

        // Failure to close follower regions is not critical.
        // When a leader region is dropped, follower regions will be unable to renew their leases via metasrv.
        // Eventually, these follower regions will be automatically closed by the region livekeeper.
        if let Err(err) = join_all(close_region_tasks)
            .await
            .into_iter()
            .collect::<Result<Vec<_>>>()
        {
            error!(err; "Failed to close follower regions on datanodes, table_id: {}", table_id);
        }

        // Deletes the leader region from registry.
        let region_ids = operating_leader_regions(region_routes);
        leader_region_registry.batch_delete(region_ids.into_iter().map(|(region_id, _)| region_id));

        Ok(())
    }

    /// Cleans leader regions on datanodes without reopening them as live regions.
    pub async fn on_cleanup_regions_offline(
        &self,
        node_manager: &NodeManagerRef,
        leader_region_registry: &LeaderRegionRegistryRef,
        table_info: &TableInfo,
        region_routes: &[RegionRoute],
        region_wal_options: &HashMap<RegionNumber, WalOptions>,
    ) -> Result<()> {
        let template = build_template_from_raw_table_info(table_info)?;
        let builder = CreateRequestBuilder::new(template, None);
        let storage_path = region_storage_path(&self.table.catalog_name, &self.table.schema_name);

        let leaders = find_leaders(region_routes);
        let mut cleanup_region_tasks = Vec::with_capacity(leaders.len());
        let table_id = self.table_id;
        for datanode in leaders {
            let requester = node_manager.datanode(&datanode).await;
            let regions = find_leader_regions(region_routes, &datanode);
            let region_ids = regions
                .iter()
                .map(|region_number| RegionId::new(table_id, *region_number))
                .collect::<Vec<_>>();

            for region_id in region_ids {
                debug!("Cleaning region {region_id} offline on Datanode {datanode:?}");
                let create_request = builder.build_one(
                    region_id,
                    storage_path.clone(),
                    region_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::CleanUp(PbCleanUpRequest {
                        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();
                cleanup_region_tasks.push(async move {
                    if let Err(err) = requester.handle(request).await
                        && err.status_code() != StatusCode::RegionNotFound
                    {
                        return Err(add_peer_context_if_needed(datanode)(err));
                    }
                    Ok(())
                });
            }
        }

        join_all(cleanup_region_tasks)
            .await
            .into_iter()
            .collect::<Result<Vec<_>>>()?;

        let region_ids = operating_leader_regions(region_routes);
        leader_region_registry.batch_delete(region_ids.into_iter().map(|(region_id, _)| region_id));

        Ok(())
    }

    /// Closes all table regions on datanodes without deleting region files or metadata tombstones.
    /// When `flush_leaders_on_close` is set, only leader regions are flushed before close.
    #[cfg(feature = "enterprise")]
    pub async fn on_close_regions(
        &self,
        node_manager: &NodeManagerRef,
        leader_region_registry: &LeaderRegionRegistryRef,
        region_routes: &[RegionRoute],
        flush_leaders_on_close: bool,
    ) -> Result<()> {
        let table_id = self.table_id;
        let mut seen_peer_ids = HashSet::new();
        let peers = find_leaders(region_routes)
            .into_iter()
            .chain(find_followers(region_routes))
            .filter(|peer| seen_peer_ids.insert(peer.id));
        let close_region_tasks = peers.map(|datanode| {
            let region_ids = find_leader_regions(region_routes, &datanode)
                .into_iter()
                .map(|region_number| {
                    (
                        RegionId::new(table_id, region_number),
                        flush_leaders_on_close,
                    )
                })
                .chain(
                    find_follower_regions(region_routes, &datanode)
                        .into_iter()
                        .map(|region_number| (RegionId::new(table_id, region_number), false)),
                )
                .collect::<Vec<_>>();

            async move {
                let requester = node_manager.datanode(&datanode).await;
                let close_region_tasks =
                    region_ids.into_iter().map(|(region_id, flush_on_close)| {
                        debug!("Closing region {region_id} on Datanode {datanode:?}");
                        let request = RegionRequest {
                            header: Some(RegionRequestHeader {
                                tracing_context: TracingContext::from_current_span().to_w3c(),
                                ..Default::default()
                            }),
                            body: Some(region_request::Body::Close(PbCloseRegionRequest {
                                region_id: region_id.as_u64(),
                                flush_on_close,
                            })),
                        };

                        let datanode = datanode.clone();
                        let requester = requester.clone();
                        async move {
                            if let Err(err) = requester.handle(request).await
                                && err.status_code() != StatusCode::RegionNotFound
                            {
                                return Err(add_peer_context_if_needed(datanode)(err));
                            }
                            Ok(())
                        }
                    });

                join_all(close_region_tasks)
                    .await
                    .into_iter()
                    .collect::<Result<Vec<_>>>()?;
                Ok(())
            }
        });

        join_all(close_region_tasks)
            .await
            .into_iter()
            .collect::<Result<Vec<_>>>()?;

        let region_ids = operating_leader_regions(region_routes);
        leader_region_registry.batch_delete(region_ids.into_iter().map(|(region_id, _)| region_id));

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use std::assert_matches;
    use std::collections::HashMap;
    use std::sync::Arc;

    use api::v1::{ColumnDataType, SemanticType};
    use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
    use table::metadata::TableInfo;
    use table::table_name::TableName;

    use super::*;
    use crate::ddl::test_util::columns::TestColumnDefBuilder;
    use crate::ddl::test_util::create_table::{
        TestCreateTableExprBuilder, build_raw_table_info_from_expr,
    };
    use crate::key::table_route::TableRouteValue;
    use crate::test_util::{MockDatanodeManager, new_ddl_context};

    fn test_create_raw_table_info(name: &str) -> TableInfo {
        let create_table = TestCreateTableExprBuilder::default()
            .column_defs([
                TestColumnDefBuilder::default()
                    .name("ts")
                    .data_type(ColumnDataType::TimestampMillisecond)
                    .semantic_type(SemanticType::Timestamp)
                    .build()
                    .unwrap()
                    .into(),
                TestColumnDefBuilder::default()
                    .name("host")
                    .data_type(ColumnDataType::String)
                    .semantic_type(SemanticType::Tag)
                    .build()
                    .unwrap()
                    .into(),
                TestColumnDefBuilder::default()
                    .name("cpu")
                    .data_type(ColumnDataType::Float64)
                    .semantic_type(SemanticType::Field)
                    .build()
                    .unwrap()
                    .into(),
            ])
            .time_index("ts")
            .primary_keys(["host".into()])
            .table_name(name)
            .build()
            .unwrap()
            .into();
        build_raw_table_info_from_expr(&create_table)
    }

    #[tokio::test]
    async fn test_on_prepare() {
        // Drops if exists
        let node_manager = Arc::new(MockDatanodeManager::new(()));
        let ctx = new_ddl_context(node_manager);
        let executor = DropTableExecutor::new(
            TableName::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, "my_table"),
            1024,
            true,
        );
        let ctrl = executor.on_prepare(&ctx).await.unwrap();
        assert!(ctrl.stop());

        // Drops a non-exists table
        let executor = DropTableExecutor::new(
            TableName::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, "my_table"),
            1024,
            false,
        );
        let err = executor.on_prepare(&ctx).await.unwrap_err();
        assert_matches!(err, error::Error::TableNotFound { .. });

        // Drops a exists table
        let executor = DropTableExecutor::new(
            TableName::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, "my_table"),
            1024,
            false,
        );
        let raw_table_info = test_create_raw_table_info("my_table");
        ctx.table_metadata_manager
            .create_table_metadata(
                raw_table_info,
                TableRouteValue::physical(vec![]),
                HashMap::new(),
            )
            .await
            .unwrap();
        let ctrl = executor.on_prepare(&ctx).await.unwrap();
        assert!(!ctrl.stop());
    }
}