tests_fuzz/generator/
repartition_expr.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 derive_builder::Builder;
16use rand::Rng;
17
18use crate::context::TableContextRef;
19use crate::error::{Error, Result};
20use crate::generator::Generator;
21use crate::ir::partition_expr::{SimplePartitions, generate_unique_bound};
22use crate::ir::repartition_expr::{MergePartitionExpr, SplitPartitionExpr};
23
24#[derive(Builder)]
25#[builder(pattern = "owned")]
26pub struct SplitPartitionExprGenerator {
27    table_ctx: TableContextRef,
28    #[builder(default = "true")]
29    wait: bool,
30}
31
32impl<R: Rng + 'static> Generator<SplitPartitionExpr, R> for SplitPartitionExprGenerator {
33    type Error = Error;
34
35    fn generate(&self, rng: &mut R) -> Result<SplitPartitionExpr> {
36        let table_name = self.table_ctx.name.clone();
37        let partition_def = self
38            .table_ctx
39            .partition
40            .as_ref()
41            .expect("expected partition def");
42        let mut partitions =
43            SimplePartitions::from_exprs(partition_def.columns[0].clone(), &partition_def.exprs)?;
44        let column = self
45            .table_ctx
46            .columns
47            .iter()
48            .find(|c| c.name.value == partitions.column_name.value)
49            .unwrap_or_else(|| {
50                panic!(
51                    "partition column not found: {}, columns: {:?}",
52                    partitions.column_name.value,
53                    self.table_ctx
54                        .columns
55                        .iter()
56                        .map(|c| &c.name.value)
57                        .collect::<Vec<_>>()
58                )
59            });
60        let new_bound = generate_unique_bound(rng, &column.column_type, &partitions.bounds)?;
61        let insert_pos = partitions.insert_bound(new_bound)?;
62
63        let from_expr = partition_def.exprs[insert_pos].clone();
64        let new_exprs = partitions.generate()?;
65        let left = new_exprs[insert_pos].clone();
66        let right = new_exprs[insert_pos + 1].clone();
67        Ok(SplitPartitionExpr {
68            table_name,
69            target: from_expr,
70            into: vec![left, right],
71            wait: self.wait,
72        })
73    }
74}
75
76#[derive(Builder)]
77#[builder(pattern = "owned")]
78pub struct MergePartitionExprGenerator {
79    table_ctx: TableContextRef,
80    #[builder(default = "true")]
81    wait: bool,
82}
83
84impl<R: Rng + 'static> Generator<MergePartitionExpr, R> for MergePartitionExprGenerator {
85    type Error = Error;
86
87    fn generate(&self, rng: &mut R) -> Result<MergePartitionExpr> {
88        let table_name = self.table_ctx.name.clone();
89        let partition_def = self
90            .table_ctx
91            .partition
92            .as_ref()
93            .expect("expected partition def");
94        let mut partitions =
95            SimplePartitions::from_exprs(partition_def.columns[0].clone(), &partition_def.exprs)?;
96        let remove_idx = rng.random_range(0..partitions.bounds.len());
97        partitions.remove_bound(remove_idx)?;
98        let left = partition_def.exprs[remove_idx].clone();
99        let right = partition_def.exprs[remove_idx + 1].clone();
100
101        Ok(MergePartitionExpr {
102            table_name,
103            targets: vec![left, right],
104            wait: self.wait,
105        })
106    }
107}
108
109#[cfg(test)]
110mod tests {
111    use std::sync::Arc;
112
113    use rand::{Rng, SeedableRng};
114
115    use crate::context::TableContext;
116    use crate::generator::Generator;
117    use crate::generator::create_expr::CreateTableExprGeneratorBuilder;
118    use crate::generator::repartition_expr::{
119        MergePartitionExprGeneratorBuilder, SplitPartitionExprGeneratorBuilder,
120    };
121    use crate::ir::repartition_expr::RepartitionExpr;
122    use crate::translator::DslTranslator;
123    use crate::translator::mysql::repartition_expr::RepartitionExprTranslator;
124
125    #[test]
126    fn test_split_partition_expr() {
127        common_telemetry::init_default_ut_logging();
128        let mut rng = rand::rng();
129        let seed = rng.random::<u64>();
130        common_telemetry::info!("initializing rng with seed: {seed}");
131        let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(seed);
132        let expr = CreateTableExprGeneratorBuilder::default()
133            .columns(10)
134            .partition(10)
135            .if_not_exists(true)
136            .engine("mito2")
137            .build()
138            .unwrap()
139            .generate(&mut rng)
140            .unwrap();
141        let table_ctx = Arc::new(TableContext::from(&expr));
142        SplitPartitionExprGeneratorBuilder::default()
143            .table_ctx(table_ctx)
144            .build()
145            .unwrap()
146            .generate(&mut rng)
147            .unwrap();
148    }
149
150    #[test]
151    fn test_split_partition_expr_deterministic() {
152        let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(1234);
153        let expr = CreateTableExprGeneratorBuilder::default()
154            .columns(10)
155            .partition(10)
156            .if_not_exists(true)
157            .engine("mito2")
158            .build()
159            .unwrap()
160            .generate(&mut rng)
161            .unwrap();
162        let table_ctx = Arc::new(TableContext::from(&expr));
163        let expr = SplitPartitionExprGeneratorBuilder::default()
164            .table_ctx(table_ctx)
165            .build()
166            .unwrap()
167            .generate(&mut rng)
168            .unwrap();
169
170        let sql = RepartitionExprTranslator
171            .translate(&RepartitionExpr::Split(expr))
172            .unwrap();
173        let expected = r#"ALTER TABLE quO SPLIT PARTITION (
174  MaGnI >= 1844674407370955160 AND MaGnI < 2767011611056432740
175) INTO (
176  MaGnI >= 1844674407370955160 AND MaGnI < 2290232243991136014,
177  MaGnI >= 2290232243991136014 AND MaGnI < 2767011611056432740
178);"#;
179        assert_eq!(expected, sql);
180    }
181
182    #[test]
183    fn test_merge_partition_expr() {
184        common_telemetry::init_default_ut_logging();
185        let mut rng = rand::rng();
186        let seed = rng.random::<u64>();
187        common_telemetry::info!("initializing rng with seed: {seed}");
188        let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(seed);
189        let expr = CreateTableExprGeneratorBuilder::default()
190            .columns(10)
191            .partition(10)
192            .if_not_exists(true)
193            .engine("mito2")
194            .build()
195            .unwrap()
196            .generate(&mut rng)
197            .unwrap();
198        let table_ctx = Arc::new(TableContext::from(&expr));
199        MergePartitionExprGeneratorBuilder::default()
200            .table_ctx(table_ctx)
201            .build()
202            .unwrap()
203            .generate(&mut rng)
204            .unwrap();
205    }
206
207    #[test]
208    fn test_merge_partition_expr_deterministic() {
209        let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(1234);
210        let expr = CreateTableExprGeneratorBuilder::default()
211            .columns(10)
212            .partition(10)
213            .if_not_exists(true)
214            .engine("mito2")
215            .build()
216            .unwrap()
217            .generate(&mut rng)
218            .unwrap();
219        let table_ctx = Arc::new(TableContext::from(&expr));
220        let expr = MergePartitionExprGeneratorBuilder::default()
221            .table_ctx(table_ctx)
222            .build()
223            .unwrap()
224            .generate(&mut rng)
225            .unwrap();
226
227        let sql = RepartitionExprTranslator
228            .translate(&RepartitionExpr::Merge(expr))
229            .unwrap();
230        let expected = r#"ALTER TABLE quO MERGE PARTITION (
231  MaGnI >= 3689348814741910320 AND MaGnI < 4611686018427387900,
232  MaGnI >= 4611686018427387900 AND MaGnI < 5534023222112865480
233);"#;
234        assert_eq!(expected, sql);
235    }
236}
tests_fuzz/generator/repartition_expr.rs

tests_fuzz/generator/
repartition_expr.rs
