Skip to content

Commit

Permalink
feat(fuzz): add fuzz_alter_table target
Browse files Browse the repository at this point in the history
  • Loading branch information
WenyXu committed Mar 13, 2024
1 parent 51dd36d commit 2bb2c24
Show file tree
Hide file tree
Showing 6 changed files with 448 additions and 20 deletions.
7 changes: 7 additions & 0 deletions tests-fuzz/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -55,3 +55,10 @@ path = "targets/fuzz_insert.rs"
test = false
bench = false
doc = false

[[bin]]
name = "fuzz_alter_table"
path = "targets/fuzz_alter_table.rs"
test = false
bench = false
doc = false
189 changes: 188 additions & 1 deletion tests-fuzz/src/context.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,13 +14,20 @@

use std::sync::Arc;

use common_query::AddColumnLocation;
use partition::partition::PartitionDef;
use rand::Rng;
use snafu::{ensure, OptionExt};

use crate::ir::{Column, CreateTableExpr, Ident};
use crate::error::{self, Result};
use crate::generator::Random;
use crate::ir::alter_expr::AlterTableOperation;
use crate::ir::{AlterTableExpr, Column, CreateTableExpr, Ident};

pub type TableContextRef = Arc<TableContext>;

/// TableContext stores table info.
#[derive(Debug, Clone)]
pub struct TableContext {
pub name: Ident,
pub columns: Vec<Column>,
Expand Down Expand Up @@ -48,3 +55,183 @@ impl From<&CreateTableExpr> for TableContext {
}
}
}

impl TableContext {
/// Applies the [AlterTableExpr].
pub fn alter(mut self, expr: AlterTableExpr) -> Result<TableContext> {
match expr.alter_options {
AlterTableOperation::AddColumn { column, location } => {
ensure!(
!self.columns.iter().any(|col| col.name == column.name),
error::UnexpectedSnafu {
violated: format!("Column {} exists", column.name),
}
);
match location {
Some(AddColumnLocation::First) => {
let mut columns = Vec::with_capacity(self.columns.len() + 1);
columns.push(column);
columns.extend(self.columns);
self.columns = columns;
}
Some(AddColumnLocation::After { column_name }) => {
let index = self
.columns
.iter()
// TODO(weny): find a better way?
.position(|col| col.name.to_string() == column_name)
.context(error::UnexpectedSnafu {
violated: format!("Column: {column_name} not found"),
})?;
self.columns.insert(index + 1, column);
}
None => self.columns.push(column),
}
// Re-generates the primary_keys
self.primary_keys = self
.columns
.iter()
.enumerate()
.flat_map(|(idx, col)| {
if col.is_primary_key() {
Some(idx)
} else {
None
}
})
.collect();
Ok(self)
}
AlterTableOperation::DropColumn { name } => {
self.columns.retain(|col| col.name != name);
// Re-generates the primary_keys
self.primary_keys = self
.columns
.iter()
.enumerate()
.flat_map(|(idx, col)| {
if col.is_primary_key() {
Some(idx)
} else {
None
}
})
.collect();
Ok(self)
}
AlterTableOperation::RenameTable { new_table_name } => {
ensure!(
new_table_name != self.name,
error::UnexpectedSnafu {
violated: "The new table name is equal the current name",
}
);
self.name = new_table_name;
Ok(self)
}
}
}

pub fn generate_unique_column_name<R: Rng>(
&self,
rng: &mut R,
generator: &dyn Random<Ident, R>,
) -> Ident {
let mut name = generator.gen(rng);
while self.columns.iter().any(|col| col.name.value == name.value) {
name = generator.gen(rng);
}
name
}

pub fn generate_unique_table_name<R: Rng>(
&self,
rng: &mut R,
generator: &dyn Random<Ident, R>,
) -> Ident {
let mut name = generator.gen(rng);
while self.name.value == name.value {
name = generator.gen(rng);
}
name
}
}

#[cfg(test)]
mod tests {
use common_query::AddColumnLocation;
use datatypes::data_type::ConcreteDataType;

use super::TableContext;
use crate::ir::alter_expr::AlterTableOperation;
use crate::ir::create_expr::ColumnOption;
use crate::ir::{AlterTableExpr, Column, Ident};

#[test]
fn test_table_context_alter() {
let table_ctx = TableContext {
name: "foo".into(),
columns: vec![],
partition: None,
primary_keys: vec![],
};
// Add a column
let expr = AlterTableExpr {
table_name: "foo".into(),
alter_options: AlterTableOperation::AddColumn {
column: Column {
name: "a".into(),
column_type: ConcreteDataType::timestamp_microsecond_datatype(),
options: vec![ColumnOption::PrimaryKey],
},
location: None,
},
};
let table_ctx = table_ctx.alter(expr).unwrap();
assert_eq!(table_ctx.columns[0].name, Ident::new("a"));
assert_eq!(table_ctx.primary_keys, vec![0]);

// Add a column at first
let expr = AlterTableExpr {
table_name: "foo".into(),
alter_options: AlterTableOperation::AddColumn {
column: Column {
name: "b".into(),
column_type: ConcreteDataType::timestamp_microsecond_datatype(),
options: vec![ColumnOption::PrimaryKey],
},
location: Some(AddColumnLocation::First),
},
};
let table_ctx = table_ctx.alter(expr).unwrap();
assert_eq!(table_ctx.columns[0].name, Ident::new("b"));
assert_eq!(table_ctx.primary_keys, vec![0, 1]);

// Add a column after "b"
let expr = AlterTableExpr {
table_name: "foo".into(),
alter_options: AlterTableOperation::AddColumn {
column: Column {
name: "c".into(),
column_type: ConcreteDataType::timestamp_microsecond_datatype(),
options: vec![ColumnOption::PrimaryKey],
},
location: Some(AddColumnLocation::After {
column_name: "b".into(),
}),
},
};
let table_ctx = table_ctx.alter(expr).unwrap();
assert_eq!(table_ctx.columns[1].name, Ident::new("c"));
assert_eq!(table_ctx.primary_keys, vec![0, 1, 2]);

// Drop the column "b"
let expr = AlterTableExpr {
table_name: "foo".into(),
alter_options: AlterTableOperation::DropColumn { name: "b".into() },
};
let table_ctx = table_ctx.alter(expr).unwrap();
assert_eq!(table_ctx.columns[1].name, Ident::new("a"));
assert_eq!(table_ctx.primary_keys, vec![0, 1]);
}
}
47 changes: 40 additions & 7 deletions tests-fuzz/src/generator/alter_expr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
use std::marker::PhantomData;

use common_query::AddColumnLocation;
use datatypes::data_type::ConcreteDataType;
use derive_builder::Builder;
use rand::Rng;
use snafu::ensure;
Expand All @@ -24,10 +25,38 @@ use crate::error::{self, Error, Result};
use crate::fake::WordGenerator;
use crate::generator::{ColumnOptionGenerator, ConcreteDataTypeGenerator, Generator, Random};
use crate::ir::alter_expr::{AlterTableExpr, AlterTableOperation};
use crate::ir::create_expr::ColumnOption;
use crate::ir::{
column_options_generator, droppable_columns, generate_columns, ColumnTypeGenerator, Ident,
droppable_columns, generate_columns, generate_random_value, ColumnTypeGenerator, Ident,
};

fn add_column_options_generator<R: Rng>(
rng: &mut R,
column_type: &ConcreteDataType,
) -> Vec<ColumnOption> {
// 0 -> NULL
// 1 -> DEFAULT VALUE
// 2 -> PRIMARY KEY + DEFAULT VALUE
let idx = rng.gen_range(0..3);
match idx {
0 => vec![ColumnOption::Null],
1 => {
vec![ColumnOption::DefaultValue(generate_random_value(
rng,
column_type,
None,
))]
}
2 => {
vec![
ColumnOption::PrimaryKey,
ColumnOption::DefaultValue(generate_random_value(rng, column_type, None)),
]
}
_ => unreachable!(),
}
}

/// Generates the [AlterTableOperation::AddColumn] of [AlterTableExpr].
#[derive(Builder)]
#[builder(pattern = "owned")]
Expand All @@ -37,7 +66,7 @@ pub struct AlterExprAddColumnGenerator<R: Rng + 'static> {
location: bool,
#[builder(default = "Box::new(WordGenerator)")]
name_generator: Box<dyn Random<Ident, R>>,
#[builder(default = "Box::new(column_options_generator)")]
#[builder(default = "Box::new(add_column_options_generator)")]
column_options_generator: ColumnOptionGenerator<R>,
#[builder(default = "Box::new(ColumnTypeGenerator)")]
column_type_generator: ConcreteDataTypeGenerator<R>,
Expand Down Expand Up @@ -65,7 +94,9 @@ impl<R: Rng + 'static> Generator<AlterTableExpr, R> for AlterExprAddColumnGenera
None
};

let name = self.name_generator.gen(rng);
let name = self
.table_ctx
.generate_unique_column_name(rng, self.name_generator.as_ref());
let column = generate_columns(
rng,
vec![name],
Expand Down Expand Up @@ -116,7 +147,9 @@ impl<R: Rng> Generator<AlterTableExpr, R> for AlterExprRenameGenerator<R> {
type Error = Error;

fn generate(&self, rng: &mut R) -> Result<AlterTableExpr> {
let new_table_name = self.name_generator.gen(rng);
let new_table_name = self
.table_ctx
.generate_unique_table_name(rng, self.name_generator.as_ref());
Ok(AlterTableExpr {
table_name: self.table_ctx.name.clone(),
alter_options: AlterTableOperation::RenameTable { new_table_name },
Expand Down Expand Up @@ -153,7 +186,7 @@ mod tests {
.generate(&mut rng)
.unwrap();
let serialized = serde_json::to_string(&expr).unwrap();
let expected = r#"{"table_name":{"value":"animI","quote_style":null},"alter_options":{"AddColumn":{"column":{"name":{"value":"velit","quote_style":null},"column_type":{"Int32":{}},"options":[{"DefaultValue":{"Int32":853246610}}]},"location":null}}}"#;
let expected = r#"{"table_name":{"value":"IuRe","quote_style":null},"alter_options":{"AddColumn":{"column":{"name":{"value":"voluptas","quote_style":null},"column_type":{"Int32":{}},"options":["PrimaryKey",{"DefaultValue":{"Int32":-343607799}}]},"location":null}}}"#;
assert_eq!(expected, serialized);

let expr = AlterExprRenameGeneratorBuilder::default()
Expand All @@ -163,7 +196,7 @@ mod tests {
.generate(&mut rng)
.unwrap();
let serialized = serde_json::to_string(&expr).unwrap();
let expected = r#"{"table_name":{"value":"animI","quote_style":null},"alter_options":{"RenameTable":{"new_table_name":{"value":"iure","quote_style":null}}}}"#;
let expected = r#"{"table_name":{"value":"IuRe","quote_style":null},"alter_options":{"RenameTable":{"new_table_name":{"value":"dolorum","quote_style":null}}}}"#;
assert_eq!(expected, serialized);

let expr = AlterExprDropColumnGeneratorBuilder::default()
Expand All @@ -173,7 +206,7 @@ mod tests {
.generate(&mut rng)
.unwrap();
let serialized = serde_json::to_string(&expr).unwrap();
let expected = r#"{"table_name":{"value":"animI","quote_style":null},"alter_options":{"DropColumn":{"name":{"value":"toTAm","quote_style":null}}}}"#;
let expected = r#"{"table_name":{"value":"IuRe","quote_style":null},"alter_options":{"DropColumn":{"name":{"value":"ADIPisCI","quote_style":null}}}}"#;
assert_eq!(expected, serialized);
}
}
40 changes: 29 additions & 11 deletions tests-fuzz/src/ir.rs
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ use rand::seq::SliceRandom;
use rand::Rng;
use serde::{Deserialize, Serialize};

use self::create_expr::ColumnOptionType;
use crate::generator::Random;
use crate::impl_random;
use crate::ir::create_expr::ColumnOption;
Expand Down Expand Up @@ -274,17 +275,34 @@ pub fn column_options_generator<R: Rng>(
// 2 -> DEFAULT VALUE
// 3 -> PRIMARY KEY
// 4 -> EMPTY
let option_idx = rng.gen_range(0..5);
match option_idx {
0 => vec![ColumnOption::Null],
1 => vec![ColumnOption::NotNull],
2 => vec![ColumnOption::DefaultValue(generate_random_value(
rng,
column_type,
None,
))],
3 => vec![ColumnOption::PrimaryKey],
_ => vec![],
make_column_options_generator(vec![
ColumnOptionType::Null,
ColumnOptionType::NotNull,
ColumnOptionType::DefaultValue,
ColumnOptionType::PrimaryKey,
])(rng, column_type)
}

/// Generates one of [ColumnOption] in `allowed_column_options` or returns empty options.
pub fn make_column_options_generator<R: Rng>(
allowed_column_options: Vec<ColumnOptionType>,
) -> impl Fn(&mut R, &ConcreteDataType) -> Vec<ColumnOption> {
move |rng, column_type| {
let idx = rng.gen_range(0..allowed_column_options.len() + 1);
if let Some(t) = allowed_column_options.get(idx) {
match t {
ColumnOptionType::Null => vec![ColumnOption::Null],
ColumnOptionType::NotNull => vec![ColumnOption::NotNull],
ColumnOptionType::DefaultValue => vec![ColumnOption::DefaultValue(
generate_random_value(rng, column_type, None),
)],
ColumnOptionType::PrimaryKey => vec![ColumnOption::PrimaryKey],
ColumnOptionType::DefaultFn => todo!(),
ColumnOptionType::TimeIndex => unreachable!(),
}
} else {
vec![]
}
}
}

Expand Down
Loading

0 comments on commit 2bb2c24

Please sign in to comment.