operator/req_convert/insert/
row_to_region.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 ahash::{HashMap, HashSet};
16use api::v1::RowInsertRequests;
17use api::v1::region::InsertRequests as RegionInsertRequests;
18use partition::manager::PartitionRuleManager;
19use snafu::OptionExt;
20use table::metadata::{TableId, TableInfoRef};
21
22use crate::error::{Result, TableNotFoundSnafu};
23use crate::insert::InstantAndNormalInsertRequests;
24use crate::req_convert::common::partitioner::Partitioner;
25
26pub struct RowToRegion<'a> {
27    tables_info: HashMap<String, TableInfoRef>,
28    instant_table_ids: HashSet<TableId>,
29    partition_manager: &'a PartitionRuleManager,
30}
31
32impl<'a> RowToRegion<'a> {
33    pub fn new(
34        tables_info: HashMap<String, TableInfoRef>,
35        instant_table_ids: HashSet<TableId>,
36        partition_manager: &'a PartitionRuleManager,
37    ) -> Self {
38        Self {
39            tables_info,
40            instant_table_ids,
41            partition_manager,
42        }
43    }
44
45    pub async fn convert(
46        &self,
47        requests: RowInsertRequests,
48        skip_wal: bool,
49    ) -> Result<InstantAndNormalInsertRequests> {
50        let mut region_request = Vec::with_capacity(requests.inserts.len());
51        let mut instant_request = Vec::with_capacity(requests.inserts.len());
52        for request in requests.inserts {
53            let Some(rows) = request.rows else { continue };
54
55            let table_info = self.get_table_info(&request.table_name)?;
56            let table_id = table_info.table_id();
57
58            let requests = Partitioner::new(self.partition_manager)
59                .partition_insert_requests(table_info, rows, skip_wal)
60                .await?;
61
62            if self.instant_table_ids.contains(&table_id) {
63                instant_request.extend(requests);
64            } else {
65                region_request.extend(requests);
66            }
67        }
68
69        Ok(InstantAndNormalInsertRequests {
70            normal_requests: RegionInsertRequests {
71                requests: region_request,
72            },
73            instant_requests: RegionInsertRequests {
74                requests: instant_request,
75            },
76        })
77    }
78
79    fn get_table_info(&self, table_name: &str) -> Result<&TableInfoRef> {
80        self.tables_info
81            .get(table_name)
82            .context(TableNotFoundSnafu { table_name })
83    }
84}
85
86#[cfg(test)]
87mod tests {
88    use std::sync::Arc;
89
90    use api::v1::helper::tag_column_schema;
91    use api::v1::value::ValueData;
92    use api::v1::{ColumnDataType, Row, RowInsertRequest, Rows, Value};
93
94    use super::*;
95    use crate::test_util::{
96        create_partition_rule_manager, new_test_table_info, prepare_mocked_backend,
97    };
98
99    #[tokio::test]
100    async fn test_partitioned_insert_skip_wal_normal_and_instant() {
101        check_partitioned_insert_skip_wal(false, false).await;
102        check_partitioned_insert_skip_wal(false, true).await;
103        check_partitioned_insert_skip_wal(true, false).await;
104        check_partitioned_insert_skip_wal(true, true).await;
105    }
106
107    async fn check_partitioned_insert_skip_wal(instant: bool, skip_wal: bool) {
108        let backend = prepare_mocked_backend().await;
109        let partition_manager = create_partition_rule_manager(backend).await;
110        let table_info = Arc::new(new_test_table_info(1, "table_1", [1, 2, 3].into_iter()));
111        let instant_table_ids = if instant {
112            HashSet::from_iter([1])
113        } else {
114            HashSet::default()
115        };
116        let converter = RowToRegion::new(
117            HashMap::from_iter([("table_1".to_string(), table_info)]),
118            instant_table_ids,
119            &partition_manager,
120        );
121        let requests = RowInsertRequests {
122            inserts: vec![RowInsertRequest {
123                table_name: "table_1".to_string(),
124                rows: Some(Rows {
125                    schema: vec![tag_column_schema("a", ColumnDataType::Int32)],
126                    rows: [1, 11, 101]
127                        .into_iter()
128                        .map(|value| Row {
129                            values: vec![Value {
130                                value_data: Some(ValueData::I32Value(value)),
131                            }],
132                        })
133                        .collect(),
134                }),
135            }],
136        };
137        let result = converter.convert(requests, skip_wal).await.unwrap();
138        let (selected, other) = if instant {
139            (result.instant_requests, result.normal_requests)
140        } else {
141            (result.normal_requests, result.instant_requests)
142        };
143        assert!(other.requests.is_empty());
144        assert_eq!(selected.requests.len(), 3);
145        assert!(
146            selected
147                .requests
148                .iter()
149                .all(|request| request.skip_wal == skip_wal)
150        );
151        assert_eq!(
152            selected
153                .requests
154                .iter()
155                .map(|request| request.rows.as_ref().unwrap().rows.len())
156                .sum::<usize>(),
157            3
158        );
159    }
160}
operator/req_convert/insert/row_to_region.rs

operator/req_convert/insert/
row_to_region.rs
