Skip to content

Commit

Permalink
feat: add insert/select generator & translator (#3240)
Browse files Browse the repository at this point in the history
* feat: add insert into expr generator & translator

* feat: add select expr generator & translator

* chore: apply suggestions from CR

* fix: fix unit tests
  • Loading branch information
WenyXu authored Jan 26, 2024
1 parent 1fab7ab commit f99b087
Show file tree
Hide file tree
Showing 20 changed files with 487 additions and 65 deletions.
2 changes: 1 addition & 1 deletion tests-fuzz/src/context.rs
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ pub struct TableContext {
impl From<&CreateTableExpr> for TableContext {
fn from(
CreateTableExpr {
name,
table_name: name,
columns,
partition,
primary_keys,
Expand Down
2 changes: 2 additions & 0 deletions tests-fuzz/src/generator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@

pub mod alter_expr;
pub mod create_expr;
pub mod insert_expr;
pub mod select_expr;

use std::fmt;

Expand Down
12 changes: 6 additions & 6 deletions tests-fuzz/src/generator/alter_expr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -74,7 +74,7 @@ impl<R: Rng + 'static> Generator<AlterTableExpr, R> for AlterExprAddColumnGenera
)
.remove(0);
Ok(AlterTableExpr {
name: self.table_ctx.name.to_string(),
table_name: self.table_ctx.name.to_string(),
alter_options: AlterTableOperation::AddColumn { column, location },
})
}
Expand All @@ -99,7 +99,7 @@ impl<R: Rng> Generator<AlterTableExpr, R> for AlterExprDropColumnGenerator<R> {
.name
.to_string();
Ok(AlterTableExpr {
name: self.table_ctx.name.to_string(),
table_name: self.table_ctx.name.to_string(),
alter_options: AlterTableOperation::DropColumn { name },
})
}
Expand All @@ -120,7 +120,7 @@ impl<R: Rng> Generator<AlterTableExpr, R> for AlterExprRenameGenerator<R> {
fn generate(&self, rng: &mut R) -> Result<AlterTableExpr> {
let new_table_name = self.name_generator.gen(rng);
Ok(AlterTableExpr {
name: self.table_ctx.name.to_string(),
table_name: self.table_ctx.name.to_string(),
alter_options: AlterTableOperation::RenameTable { new_table_name },
})
}
Expand Down Expand Up @@ -155,7 +155,7 @@ mod tests {
.generate(&mut rng)
.unwrap();
let serialized = serde_json::to_string(&expr).unwrap();
let expected = r#"{"name":"DigNissIMOS","alter_options":{"AddColumn":{"column":{"name":"sit","column_type":{"Boolean":null},"options":["PrimaryKey"]},"location":null}}}"#;
let expected = r#"{"table_name":"DigNissIMOS","alter_options":{"AddColumn":{"column":{"name":"sit","column_type":{"Boolean":null},"options":["PrimaryKey"]},"location":null}}}"#;
assert_eq!(expected, serialized);

let expr = AlterExprRenameGeneratorBuilder::default()
Expand All @@ -165,7 +165,7 @@ mod tests {
.generate(&mut rng)
.unwrap();
let serialized = serde_json::to_string(&expr).unwrap();
let expected = r#"{"name":"DigNissIMOS","alter_options":{"RenameTable":{"new_table_name":"excepturi"}}}"#;
let expected = r#"{"table_name":"DigNissIMOS","alter_options":{"RenameTable":{"new_table_name":"excepturi"}}}"#;
assert_eq!(expected, serialized);

let expr = AlterExprDropColumnGeneratorBuilder::default()
Expand All @@ -176,7 +176,7 @@ mod tests {
.unwrap();
let serialized = serde_json::to_string(&expr).unwrap();
let expected =
r#"{"name":"DigNissIMOS","alter_options":{"DropColumn":{"name":"INVentORE"}}}"#;
r#"{"table_name":"DigNissIMOS","alter_options":{"DropColumn":{"name":"INVentORE"}}}"#;
assert_eq!(expected, serialized);
}
}
8 changes: 5 additions & 3 deletions tests-fuzz/src/generator/create_expr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -106,6 +106,7 @@ impl<R: Rng + 'static> Generator<CreateTableExpr, R> for CreateTableExprGenerato
partition_bounds.push(PartitionBound::Value(generate_random_value(
rng,
&column.column_type,
None,
)));
partition_bounds.sort();
}
Expand Down Expand Up @@ -133,6 +134,7 @@ impl<R: Rng + 'static> Generator<CreateTableExpr, R> for CreateTableExprGenerato
partition_bounds.push(PartitionBound::Value(generate_random_value(
rng,
&column.column_type,
None,
)));
partition_bounds.sort();
}
Expand Down Expand Up @@ -171,9 +173,9 @@ impl<R: Rng + 'static> Generator<CreateTableExpr, R> for CreateTableExprGenerato
builder.engine(self.engine.to_string());
builder.if_not_exists(self.if_not_exists);
if self.name.is_empty() {
builder.name(self.name_generator.gen(rng));
builder.table_name(self.name_generator.gen(rng));
} else {
builder.name(self.name.to_string());
builder.table_name(self.name.to_string());
}
builder.build().context(error::BuildCreateTableExprSnafu)
}
Expand Down Expand Up @@ -228,7 +230,7 @@ mod tests {
.unwrap();

let serialized = serde_json::to_string(&expr).unwrap();
let expected = r#"{"name":"iN","columns":[{"name":"CUlpa","column_type":{"Int16":{}},"options":["PrimaryKey","NotNull"]},{"name":"dEBiTiS","column_type":{"Timestamp":{"Second":null}},"options":["TimeIndex"]},{"name":"HArum","column_type":{"Int16":{}},"options":["NotNull"]},{"name":"NObIS","column_type":{"Int32":{}},"options":["PrimaryKey"]},{"name":"IMPEDiT","column_type":{"Int16":{}},"options":[{"DefaultValue":{"Int16":-25151}}]},{"name":"bLanDITIis","column_type":{"Boolean":null},"options":[{"DefaultValue":{"Boolean":true}}]},{"name":"Dolores","column_type":{"Float32":{}},"options":["PrimaryKey"]},{"name":"eSt","column_type":{"Float32":{}},"options":[{"DefaultValue":{"Float32":0.9152612}}]},{"name":"INVentORE","column_type":{"Int64":{}},"options":["PrimaryKey"]},{"name":"aDIpiSci","column_type":{"Float64":{}},"options":["Null"]}],"if_not_exists":true,"partition":{"partition_columns":["CUlpa"],"partition_bounds":[{"Value":{"Int16":15966}},{"Value":{"Int16":31925}},"MaxValue"]},"engine":"mito2","options":{},"primary_keys":[6,0,8,3]}"#;
let expected = r#"{"table_name":"iN","columns":[{"name":"CUlpa","column_type":{"Int16":{}},"options":["PrimaryKey","NotNull"]},{"name":"dEBiTiS","column_type":{"Timestamp":{"Second":null}},"options":["TimeIndex"]},{"name":"HArum","column_type":{"Int16":{}},"options":["NotNull"]},{"name":"NObIS","column_type":{"Int32":{}},"options":["PrimaryKey"]},{"name":"IMPEDiT","column_type":{"Int16":{}},"options":[{"DefaultValue":{"Int16":-25151}}]},{"name":"bLanDITIis","column_type":{"Boolean":null},"options":[{"DefaultValue":{"Boolean":true}}]},{"name":"Dolores","column_type":{"Float32":{}},"options":["PrimaryKey"]},{"name":"eSt","column_type":{"Float32":{}},"options":[{"DefaultValue":{"Float32":0.9152612}}]},{"name":"INVentORE","column_type":{"Int64":{}},"options":["PrimaryKey"]},{"name":"aDIpiSci","column_type":{"Float64":{}},"options":["Null"]}],"if_not_exists":true,"partition":{"partition_columns":["CUlpa"],"partition_bounds":[{"Value":{"Int16":15966}},{"Value":{"Int16":31925}},"MaxValue"]},"engine":"mito2","options":{},"primary_keys":[6,0,8,3]}"#;
assert_eq!(expected, serialized);
}
}
70 changes: 70 additions & 0 deletions tests-fuzz/src/generator/insert_expr.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,70 @@
// 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 derive_builder::Builder;
use rand::seq::SliceRandom;
use rand::Rng;

use crate::context::TableContextRef;
use crate::error::{Error, Result};
use crate::fake::WordGenerator;
use crate::generator::{Generator, Random};
use crate::ir::generate_random_value;
use crate::ir::insert_expr::InsertIntoExpr;

/// Generates [InsertIntoExpr].
#[derive(Builder)]
#[builder(pattern = "owned")]
pub struct InsertExprGenerator<R: Rng + 'static> {
table_ctx: TableContextRef,
#[builder(default = "1")]
rows: usize,
#[builder(default = "Box::new(WordGenerator)")]
word_generator: Box<dyn Random<String, R>>,
#[builder(default)]
_phantom: PhantomData<R>,
}

impl<R: Rng + 'static> Generator<InsertIntoExpr, R> for InsertExprGenerator<R> {
type Error = Error;

/// Generates the [CreateTableExpr].
fn generate(&self, rng: &mut R) -> Result<InsertIntoExpr> {
let mut columns = self.table_ctx.columns.clone();
columns.shuffle(rng);

let mut rows = Vec::with_capacity(self.rows);
for _ in 0..self.rows {
let mut row = Vec::with_capacity(columns.len());
for column in &columns {
// TODO(weny): generates the special cases
row.push(generate_random_value(
rng,
&column.column_type,
Some(self.word_generator.as_ref()),
));
}

rows.push(row);
}

Ok(InsertIntoExpr {
table_name: self.table_ctx.name.to_string(),
columns,
rows,
})
}
}
72 changes: 72 additions & 0 deletions tests-fuzz/src/generator/select_expr.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,72 @@
// 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 derive_builder::Builder;
use rand::seq::SliceRandom;
use rand::Rng;

use crate::context::TableContextRef;
use crate::error::{Error, Result};
use crate::generator::Generator;
use crate::ir::select_expr::{Direction, SelectExpr};

#[derive(Builder)]
#[builder(pattern = "owned")]
pub struct SelectExprGenerator<R: Rng + 'static> {
table_ctx: TableContextRef,
#[builder(default = "8192")]
max_limit: usize,
#[builder(default)]
_phantom: PhantomData<R>,
}

impl<R: Rng + 'static> Generator<SelectExpr, R> for SelectExprGenerator<R> {
type Error = Error;

fn generate(&self, rng: &mut R) -> Result<SelectExpr> {
let selection = rng.gen_range(1..self.table_ctx.columns.len());
let mut selected_columns = self
.table_ctx
.columns
.choose_multiple(rng, selection)
.cloned()
.collect::<Vec<_>>();
selected_columns.shuffle(rng);

let order_by_selection = rng.gen_range(1..selection);

let order_by = selected_columns
.choose_multiple(rng, order_by_selection)
.map(|c| c.name.to_string())
.collect::<Vec<_>>();

let limit = rng.gen_range(1..self.max_limit);

let direction = if rng.gen_bool(1.0 / 2.0) {
Direction::Asc
} else {
Direction::Desc
};

Ok(SelectExpr {
table_name: self.table_ctx.name.to_string(),
columns: selected_columns,
order_by,
direction,
limit,
})
}
}
17 changes: 14 additions & 3 deletions tests-fuzz/src/ir.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,8 @@
pub(crate) mod alter_expr;
pub(crate) mod create_expr;
pub(crate) mod insert_expr;
pub(crate) mod select_expr;

pub use alter_expr::AlterTableExpr;
pub use create_expr::CreateTableExpr;
Expand Down Expand Up @@ -71,17 +73,25 @@ pub struct TsColumnTypeGenerator;
pub struct PartibleColumnTypeGenerator;

/// Generates a random [Value].
pub fn generate_random_value<R: Rng>(rng: &mut R, datatype: &ConcreteDataType) -> Value {
pub fn generate_random_value<R: Rng>(
rng: &mut R,
datatype: &ConcreteDataType,
random_str: Option<&dyn Random<String, R>>,
) -> Value {
match datatype {
&ConcreteDataType::Boolean(_) => Value::from(rng.gen::<bool>()),
ConcreteDataType::Int16(_) => Value::from(rng.gen::<i16>()),
ConcreteDataType::Int32(_) => Value::from(rng.gen::<i32>()),
ConcreteDataType::Int64(_) => Value::from(rng.gen::<i64>()),
ConcreteDataType::Float32(_) => Value::from(rng.gen::<f32>()),
ConcreteDataType::Float64(_) => Value::from(rng.gen::<f64>()),
ConcreteDataType::String(_) => Value::from(rng.gen::<char>().to_string()),
ConcreteDataType::String(_) => match random_str {
Some(random) => Value::from(random.gen(rng)),
None => Value::from(rng.gen::<char>().to_string()),
},
ConcreteDataType::Date(_) => Value::from(rng.gen::<i32>()),
ConcreteDataType::DateTime(_) => Value::from(rng.gen::<i64>()),
&ConcreteDataType::Timestamp(_) => Value::from(rng.gen::<u64>()),

_ => unimplemented!("unsupported type: {datatype}"),
}
Expand Down Expand Up @@ -142,6 +152,7 @@ pub fn column_options_generator<R: Rng>(
2 => vec![ColumnOption::DefaultValue(generate_random_value(
rng,
column_type,
None,
))],
3 => vec![ColumnOption::PrimaryKey],
_ => vec![],
Expand All @@ -163,7 +174,7 @@ pub fn partible_column_options_generator<R: Rng + 'static>(
1 => vec![ColumnOption::PrimaryKey, ColumnOption::NotNull],
2 => vec![
ColumnOption::PrimaryKey,
ColumnOption::DefaultValue(generate_random_value(rng, column_type)),
ColumnOption::DefaultValue(generate_random_value(rng, column_type, None)),
],
3 => vec![ColumnOption::PrimaryKey],
_ => unreachable!(),
Expand Down
2 changes: 1 addition & 1 deletion tests-fuzz/src/ir/alter_expr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ use crate::ir::Column;

#[derive(Debug, Builder, Clone, Serialize, Deserialize)]
pub struct AlterTableExpr {
pub name: String,
pub table_name: String,
pub alter_options: AlterTableOperation,
}

Expand Down
2 changes: 1 addition & 1 deletion tests-fuzz/src/ir/create_expr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,7 @@ impl Display for ColumnOption {
#[derive(Debug, Builder, Clone, Serialize, Deserialize)]
pub struct CreateTableExpr {
#[builder(setter(into))]
pub name: String,
pub table_name: String,
pub columns: Vec<Column>,
#[builder(default)]
pub if_not_exists: bool,
Expand Down
25 changes: 25 additions & 0 deletions tests-fuzz/src/ir/insert_expr.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
// 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 datatypes::value::Value;

use crate::ir::Column;

pub type RowValue = Vec<Value>;

pub struct InsertIntoExpr {
pub table_name: String,
pub columns: Vec<Column>,
pub rows: Vec<RowValue>,
}
39 changes: 39 additions & 0 deletions tests-fuzz/src/ir/select_expr.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,39 @@
// 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::fmt::Display;

use crate::ir::Column;

pub enum Direction {
Asc,
Desc,
}

impl Display for Direction {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Direction::Asc => write!(f, "ASC"),
Direction::Desc => write!(f, "DESC"),
}
}
}

pub struct SelectExpr {
pub table_name: String,
pub columns: Vec<Column>,
pub order_by: Vec<String>,
pub direction: Direction,
pub limit: usize,
}
3 changes: 3 additions & 0 deletions tests-fuzz/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,3 +21,6 @@ pub mod fake;
pub mod generator;
pub mod ir;
pub mod translator;

#[cfg(test)]
pub mod test_utils;
Loading

0 comments on commit f99b087

Please sign in to comment.