mirror of
https://github.com/GreptimeTeam/greptimedb.git
synced 2026-05-22 07:50:38 +00:00
2b2ea5bf72
* chore: upgrade some dependencies * chore: upgrade some dependencies * fix: cr * fix: ci * fix: test * fix: cargo fmt
113 lines
4.0 KiB
Rust
113 lines
4.0 KiB
Rust
// Copyright 2023 Greptime Team
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
use std::marker::PhantomData;
|
|
|
|
use datatypes::value::Value;
|
|
use derive_builder::Builder;
|
|
use rand::seq::{IndexedRandom, SliceRandom};
|
|
use rand::Rng;
|
|
|
|
use super::TsValueGenerator;
|
|
use crate::context::TableContextRef;
|
|
use crate::error::{Error, Result};
|
|
use crate::fake::WordGenerator;
|
|
use crate::generator::{Generator, Random, ValueGenerator};
|
|
use crate::ir::insert_expr::{InsertIntoExpr, RowValue};
|
|
use crate::ir::{generate_random_timestamp, generate_random_value, Ident};
|
|
|
|
/// Generates [InsertIntoExpr].
|
|
#[derive(Builder)]
|
|
#[builder(pattern = "owned")]
|
|
pub struct InsertExprGenerator<R: Rng + 'static> {
|
|
table_ctx: TableContextRef,
|
|
// Whether to omit all columns, i.e. INSERT INTO table_name VALUES (...)
|
|
omit_column_list: bool,
|
|
#[builder(default = "1")]
|
|
rows: usize,
|
|
#[builder(default = "Box::new(WordGenerator)")]
|
|
word_generator: Box<dyn Random<Ident, R>>,
|
|
#[builder(default = "Box::new(generate_random_value)")]
|
|
value_generator: ValueGenerator<R>,
|
|
#[builder(default = "Box::new(generate_random_timestamp)")]
|
|
ts_value_generator: TsValueGenerator<R>,
|
|
#[builder(default)]
|
|
_phantom: PhantomData<R>,
|
|
}
|
|
|
|
impl<R: Rng + 'static> Generator<InsertIntoExpr, R> for InsertExprGenerator<R> {
|
|
type Error = Error;
|
|
|
|
/// Generates the [InsertIntoExpr].
|
|
fn generate(&self, rng: &mut R) -> Result<InsertIntoExpr> {
|
|
let mut values_columns = vec![];
|
|
if self.omit_column_list {
|
|
// If omit column list, then all columns are required in the values list
|
|
values_columns.clone_from(&self.table_ctx.columns);
|
|
} else {
|
|
for column in &self.table_ctx.columns {
|
|
let can_omit = column.is_nullable() || column.has_default_value();
|
|
|
|
// 50% chance to omit a column if it's not required
|
|
if !can_omit || rng.random_bool(0.5) {
|
|
values_columns.push(column.clone());
|
|
}
|
|
}
|
|
values_columns.shuffle(rng);
|
|
|
|
// If all columns are omitted, pick a random column
|
|
if values_columns.is_empty() {
|
|
values_columns.push(self.table_ctx.columns.choose(rng).unwrap().clone());
|
|
}
|
|
}
|
|
|
|
let mut values_list = Vec::with_capacity(self.rows);
|
|
for _ in 0..self.rows {
|
|
let mut row = Vec::with_capacity(values_columns.len());
|
|
for column in &values_columns {
|
|
if column.is_nullable() && rng.random_bool(0.2) {
|
|
row.push(RowValue::Value(Value::Null));
|
|
continue;
|
|
}
|
|
|
|
if column.has_default_value() && rng.random_bool(0.2) {
|
|
row.push(RowValue::Default);
|
|
continue;
|
|
}
|
|
if column.is_time_index() {
|
|
row.push(RowValue::Value((self.ts_value_generator)(
|
|
rng,
|
|
column.timestamp_type().unwrap(),
|
|
)));
|
|
} else {
|
|
row.push(RowValue::Value((self.value_generator)(
|
|
rng,
|
|
&column.column_type,
|
|
Some(self.word_generator.as_ref()),
|
|
)));
|
|
}
|
|
}
|
|
|
|
values_list.push(row);
|
|
}
|
|
|
|
Ok(InsertIntoExpr {
|
|
table_name: self.table_ctx.name.clone(),
|
|
omit_column_list: self.omit_column_list,
|
|
columns: values_columns,
|
|
values_list,
|
|
})
|
|
}
|
|
}
|