tests_fuzz/generator/
select_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 std::marker::PhantomData;
16
17use derive_builder::Builder;
18use rand::seq::{IndexedRandom, SliceRandom};
19use rand::Rng;
20
21use crate::context::TableContextRef;
22use crate::error::{Error, Result};
23use crate::generator::Generator;
24use crate::ir::select_expr::{Direction, SelectExpr};
25
26#[derive(Builder)]
27#[builder(pattern = "owned")]
28pub struct SelectExprGenerator<R: Rng + 'static> {
29    table_ctx: TableContextRef,
30    #[builder(default = "8192")]
31    max_limit: usize,
32    #[builder(default)]
33    _phantom: PhantomData<R>,
34}
35
36impl<R: Rng + 'static> Generator<SelectExpr, R> for SelectExprGenerator<R> {
37    type Error = Error;
38
39    fn generate(&self, rng: &mut R) -> Result<SelectExpr> {
40        let selection = rng.random_range(1..self.table_ctx.columns.len());
41        let mut selected_columns = self
42            .table_ctx
43            .columns
44            .choose_multiple(rng, selection)
45            .cloned()
46            .collect::<Vec<_>>();
47        selected_columns.shuffle(rng);
48
49        let order_by_selection = rng.random_range(1..selection);
50
51        let order_by = selected_columns
52            .choose_multiple(rng, order_by_selection)
53            .map(|c| c.name.to_string())
54            .collect::<Vec<_>>();
55
56        let limit = rng.random_range(1..self.max_limit);
57
58        let direction = if rng.random_bool(1.0 / 2.0) {
59            Direction::Asc
60        } else {
61            Direction::Desc
62        };
63
64        Ok(SelectExpr {
65            table_name: self.table_ctx.name.to_string(),
66            columns: selected_columns,
67            order_by,
68            direction,
69            limit,
70        })
71    }
72}