tests_fuzz/generator/
insert_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 datatypes::value::Value;
18use derive_builder::Builder;
19use rand::seq::{IndexedRandom, SliceRandom};
20use rand::Rng;
21
22use super::TsValueGenerator;
23use crate::context::TableContextRef;
24use crate::error::{Error, Result};
25use crate::fake::WordGenerator;
26use crate::generator::{Generator, Random, ValueGenerator};
27use crate::ir::insert_expr::{InsertIntoExpr, RowValue};
28use crate::ir::{generate_random_timestamp, generate_random_value, Ident};
29
30/// Generates [InsertIntoExpr].
31#[derive(Builder)]
32#[builder(pattern = "owned")]
33pub struct InsertExprGenerator<R: Rng + 'static> {
34    table_ctx: TableContextRef,
35    // Whether to omit all columns, i.e. INSERT INTO table_name VALUES (...)
36    omit_column_list: bool,
37    #[builder(default = "1")]
38    rows: usize,
39    #[builder(default = "Box::new(WordGenerator)")]
40    word_generator: Box<dyn Random<Ident, R>>,
41    #[builder(default = "Box::new(generate_random_value)")]
42    value_generator: ValueGenerator<R>,
43    #[builder(default = "Box::new(generate_random_timestamp)")]
44    ts_value_generator: TsValueGenerator<R>,
45    #[builder(default)]
46    _phantom: PhantomData<R>,
47}
48
49impl<R: Rng + 'static> Generator<InsertIntoExpr, R> for InsertExprGenerator<R> {
50    type Error = Error;
51
52    /// Generates the [InsertIntoExpr].
53    fn generate(&self, rng: &mut R) -> Result<InsertIntoExpr> {
54        let mut values_columns = vec![];
55        if self.omit_column_list {
56            // If omit column list, then all columns are required in the values list
57            values_columns.clone_from(&self.table_ctx.columns);
58        } else {
59            for column in &self.table_ctx.columns {
60                let can_omit = column.is_nullable() || column.has_default_value();
61
62                // 50% chance to omit a column if it's not required
63                if !can_omit || rng.random_bool(0.5) {
64                    values_columns.push(column.clone());
65                }
66            }
67            values_columns.shuffle(rng);
68
69            // If all columns are omitted, pick a random column
70            if values_columns.is_empty() {
71                values_columns.push(self.table_ctx.columns.choose(rng).unwrap().clone());
72            }
73        }
74
75        let mut values_list = Vec::with_capacity(self.rows);
76        for _ in 0..self.rows {
77            let mut row = Vec::with_capacity(values_columns.len());
78            for column in &values_columns {
79                if column.is_nullable() && rng.random_bool(0.2) {
80                    row.push(RowValue::Value(Value::Null));
81                    continue;
82                }
83
84                if column.has_default_value() && rng.random_bool(0.2) {
85                    row.push(RowValue::Default);
86                    continue;
87                }
88                if column.is_time_index() {
89                    row.push(RowValue::Value((self.ts_value_generator)(
90                        rng,
91                        column.timestamp_type().unwrap(),
92                    )));
93                } else {
94                    row.push(RowValue::Value((self.value_generator)(
95                        rng,
96                        &column.column_type,
97                        Some(self.word_generator.as_ref()),
98                    )));
99                }
100            }
101
102            values_list.push(row);
103        }
104
105        Ok(InsertIntoExpr {
106            table_name: self.table_ctx.name.clone(),
107            omit_column_list: self.omit_column_list,
108            columns: values_columns,
109            values_list,
110        })
111    }
112}
tests_fuzz/generator/insert_expr.rs

tests_fuzz/generator/
insert_expr.rs
