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

tests_fuzz/generator/
insert_expr.rs
