metric_engine/engine/

alter.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 extract_new_columns;
mod validate;

use std::collections::{BTreeSet, HashMap, HashSet};

use api::v1::SemanticType;
use common_query::native_histogram::is_native_histogram_value_type;
use extract_new_columns::extract_new_columns;
use snafu::{OptionExt, ResultExt, ensure};
use store_api::metadata::ColumnMetadata;
use store_api::metric_engine_consts::ALTER_PHYSICAL_EXTENSION_KEY;
use store_api::region_request::{AffectedRows, AlterKind, RegionAlterRequest};
use store_api::storage::RegionId;
use validate::validate_alter_region_requests;

use crate::engine::MetricEngineInner;
use crate::error::{
    AddingFieldColumnSnafu, LogicalRegionNotFoundSnafu, PhysicalRegionNotFoundSnafu, Result,
    SerializeColumnMetadataSnafu, UnexpectedRequestSnafu,
};
use crate::utils::{append_manifest_info, encode_manifest_info_to_extensions, to_data_region_id};

impl MetricEngineInner {
    pub async fn alter_regions(
        &self,
        mut requests: Vec<(RegionId, RegionAlterRequest)>,
        extension_return_value: &mut HashMap<String, Vec<u8>>,
    ) -> Result<AffectedRows> {
        if requests.is_empty() {
            return Ok(0);
        }

        let first_region_id = &requests.first().unwrap().0;
        if self.is_physical_region(*first_region_id) {
            ensure!(
                requests.len() == 1,
                UnexpectedRequestSnafu {
                    reason: "Physical table must be altered with single request".to_string(),
                }
            );
            let (region_id, request) = requests.pop().unwrap();
            self.alter_physical_region(region_id, request).await?;
        } else {
            // Fast path for single logical region alter request
            if requests.len() == 1 {
                // Safety: requests is not empty
                let region_id = requests.first().unwrap().0;
                let physical_region_id = self
                    .state
                    .read()
                    .unwrap()
                    .get_physical_region_id(region_id)
                    .with_context(|| LogicalRegionNotFoundSnafu { region_id })?;
                let mut manifest_infos = Vec::with_capacity(1);
                self.alter_logical_regions(physical_region_id, requests, extension_return_value)
                    .await?;
                append_manifest_info(&self.mito, physical_region_id, &mut manifest_infos);
                encode_manifest_info_to_extensions(&manifest_infos, extension_return_value)?;
            } else {
                let grouped_requests =
                    self.group_logical_region_requests_by_physical_region_id(requests)?;
                let mut manifest_infos = Vec::with_capacity(grouped_requests.len());
                for (physical_region_id, requests) in grouped_requests {
                    self.alter_logical_regions(
                        physical_region_id,
                        requests,
                        extension_return_value,
                    )
                    .await?;
                    append_manifest_info(&self.mito, physical_region_id, &mut manifest_infos);
                }
                encode_manifest_info_to_extensions(&manifest_infos, extension_return_value)?;
            }
        }
        Ok(0)
    }

    /// Groups the alter logical region requests by physical region id.
    fn group_logical_region_requests_by_physical_region_id(
        &self,
        requests: Vec<(RegionId, RegionAlterRequest)>,
    ) -> Result<HashMap<RegionId, Vec<(RegionId, RegionAlterRequest)>>> {
        let mut result = HashMap::with_capacity(requests.len());
        let state = self.state.read().unwrap();

        for (region_id, request) in requests {
            let physical_region_id = state
                .get_physical_region_id(region_id)
                .with_context(|| LogicalRegionNotFoundSnafu { region_id })?;
            result
                .entry(physical_region_id)
                .or_insert_with(Vec::new)
                .push((region_id, request));
        }

        Ok(result)
    }

    /// Alter multiple logical regions on the same physical region.
    pub async fn alter_logical_regions(
        &self,
        physical_region_id: RegionId,
        requests: Vec<(RegionId, RegionAlterRequest)>,
        extension_return_value: &mut HashMap<String, Vec<u8>>,
    ) -> Result<AffectedRows> {
        // Checks all alter requests are add columns.
        validate_alter_region_requests(&requests)?;
        self.validate_logical_field_alters(physical_region_id, &requests)
            .await?;

        // Finds new columns to add
        let mut new_column_names = HashSet::new();
        let mut new_columns_to_add = vec![];

        let index_options = {
            let state = &self.state.read().unwrap();
            let region_state = state
                .physical_region_states()
                .get(&physical_region_id)
                .with_context(|| PhysicalRegionNotFoundSnafu {
                    region_id: physical_region_id,
                })?;
            let physical_columns = region_state.physical_columns();

            extract_new_columns(
                &requests,
                physical_columns,
                &mut new_column_names,
                &mut new_columns_to_add,
            )?;

            region_state.options().index
        };
        let data_region_id = to_data_region_id(physical_region_id);

        // Acquire logical region locks in a deterministic order to avoid deadlocks when multiple
        // alter operations target overlapping regions concurrently.
        let region_ids = requests
            .iter()
            .map(|(region_id, _)| *region_id)
            .collect::<BTreeSet<_>>();

        let mut write_guards = Vec::with_capacity(region_ids.len());
        for region_id in region_ids {
            write_guards.push(
                self.metadata_region
                    .write_lock_logical_region(region_id)
                    .await?,
            );
        }

        self.data_region
            .add_columns(data_region_id, new_columns_to_add, index_options)
            .await?;

        let physical_columns = self.data_region.physical_columns(data_region_id).await?;
        let physical_schema_map = physical_columns
            .iter()
            .map(|metadata| (metadata.column_schema.name.as_str(), metadata))
            .collect::<HashMap<_, _>>();

        let logical_region_columns = requests.iter().map(|(region_id, request)| {
            let AlterKind::AddColumns { columns } = &request.kind else {
                unreachable!()
            };
            (
                *region_id,
                columns
                    .iter()
                    .map(|col| {
                        let column_name = col.column_metadata.column_schema.name.as_str();
                        let column_metadata = *physical_schema_map.get(column_name).unwrap();
                        (column_name, column_metadata)
                    })
                    .collect::<HashMap<_, _>>(),
            )
        });

        let new_add_columns = new_column_names.iter().map(|name| {
            // Safety: previous steps ensure the physical region exist
            let column_metadata = *physical_schema_map.get(name).unwrap();
            (name.to_string(), column_metadata.clone())
        });

        // Writes logical regions metadata to metadata region
        self.metadata_region
            .add_logical_regions(physical_region_id, false, logical_region_columns)
            .await?;

        extension_return_value.insert(
            ALTER_PHYSICAL_EXTENSION_KEY.to_string(),
            ColumnMetadata::encode_list(&physical_columns).context(SerializeColumnMetadataSnafu)?,
        );

        let mut state = self.state.write().unwrap();
        state.add_physical_columns(data_region_id, new_add_columns);
        state.invalid_logical_regions_cache(requests.iter().map(|(region_id, _)| *region_id));

        Ok(0)
    }

    async fn validate_logical_field_alters(
        &self,
        physical_region_id: RegionId,
        requests: &[(RegionId, RegionAlterRequest)],
    ) -> Result<()> {
        // Logical metric tables have one field column. Native histograms are a
        // special struct field, so field alters must leave exactly that field.
        for (region_id, request) in requests {
            let AlterKind::AddColumns { columns } = &request.kind else {
                unreachable!()
            };
            let added_fields = columns
                .iter()
                .filter(|col| col.column_metadata.semantic_type == SemanticType::Field)
                .collect::<Vec<_>>();
            let Some(&first_added_field) = added_fields.first() else {
                continue;
            };

            let mut fields = self
                .load_logical_columns(physical_region_id, *region_id)
                .await?
                .into_iter()
                .filter(|col| col.semantic_type == SemanticType::Field)
                .collect::<Vec<_>>();
            fields.extend(
                added_fields
                    .into_iter()
                    .map(|col| col.column_metadata.clone()),
            );

            ensure!(
                fields.len() == 1
                    && is_native_histogram_value_type(&fields[0].column_schema.data_type),
                AddingFieldColumnSnafu {
                    name: first_added_field.column_metadata.column_schema.name.clone(),
                }
            );
        }

        Ok(())
    }

    async fn alter_physical_region(
        &self,
        region_id: RegionId,
        request: RegionAlterRequest,
    ) -> Result<()> {
        self.data_region
            .alter_region_options(region_id, request)
            .await?;
        Ok(())
    }
}

#[cfg(test)]
mod test {
    use api::v1::SemanticType;
    use common_meta::ddl::test_util::assert_column_name_and_id;
    use common_meta::ddl::utils::{parse_column_metadatas, parse_manifest_infos_from_extensions};
    use common_query::prelude::{greptime_timestamp, greptime_value};
    use store_api::metric_engine_consts::ALTER_PHYSICAL_EXTENSION_KEY;
    use store_api::region_engine::RegionEngine;
    use store_api::region_request::{
        AlterKind, BatchRegionDdlRequest, RegionAlterRequest, SetRegionOption,
    };
    use store_api::storage::RegionId;
    use store_api::storage::consts::ReservedColumnId;

    use crate::test_util::{TestEnv, alter_logical_region_request, create_logical_region_request};

    #[tokio::test]
    async fn test_alter_region() {
        let env = TestEnv::new().await;
        env.init_metric_region().await;
        let engine = env.metric();
        let engine_inner = engine.inner;

        // alter physical region
        let physical_region_id = env.default_physical_region_id();
        let request = alter_logical_region_request(&["tag1"]);

        let result = engine_inner
            .alter_physical_region(physical_region_id, request.clone())
            .await;
        assert!(result.is_err());
        assert_eq!(
            result.unwrap_err().to_string(),
            "Alter request to physical region is forbidden".to_string()
        );

        // skip WAL on the physical region should be forwarded to the data region
        let alter_region_option_request = RegionAlterRequest {
            kind: AlterKind::SetRegionOptions {
                options: vec![SetRegionOption::SkipWal],
            },
        };
        engine_inner
            .alter_physical_region(physical_region_id, alter_region_option_request)
            .await
            .unwrap();

        // alter logical region
        let metadata_region = env.metadata_region();
        let logical_region_id = env.default_logical_region_id();
        let is_column_exist = metadata_region
            .column_semantic_type(physical_region_id, logical_region_id, "tag1")
            .await
            .unwrap()
            .is_some();
        assert!(!is_column_exist);

        let region_id = env.default_logical_region_id();
        let response = env
            .metric()
            .handle_batch_ddl_requests(BatchRegionDdlRequest::Alter(vec![(
                region_id,
                request.clone(),
            )]))
            .await
            .unwrap();
        let manifest_infos = parse_manifest_infos_from_extensions(&response.extensions).unwrap();
        assert_eq!(manifest_infos[0].0, physical_region_id);
        assert!(manifest_infos[0].1.is_metric());

        let semantic_type = metadata_region
            .column_semantic_type(physical_region_id, logical_region_id, "tag1")
            .await
            .unwrap()
            .unwrap();
        assert_eq!(semantic_type, SemanticType::Tag);
        let timestamp_index = metadata_region
            .column_semantic_type(physical_region_id, logical_region_id, greptime_timestamp())
            .await
            .unwrap()
            .unwrap();
        assert_eq!(timestamp_index, SemanticType::Timestamp);
        let column_metadatas =
            parse_column_metadatas(&response.extensions, ALTER_PHYSICAL_EXTENSION_KEY).unwrap();
        assert_column_name_and_id(
            &column_metadatas,
            &[
                (greptime_timestamp(), 0),
                (greptime_value(), 1),
                ("__table_id", ReservedColumnId::table_id()),
                ("__tsid", ReservedColumnId::tsid()),
                ("job", 2),
                ("tag1", 3),
            ],
        );
    }

    #[tokio::test]
    async fn test_alter_logical_regions() {
        let env = TestEnv::new().await;
        let engine = env.metric();
        let physical_region_id1 = RegionId::new(1024, 0);
        let physical_region_id2 = RegionId::new(1024, 1);
        let logical_region_id1 = RegionId::new(1025, 0);
        let logical_region_id2 = RegionId::new(1025, 1);
        env.create_physical_region(physical_region_id1, "/test_dir1", vec![])
            .await;
        env.create_physical_region(physical_region_id2, "/test_dir2", vec![])
            .await;

        let region_create_request1 = crate::test_util::create_logical_region_request(
            &["job"],
            physical_region_id1,
            "logical1",
        );
        let region_create_request2 =
            create_logical_region_request(&["job"], physical_region_id2, "logical2");
        engine
            .handle_batch_ddl_requests(BatchRegionDdlRequest::Create(vec![
                (logical_region_id1, region_create_request1),
                (logical_region_id2, region_create_request2),
            ]))
            .await
            .unwrap();

        let region_alter_request1 = alter_logical_region_request(&["tag1"]);
        let region_alter_request2 = alter_logical_region_request(&["tag1"]);
        let response = engine
            .handle_batch_ddl_requests(BatchRegionDdlRequest::Alter(vec![
                (logical_region_id1, region_alter_request1),
                (logical_region_id2, region_alter_request2),
            ]))
            .await
            .unwrap();

        let manifest_infos = parse_manifest_infos_from_extensions(&response.extensions).unwrap();
        assert_eq!(manifest_infos.len(), 2);
        let region_ids = manifest_infos.into_iter().map(|i| i.0).collect::<Vec<_>>();
        assert!(region_ids.contains(&physical_region_id1));
        assert!(region_ids.contains(&physical_region_id2));

        let column_metadatas =
            parse_column_metadatas(&response.extensions, ALTER_PHYSICAL_EXTENSION_KEY).unwrap();
        assert_column_name_and_id(
            &column_metadatas,
            &[
                (greptime_timestamp(), 0),
                (greptime_value(), 1),
                ("__table_id", ReservedColumnId::table_id()),
                ("__tsid", ReservedColumnId::tsid()),
                ("job", 2),
                ("tag1", 3),
            ],
        );
    }
}