1pub mod alter_expr;
16pub mod create_expr;
17pub mod insert_expr;
18pub mod repartition_expr;
19pub mod select_expr;
20
21use std::fmt;
22
23use datatypes::data_type::ConcreteDataType;
24use datatypes::types::TimestampType;
25use datatypes::value::Value;
26use rand::Rng;
27
28use crate::error::Error;
29use crate::ir::create_expr::ColumnOption;
30use crate::ir::{AlterTableExpr, Column, CreateTableExpr, Ident, RowValue};
31
32pub type CreateTableExprGenerator<R> =
33 Box<dyn Generator<CreateTableExpr, R, Error = Error> + Sync + Send>;
34
35pub type AlterTableExprGenerator<R> =
36 Box<dyn Generator<AlterTableExpr, R, Error = Error> + Sync + Send>;
37
38pub type ColumnOptionGenerator<R> = Box<dyn Fn(&mut R, &ConcreteDataType) -> Vec<ColumnOption>>;
39
40pub type ConcreteDataTypeGenerator<R> = Box<dyn Random<ConcreteDataType, R>>;
41
42pub type ValueGenerator<R> =
43 Box<dyn Fn(&mut R, &ConcreteDataType, Option<&dyn Random<Ident, R>>) -> Value>;
44
45pub type TsValueGenerator<R> = Box<dyn Fn(&mut R, TimestampType) -> Value>;
46
47pub type ValueOverride<R> = Box<dyn Fn(&Column, &mut R) -> Option<RowValue>>;
48
49pub trait Generator<T, R: Rng> {
50 type Error: Sync + Send + fmt::Debug;
51
52 fn generate(&self, rng: &mut R) -> Result<T, Self::Error>;
53}
54
55pub trait Random<T, R: Rng> {
56 fn generate(&self, rng: &mut R) -> T {
58 self.choose(rng, 1).remove(0)
59 }
60
61 fn choose(&self, rng: &mut R, amount: usize) -> Vec<T>;
63}
64
65#[macro_export]
66macro_rules! impl_random {
67 ($type: ident, $value:ident, $values: ident) => {
68 impl<R: Rng> Random<$type, R> for $value {
69 fn choose(&self, rng: &mut R, amount: usize) -> Vec<$type> {
70 let mut result = std::collections::BTreeSet::new();
72 while result.len() != amount {
73 result.insert($values.choose(rng).unwrap().clone());
74 }
75 let mut result = result.into_iter().map(Into::into).collect::<Vec<_>>();
76 result.shuffle(rng);
78 result
79 }
80 }
81 };
82}