Skip to main content

metric_engine/engine/alter/
extract_new_columns.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::{HashMap, HashSet};
16
17use api::v1::SemanticType;
18use snafu::ensure;
19use store_api::metadata::ColumnMetadata;
20use store_api::region_request::{AlterKind, RegionAlterRequest};
21use store_api::storage::RegionId;
22
23use crate::error::{AddingFieldColumnSnafu, Result};
24
25/// Extract new columns from the create requests.
26///
27/// # Panics
28///
29/// This function will panic if the alter kind is not `AddColumns`.
30pub fn extract_new_columns<'a>(
31    requests: &'a [(RegionId, RegionAlterRequest)],
32    physical_columns: &HashMap<String, ColumnMetadata>,
33    new_column_names: &mut HashSet<&'a str>,
34    new_columns: &mut Vec<ColumnMetadata>,
35) -> Result<()> {
36    for (_, request) in requests {
37        let AlterKind::AddColumns { columns } = &request.kind else {
38            unreachable!()
39        };
40        for col in columns {
41            let column_name = col.column_metadata.column_schema.name.as_str();
42            if !physical_columns.contains_key(column_name)
43                && !new_column_names.contains(column_name)
44            {
45                ensure!(
46                    col.column_metadata.semantic_type != SemanticType::Field,
47                    AddingFieldColumnSnafu {
48                        name: column_name.to_string(),
49                    }
50                );
51                new_column_names.insert(column_name);
52                // TODO(weny): avoid clone
53                new_columns.push(col.column_metadata.clone());
54            }
55        }
56    }
57
58    Ok(())
59}
60
61#[cfg(test)]
62mod tests {
63    use std::collections::{HashMap, HashSet};
64
65    use api::v1::SemanticType;
66    use datatypes::prelude::ConcreteDataType;
67    use datatypes::schema::ColumnSchema;
68    use store_api::metadata::ColumnMetadata;
69    use store_api::region_request::{AddColumn, AlterKind, RegionAlterRequest};
70    use store_api::storage::RegionId;
71
72    use super::*;
73    use crate::error::Error;
74
75    #[test]
76    fn test_extract_new_columns_with_field_type() {
77        let requests = vec![(
78            RegionId::new(1, 1),
79            RegionAlterRequest {
80                kind: AlterKind::AddColumns {
81                    columns: vec![AddColumn {
82                        column_metadata: ColumnMetadata {
83                            column_schema: ColumnSchema::new(
84                                "new_column".to_string(),
85                                ConcreteDataType::string_datatype(),
86                                false,
87                            ),
88                            semantic_type: SemanticType::Field,
89                            column_id: 0,
90                        },
91                        location: None,
92                    }],
93                },
94            },
95        )];
96
97        let physical_columns = HashMap::new();
98        let mut new_column_names = HashSet::new();
99        let mut new_columns = Vec::new();
100
101        let err = extract_new_columns(
102            &requests,
103            &physical_columns,
104            &mut new_column_names,
105            &mut new_columns,
106        )
107        .unwrap_err();
108
109        assert!(matches!(err, Error::AddingFieldColumn { .. }));
110    }
111}