common_meta/ddl/

alter_table.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.

mod executor;
mod metadata;
mod partition_compat;
mod region_request;

pub use partition_compat::validate_region_partition_bounds;

use std::collections::HashSet;
use std::vec;

use api::region::RegionResponse;
use api::v1::alter_table_expr::Kind;
use api::v1::{RenameTable, SetTableOptions};
use async_trait::async_trait;
use common_catalog::consts::{METRIC_ENGINE, MITO_ENGINE};
use common_error::ext::BoxedError;
use common_procedure::error::{FromJsonSnafu, Result as ProcedureResult, ToJsonSnafu};
use common_procedure::{
    Context as ProcedureContext, ContextProvider, Error as ProcedureError, EventContext,
    EventTrigger, LockKey, PoisonKey, PoisonKeys, Procedure, ProcedureId, Status, StringKey,
};
use common_telemetry::{error, info, warn};
use common_wal::options::WalOptions;
use serde::{Deserialize, Serialize};
use snafu::{ResultExt, ensure};
use store_api::metadata::ColumnMetadata;
use store_api::metric_engine_consts::TABLE_COLUMN_METADATA_EXTENSION_KEY;
use store_api::storage::RegionId;
use strum::AsRefStr;
use table::metadata::{TableId, TableInfo};
use table::requests::SKIP_WAL_KEY;
use table::table_reference::TableReference;

use crate::ddl::DdlContext;
use crate::ddl::alter_table::executor::AlterTableExecutor;
use crate::ddl::event::table::{
    TableDdlEvent, TableDdlEventType, TableDdlLocator, alter_table_kind_name,
};
use crate::ddl::utils::{
    MultipleResults, extract_column_metadatas, extract_region_wal_options, handle_multiple_results,
    map_to_procedure_error, sync_follower_regions,
};
use crate::error::{
    AbortProcedureSnafu, ConvertAlterTableRequestSnafu, NoLeaderSnafu, PutPoisonSnafu, Result,
    RetryLaterSnafu, UnexpectedSnafu, UnsupportedSnafu,
};
use crate::key::table_info::TableInfoValue;
use crate::key::{DeserializedValueWithBytes, RegionDistribution};
use crate::lock_key::{CatalogLock, RegionLock, SchemaLock, TableLock, TableNameLock};
use crate::metrics;
use crate::poison_key::table_poison_key;
use crate::rpc::ddl::AlterTableTask;
use crate::rpc::router::{RegionRoute, find_leaders, region_distribution};

/// The alter table procedure
pub struct AlterTableProcedure {
    /// The runtime context.
    context: DdlContext,
    /// The serialized data.
    data: AlterTableData,
    /// Cached new table metadata in the prepare step.
    /// If we recover the procedure from json, then the table info value is not cached.
    /// But we already validated it in the prepare step.
    new_table_info: Option<TableInfo>,
    /// The alter table executor.
    executor: AlterTableExecutor,
}

#[derive(Debug)]
pub(crate) struct RegionRouteChanged;

/// Builds the executor from the [`AlterTableData`].
///
/// # Panics
/// - If the alter kind is not set.
fn build_executor_from_alter_expr(alter_data: &AlterTableData) -> AlterTableExecutor {
    let table_name = alter_data.table_ref().into();
    let table_id = alter_data.table_id;
    let alter_kind = alter_data.task.alter_table.kind.as_ref().unwrap();
    let new_table_name = if let Kind::RenameTable(RenameTable { new_table_name }) = alter_kind {
        Some(new_table_name.clone())
    } else {
        None
    };
    AlterTableExecutor::new(table_name, table_id, new_table_name)
}

impl AlterTableProcedure {
    pub const TYPE_NAME: &'static str = "metasrv-procedure::AlterTable";

    pub fn new(table_id: TableId, task: AlterTableTask, context: DdlContext) -> Result<Self> {
        Self::new_with_region_locks(table_id, task, vec![], context)
    }

    pub(crate) fn new_with_region_locks(
        table_id: TableId,
        task: AlterTableTask,
        region_locks: Vec<RegionId>,
        context: DdlContext,
    ) -> Result<Self> {
        task.validate()?;
        let data = AlterTableData::new(task, table_id, region_locks);
        let executor = build_executor_from_alter_expr(&data);
        Ok(Self {
            context,
            data,
            new_table_info: None,
            executor,
        })
    }

    pub fn from_json(json: &str, context: DdlContext) -> ProcedureResult<Self> {
        let data: AlterTableData = serde_json::from_str(json).context(FromJsonSnafu)?;
        let executor = build_executor_from_alter_expr(&data);

        Ok(AlterTableProcedure {
            context,
            data,
            new_table_info: None,
            executor,
        })
    }

    // Checks whether the table exists.
    pub(crate) async fn on_prepare(&mut self) -> Result<Status> {
        self.executor
            .on_prepare(&self.context.table_metadata_manager)
            .await?;
        self.fill_table_info().await?;

        // Safety: filled in `fill_table_info`.
        let table_info_value = self.data.table_info_value.as_ref().unwrap();
        let new_table_info = AlterTableExecutor::validate_alter_table_expr(
            &table_info_value.table_info,
            self.data.task.alter_table.clone(),
        )?;
        self.new_table_info = Some(new_table_info);

        // Reject schema changes that the persisted partition expressions of some
        // region cannot survive, before any region is altered.
        self.validate_region_partition_exprs().await?;

        // Safety: Checked in `AlterTableProcedure::new`.
        let alter_kind = self.data.task.alter_table.kind.as_ref().unwrap();
        if sets_skip_wal(alter_kind) {
            ensure!(
                only_sets_skip_wal(alter_kind),
                UnsupportedSnafu {
                    operation: "combining skip_wal with other table options".to_string()
                }
            );
            let engine = table_info_value.table_info.meta.engine.as_str();
            ensure!(
                engine == MITO_ENGINE || engine == METRIC_ENGINE,
                UnsupportedSnafu {
                    operation: format!("setting skip_wal on {engine} engine tables")
                }
            );
            // Keep skip-WAL changes serialized with region migration and use the same
            // route snapshot for validation and RPC delivery.
            let (physical_table_id, physical_table_route) = self
                .context
                .table_metadata_manager
                .table_route_manager()
                .get_physical_table_route(self.data.table_id())
                .await?;
            ensure!(
                physical_table_id == self.data.table_id(),
                UnsupportedSnafu {
                    operation: "setting skip_wal on logical tables".to_string()
                }
            );
            let current_region_ids = physical_table_route
                .region_routes
                .iter()
                .map(|route| route.region.id)
                .collect::<HashSet<_>>();
            let locked_region_ids = self
                .data
                .region_locks
                .iter()
                .copied()
                .collect::<HashSet<_>>();
            if physical_table_route.region_routes.len() != self.data.region_locks.len()
                || current_region_ids != locked_region_ids
            {
                // Procedure lock keys are fixed at submission. Finish this attempt so the
                // DDL manager can submit another procedure with locks for the current route.
                return Ok(Status::done_with_output(RegionRouteChanged));
            }
            ensure!(
                !find_leaders(&physical_table_route.region_routes).is_empty(),
                NoLeaderSnafu {
                    table_id: physical_table_id
                }
            );
            // validate_alter_table_expr() has already parsed this single skip-WAL option.
            let Some(skip_wal) = skip_wal_value(alter_kind) else {
                return UnexpectedSnafu {
                    err_msg: "missing or invalid skip_wal value after validation".to_string(),
                }
                .fail();
            };
            if !skip_wal {
                let datanode_table_values = self
                    .context
                    .table_metadata_manager
                    .datanode_table_manager()
                    .regions(physical_table_id, &physical_table_route)
                    .await?;
                let region_wal_options = extract_region_wal_options(&datanode_table_values)?;
                ensure!(
                    current_region_ids.iter().all(|region_id| {
                        matches!(
                            region_wal_options
                                .get(&region_id.region_number())
                                .cloned()
                                .unwrap_or_default(),
                            WalOptions::RaftEngine
                                | WalOptions::Kafka(_)
                                | WalOptions::ObjectStore(_)
                        )
                    }),
                    UnsupportedSnafu {
                        operation: "setting skip_wal = 'false' without an existing WAL provider"
                            .to_string()
                    }
                );
            }
            self.data.region_distribution =
                Some(region_distribution(&physical_table_route.region_routes));
        }
        self.data.state = self.data.flow()?.after_prepare();
        Ok(Status::executing(true))
    }

    /// Validates the persisted partition expressions of every region against the
    /// schema proposed by this ALTER TABLE, so an ALTER that would break a region's
    /// partition predicate is rejected before any region is modified.
    async fn validate_region_partition_exprs(&self) -> Result<()> {
        // Safety: filled in `fill_table_info`.
        let old_schema = &self
            .data
            .table_info_value
            .as_ref()
            .unwrap()
            .table_info
            .meta
            .schema;
        // Safety: set right after the new table info is built in `on_prepare`.
        let new_schema = &self.new_table_info.as_ref().unwrap().meta.schema;
        let changed_columns = partition_compat::changed_column_types(old_schema, new_schema);
        if changed_columns.is_empty() {
            return Ok(());
        }

        let (physical_table_id, physical_table_route) = self
            .context
            .table_metadata_manager
            .table_route_manager()
            .get_physical_table_route(self.data.table_id())
            .await?;
        if physical_table_id != self.data.table_id() {
            // Partition expressions are persisted on the regions of the physical
            // table; the altered schema here belongs to a logical table.
            return Ok(());
        }

        partition_compat::validate_region_partition_bounds(
            &changed_columns,
            &physical_table_route.region_routes,
        )
    }

    fn table_poison_key(&self) -> PoisonKey {
        table_poison_key(self.data.table_id())
    }

    async fn put_poison(
        &self,
        ctx_provider: &dyn ContextProvider,
        procedure_id: ProcedureId,
    ) -> Result<()> {
        let poison_key = self.table_poison_key();
        ctx_provider
            .try_put_poison(&poison_key, procedure_id)
            .await
            .context(PutPoisonSnafu)
    }

    pub async fn submit_alter_region_requests(
        &mut self,
        procedure_id: ProcedureId,
        ctx_provider: &dyn ContextProvider,
    ) -> Result<Status> {
        let table_id = self.data.table_id();
        let (_, physical_table_route) = self
            .context
            .table_metadata_manager
            .table_route_manager()
            .get_physical_table_route(table_id)
            .await?;

        self.data.region_distribution =
            Some(region_distribution(&physical_table_route.region_routes));
        let leaders = find_leaders(&physical_table_route.region_routes);
        let alter_kind = self.make_region_alter_kind()?;

        info!(
            "Submitting alter region requests for table {}, table_id: {}, alter_kind: {:?}",
            self.data.table_ref(),
            table_id,
            alter_kind,
        );

        ensure!(!leaders.is_empty(), NoLeaderSnafu { table_id });
        // Puts the poison before submitting alter region requests to datanodes.
        self.put_poison(ctx_provider, procedure_id).await?;
        let flow = self.data.flow()?;
        if flow == AlterTableFlow::MetadataFirst {
            let results = self
                .executor
                .on_alter_skip_wal_regions(
                    &self.context.node_manager,
                    &physical_table_route.region_routes,
                    alter_kind,
                )
                .await;

            return match handle_multiple_results(results) {
                MultipleResults::PartialRetryable(error) => Err(error),
                MultipleResults::PartialNonRetryable(error)
                | MultipleResults::AllNonRetryable(error) => {
                    // The metadata already contains the requested skip-WAL value. Retry the
                    // idempotent request so later attempts can update the remaining replicas.
                    Err(BoxedError::new(error)).context(RetryLaterSnafu {
                        clean_poisons: true,
                    })
                }
                MultipleResults::AllRetryable(error) => {
                    Err(BoxedError::new(error)).context(RetryLaterSnafu {
                        clean_poisons: true,
                    })
                }
                MultipleResults::Ok(results) => {
                    self.handle_alter_region_response(results)?;
                    Ok(Status::executing_with_clean_poisons(true))
                }
            };
        }

        let results = self
            .executor
            .on_alter_regions(
                &self.context.node_manager,
                &physical_table_route.region_routes,
                alter_kind,
            )
            .await;

        match handle_multiple_results(results) {
            MultipleResults::PartialRetryable(error) => {
                // Just returns the error, and wait for the next try.
                Err(error)
            }
            MultipleResults::PartialNonRetryable(error) => {
                error!(error; "Partial non-retryable errors occurred during alter table, table {}, table_id: {}", self.data.table_ref(), self.data.table_id());
                // No retry will be done.
                Ok(Status::poisoned(
                    Some(self.table_poison_key()),
                    ProcedureError::external(error),
                ))
            }
            MultipleResults::AllRetryable(error) => {
                // Just returns the error, and wait for the next try.
                let err = BoxedError::new(error);
                Err(err).context(RetryLaterSnafu {
                    clean_poisons: true,
                })
            }
            MultipleResults::Ok(results) => {
                self.submit_sync_region_requests(&results, &physical_table_route.region_routes)
                    .await;
                self.handle_alter_region_response(results)?;
                Ok(Status::executing_with_clean_poisons(true))
            }
            MultipleResults::AllNonRetryable(error) => {
                error!(error; "All alter requests returned non-retryable errors for table {}, table_id: {}", self.data.table_ref(), self.data.table_id());
                // It assumes the metadata on datanode is not changed.
                // Case: The alter region request is sent but not applied. (e.g., InvalidArgument)

                let err = BoxedError::new(error);
                Err(err).context(AbortProcedureSnafu {
                    clean_poisons: true,
                })
            }
        }
    }

    fn handle_alter_region_response(&mut self, mut results: Vec<RegionResponse>) -> Result<()> {
        if let Some(column_metadatas) =
            extract_column_metadatas(&mut results, TABLE_COLUMN_METADATA_EXTENSION_KEY)?
        {
            self.data.column_metadatas = column_metadatas;
        } else {
            warn!(
                "altering table result doesn't contains extension key `{TABLE_COLUMN_METADATA_EXTENSION_KEY}`,leaving the table's column metadata unchanged"
            );
        }
        self.data.state = self.data.flow()?.after_regions();
        Ok(())
    }

    async fn submit_sync_region_requests(
        &mut self,
        results: &[RegionResponse],
        region_routes: &[RegionRoute],
    ) {
        // Safety: filled in `prepare` step.
        let table_info = self.data.table_info().unwrap();
        if let Err(err) = sync_follower_regions(
            &self.context,
            self.data.table_id(),
            results,
            region_routes,
            table_info.meta.engine.as_str(),
        )
        .await
        {
            error!(err; "Failed to sync regions for table {}, table_id: {}", self.data.table_ref(), self.data.table_id());
        }
    }

    /// Update table metadata.
    pub(crate) async fn on_update_metadata(&mut self) -> Result<Status> {
        let table_id = self.data.table_id();
        let table_ref = self.data.table_ref();
        // Safety: filled in `fill_table_info`.
        let table_info_value = self.data.table_info_value.as_ref().unwrap();
        // Safety: Checked in `AlterTableProcedure::new`.
        let alter_kind = self.data.task.alter_table.kind.as_ref().unwrap();
        let flow = self.data.flow()?;
        let metadata_only_alter = flow == AlterTableFlow::MetadataOnly;

        // Gets the table info from the cache or builds it.
        let  new_info = match &self.new_table_info {
            Some(cached) => cached.clone(),
            None => AlterTableExecutor::validate_alter_table_expr(
                &table_info_value.table_info,
                self.data.task.alter_table.clone(),
               )
                .inspect_err(|e| {
                    // We already check the table info in the prepare step so this should not happen.
                    error!(e; "Unable to build info for table {} in update metadata step, table_id: {}", table_ref, table_id);
                })?,
        };

        // Safety: region distribution is set in `submit_alter_region_requests`.
        self.executor
            .on_alter_metadata(
                &self.context.table_metadata_manager,
                table_info_value,
                self.data.region_distribution.as_ref(),
                new_info,
                &self.data.column_metadatas,
                metadata_only_alter,
            )
            .await?;

        info!(
            "Updated table metadata for table {table_ref}, table_id: {table_id}, kind: {alter_kind:?}"
        );
        self.data.state = flow.after_metadata();
        Ok(Status::executing(true))
    }

    /// Broadcasts the invalidating table cache instructions.
    async fn on_broadcast(&mut self) -> Result<Status> {
        self.executor
            .invalidate_table_cache(&self.context.cache_invalidator)
            .await?;
        Ok(Status::done())
    }

    fn lock_key_inner(&self) -> Vec<StringKey> {
        let mut lock_key = vec![];
        let table_ref = self.data.table_ref();
        let table_id = self.data.table_id();
        lock_key.push(CatalogLock::Read(table_ref.catalog).into());
        lock_key.push(SchemaLock::read(table_ref.catalog, table_ref.schema).into());
        lock_key.push(TableLock::Write(table_id).into());

        for region_id in &self.data.region_locks {
            lock_key.push(RegionLock::Write(*region_id).into());
        }

        // Safety: Checked in `AlterTableProcedure::new`.
        let alter_kind = self.data.task.alter_table.kind.as_ref().unwrap();
        if let Kind::RenameTable(RenameTable { new_table_name }) = alter_kind {
            lock_key.push(
                TableNameLock::new(table_ref.catalog, table_ref.schema, new_table_name).into(),
            )
        }

        lock_key
    }

    #[cfg(test)]
    pub(crate) fn data(&self) -> &AlterTableData {
        &self.data
    }

    #[cfg(test)]
    pub(crate) fn mut_data(&mut self) -> &mut AlterTableData {
        &mut self.data
    }
}

fn sets_skip_wal(alter_kind: &Kind) -> bool {
    let Kind::SetTableOptions(SetTableOptions { table_options }) = alter_kind else {
        return false;
    };

    table_options
        .iter()
        .any(|option| option.key == SKIP_WAL_KEY)
}

pub(crate) fn only_sets_skip_wal(alter_kind: &Kind) -> bool {
    let Kind::SetTableOptions(SetTableOptions { table_options }) = alter_kind else {
        return false;
    };

    table_options.len() == 1 && table_options[0].key == SKIP_WAL_KEY
}

fn skip_wal_value(alter_kind: &Kind) -> Option<bool> {
    let Kind::SetTableOptions(SetTableOptions { table_options }) = alter_kind else {
        return None;
    };
    let [option] = table_options.as_slice() else {
        return None;
    };
    if option.key != SKIP_WAL_KEY {
        return None;
    }
    option.value.parse().ok()
}

fn is_metadata_only_alter(alter_kind: &Kind) -> Result<bool> {
    // A mixed annotation batch is an error, never "not metadata-only": falling
    // through to the region-first flow would dispatch it to regions.
    let family = common_grpc_expr::annotation_alter_family(alter_kind)
        .context(ConvertAlterTableRequestSnafu)?;
    Ok(family.is_some() || matches!(alter_kind, Kind::RenameTable { .. }))
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum AlterTableFlow {
    RegionFirst,
    MetadataFirst,
    MetadataOnly,
}

impl AlterTableFlow {
    fn from_kind(kind: &Kind) -> Result<Self> {
        Ok(if only_sets_skip_wal(kind) {
            Self::MetadataFirst
        } else if is_metadata_only_alter(kind)? {
            Self::MetadataOnly
        } else {
            Self::RegionFirst
        })
    }

    fn after_prepare(self) -> AlterTableState {
        match self {
            Self::RegionFirst => AlterTableState::SubmitAlterRegionRequests,
            Self::MetadataFirst | Self::MetadataOnly => AlterTableState::UpdateMetadata,
        }
    }

    fn after_regions(self) -> AlterTableState {
        match self {
            Self::RegionFirst => AlterTableState::UpdateMetadata,
            Self::MetadataFirst | Self::MetadataOnly => AlterTableState::InvalidateTableCache,
        }
    }

    fn after_metadata(self) -> AlterTableState {
        match self {
            Self::MetadataFirst => AlterTableState::SubmitAlterRegionRequests,
            Self::RegionFirst | Self::MetadataOnly => AlterTableState::InvalidateTableCache,
        }
    }
}

#[async_trait]
impl Procedure for AlterTableProcedure {
    fn type_name(&self) -> &str {
        Self::TYPE_NAME
    }

    async fn execute(&mut self, ctx: &ProcedureContext) -> ProcedureResult<Status> {
        let state = &self.data.state;

        let step = state.as_ref();

        let _timer = metrics::METRIC_META_PROCEDURE_ALTER_TABLE
            .with_label_values(&[step])
            .start_timer();

        match state {
            AlterTableState::Prepare => self.on_prepare().await,
            AlterTableState::SubmitAlterRegionRequests => {
                self.submit_alter_region_requests(ctx.procedure_id, ctx.provider.as_ref())
                    .await
            }
            AlterTableState::UpdateMetadata => self.on_update_metadata().await,
            AlterTableState::InvalidateTableCache => self.on_broadcast().await,
        }
        .map_err(map_to_procedure_error)
    }

    fn dump(&self) -> ProcedureResult<String> {
        serde_json::to_string(&self.data).context(ToJsonSnafu)
    }

    fn lock_key(&self) -> LockKey {
        let key = self.lock_key_inner();

        LockKey::new(key)
    }

    fn poison_keys(&self) -> PoisonKeys {
        PoisonKeys::new(vec![self.table_poison_key()])
    }

    fn event(&self, ctx: &EventContext<'_>) -> Option<Box<dyn common_event_recorder::Event>> {
        if !ctx
            .event_type_filter
            .allows(TableDdlEventType::AlterTable.as_str())
        {
            return None;
        }
        let table_ref = self.data.table_ref();
        let locator = TableDdlLocator::new(table_ref.catalog, table_ref.schema, table_ref.table)
            .with_table_id(self.data.table_id());
        let event = match &ctx.trigger {
            EventTrigger::Submitted => {
                let kind = self
                    .data
                    .task
                    .alter_table
                    .kind
                    .as_ref()
                    .and_then(alter_table_kind_name);
                TableDdlEvent::alter_table_submitted(locator, kind)
            }
            _ => TableDdlEvent::lifecycle(TableDdlEventType::AlterTable, [locator]),
        };

        Some(Box::new(event))
    }
}

#[derive(Debug, Serialize, Deserialize, AsRefStr)]
enum AlterTableState {
    /// Prepares to alter the table.
    Prepare,
    /// Sends alter region requests to Datanode.
    SubmitAlterRegionRequests,
    /// Updates table metadata.
    UpdateMetadata,
    /// Broadcasts the invalidating table cache instruction.
    InvalidateTableCache,
}

// The serialized data of alter table.
#[derive(Debug, Serialize, Deserialize)]
pub struct AlterTableData {
    state: AlterTableState,
    task: AlterTableTask,
    table_id: TableId,
    #[serde(default)]
    column_metadatas: Vec<ColumnMetadata>,
    /// Table info value before alteration.
    table_info_value: Option<DeserializedValueWithBytes<TableInfoValue>>,
    /// Region distribution for table in case we need to update region options.
    region_distribution: Option<RegionDistribution>,
    /// Region locks held by skip-WAL alters.
    #[serde(default)]
    region_locks: Vec<RegionId>,
}

impl AlterTableData {
    pub fn new(task: AlterTableTask, table_id: TableId, region_locks: Vec<RegionId>) -> Self {
        Self {
            state: AlterTableState::Prepare,
            task,
            table_id,
            column_metadatas: vec![],
            table_info_value: None,
            region_distribution: None,
            region_locks,
        }
    }

    fn table_ref(&self) -> TableReference<'_> {
        self.task.table_ref()
    }

    fn table_id(&self) -> TableId {
        self.table_id
    }

    fn table_info(&self) -> Option<&TableInfo> {
        self.table_info_value
            .as_ref()
            .map(|value| &value.table_info)
    }

    fn flow(&self) -> Result<AlterTableFlow> {
        // Safety: Checked in `AlterTableProcedure::new`.
        AlterTableFlow::from_kind(self.task.alter_table.kind.as_ref().unwrap())
    }

    #[cfg(test)]
    pub(crate) fn column_metadatas(&self) -> &[ColumnMetadata] {
        &self.column_metadatas
    }

    #[cfg(test)]
    pub(crate) fn set_column_metadatas(&mut self, column_metadatas: Vec<ColumnMetadata>) {
        self.column_metadatas = column_metadatas;
    }
}