Skip to content

Commit

Permalink
feat: add AlterTableExprGenerator
Browse files Browse the repository at this point in the history
  • Loading branch information
WenyXu committed Jan 17, 2024
1 parent 2f9ab93 commit 568c717
Show file tree
Hide file tree
Showing 7 changed files with 247 additions and 22 deletions.
3 changes: 3 additions & 0 deletions tests-fuzz/src/error.rs
Original file line number Diff line number Diff line change
Expand Up @@ -35,4 +35,7 @@ pub enum Error {
error: CreateTableExprBuilderError,
location: Location,
},

#[snafu(display("No droppable columns"))]
DroppableColumns { location: Location },
}
9 changes: 6 additions & 3 deletions tests-fuzz/src/generator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,18 +12,21 @@
// See the License for the specific language governing permissions and
// limitations under the License.

pub mod alter_expr;
pub mod create_expr;

use std::fmt;

use crate::error::Error;
use crate::ir::CreateTableExpr;
use crate::ir::{AlterTableExpr, CreateTableExpr};

pub type CreateTableExprGenerator =
Box<dyn Generator<CreateTableExpr, Error = Error> + Sync + Send>;

#[async_trait::async_trait]
pub type AlterTableExprGenerator = Box<dyn Generator<AlterTableExpr, Error = Error> + Sync + Send>;

pub(crate) trait Generator<T> {
type Error: Sync + Send + fmt::Debug;

async fn generate(&self) -> Result<T, Self::Error>;
fn generate(&self) -> Result<T, Self::Error>;
}
142 changes: 142 additions & 0 deletions tests-fuzz/src/generator/alter_expr.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,142 @@
// 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::f64::consts::E;
use std::sync::Arc;

use common_query::AddColumnLocation;
use faker_rand::lorem::Word;
use rand::{random, Rng};
use snafu::ensure;

use crate::error::{self, Result};
use crate::ir::alter_expr::{AlterTableExpr, AlterTableOperation};
use crate::ir::{droppable_columns, Column};

pub struct AlterTableExprGenerator {
name: String,
columns: Arc<Vec<Column>>,
ignore_no_droppable_column: bool,
}

impl AlterTableExprGenerator {
pub fn new(name: String, columns: Arc<Vec<Column>>) -> Self {
Self {
name,
columns,
ignore_no_droppable_column: false,
}
}

/// If the `ignore_no_droppable_column` is true, it retries if there is no droppable column.
pub fn ignore_no_droppable_column(mut self, v: bool) -> Self {
self.ignore_no_droppable_column = v;
self
}

fn generate_inner(&self) -> Result<AlterTableExpr> {
let mut rng = rand::thread_rng();
let idx = rng.gen_range(0..3);
// 0 -> AddColumn
// 1 -> DropColumn(invariant: We can't non-primary key columns, non-ts columns)
// 2 -> RenameTable
let alter_expr = match idx {
0 => {
let with_location = rng.gen::<bool>();
let location = if with_location {
let use_first = rng.gen::<bool>();
let location = if use_first {
AddColumnLocation::First
} else {
AddColumnLocation::After {
column_name: self.columns[rng.gen_range(0..self.columns.len())]
.name
.to_string(),
}
};
Some(location)
} else {
None
};
let column = rng.gen::<Column>();
AlterTableExpr {
name: self.name.to_string(),
alter_options: AlterTableOperation::AddColumn { column, location },
}
}
1 => {
let droppable = droppable_columns(&self.columns);
ensure!(!droppable.is_empty(), error::DroppableColumnsSnafu);
let name = droppable[rng.gen_range(0..droppable.len())]
.name
.to_string();
AlterTableExpr {
name: self.name.to_string(),
alter_options: AlterTableOperation::DropColumn { name },
}
}
2 => {
let mut new_table_name = rng.gen::<Word>().to_string();
if new_table_name == self.name {
new_table_name = format!("{}-{}", self.name, rng.gen::<u64>());
}
AlterTableExpr {
name: self.name.to_string(),
alter_options: AlterTableOperation::RenameTable { new_table_name },
}
}
_ => unreachable!(),
};

Ok(alter_expr)
}

/// Generates the [AlterTableExpr].
pub fn generate(&self) -> Result<AlterTableExpr> {
match self.generate_inner() {
Ok(expr) => Ok(expr),
Err(err) => {
if matches!(err, error::Error::DroppableColumns { .. }) {
return self.generate();
}
Err(err)
}
}
}
}

#[cfg(test)]
mod tests {
use std::sync::Arc;

use super::AlterTableExprGenerator;
use crate::generator::create_expr::CreateTableExprGenerator;
use crate::generator::Generator;

#[test]
fn test_alter_table_expr_generator() {
let create_expr = CreateTableExprGenerator::default()
.columns(10)
.generate()
.unwrap();

let alter_expr = AlterTableExprGenerator::new(
create_expr.name.to_string(),
Arc::new(create_expr.columns),
)
.ignore_no_droppable_column(true)
.generate()
.unwrap();
}
}
24 changes: 7 additions & 17 deletions tests-fuzz/src/generator/create_expr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -74,7 +74,7 @@ impl CreateTableExprGenerator {
}

/// Generates the [CreateTableExpr].
fn generate_inner(&self) -> Result<CreateTableExpr> {
pub fn generate(&self) -> Result<CreateTableExpr> {
ensure!(
self.columns != 0,
error::UnexpectedSnafu {
Expand All @@ -90,12 +90,11 @@ impl CreateTableExprGenerator {
// Generates columns.
for i in 0..self.columns {
let mut column = rng.gen::<Column>();
// Removes the primary key option.
let options = column
let is_primary_key = column
.options
.extract_if(|option| option == &ColumnOption::PrimaryKey)
.collect::<Vec<_>>();
if !options.is_empty() {
.iter()
.any(|option| option == &ColumnOption::PrimaryKey);
if is_primary_key {
primary_keys.push(i);
}
columns.push(column);
Expand Down Expand Up @@ -154,15 +153,6 @@ impl CreateTableExprGenerator {
}
}

#[async_trait::async_trait]
impl Generator<CreateTableExpr> for CreateTableExprGenerator {
type Error = Error;

async fn generate(&self) -> Result<CreateTableExpr> {
self.generate_inner()
}
}

#[cfg(test)]
mod tests {
use super::*;
Expand All @@ -174,7 +164,7 @@ mod tests {
.partitions(3)
.create_is_not_exists(true)
.engine("mito2")
.generate_inner()
.generate()
.unwrap();
assert_eq!(expr.engine, "mito2");
assert!(expr.if_not_exists);
Expand All @@ -184,7 +174,7 @@ mod tests {
let expr = CreateTableExprGenerator::default()
.columns(10)
.partitions(1)
.generate_inner()
.generate()
.unwrap();
assert_eq!(expr.columns.len(), 11);
assert_eq!(expr.partitions.len(), 0);
Expand Down
52 changes: 52 additions & 0 deletions tests-fuzz/src/ir.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,8 +14,10 @@

//! The intermediate representation
pub(crate) mod alter_expr;
pub(crate) mod create_expr;

pub use alter_expr::AlterTableExpr;
pub use create_expr::CreateTableExpr;
use datatypes::data_type::ConcreteDataType;
use datatypes::value::Value;
Expand Down Expand Up @@ -80,6 +82,18 @@ pub struct Column {
pub options: Vec<ColumnOption>,
}

/// Returns droppable columns. i.e., non-primary key columns, non-ts columns.
pub fn droppable_columns(columns: &[Column]) -> Vec<&Column> {
columns
.iter()
.filter(|column| {
!column.options.iter().any(|option| {
option == &ColumnOption::PrimaryKey || option == &ColumnOption::TimeIndex
})
})
.collect::<Vec<_>>()
}

impl Distribution<Column> for Standard {
fn sample<R: rand::prelude::Rng + ?Sized>(&self, rng: &mut R) -> Column {
let column_type = DATA_TYPES[rng.gen_range(0..DATA_TYPES.len())].clone();
Expand Down Expand Up @@ -134,3 +148,41 @@ impl Distribution<PartibleColumn> for Standard {
})
}
}

#[cfg(test)]
mod tests {
use super::*;

#[test]
fn test_droppable_columns() {
let columns = vec![
Column {
name: "hi".to_string(),
column_type: ConcreteDataType::uint64_datatype(),
options: vec![ColumnOption::PrimaryKey],
},
Column {
name: "foo".to_string(),
column_type: ConcreteDataType::uint64_datatype(),
options: vec![ColumnOption::TimeIndex],
},
];
let droppable = droppable_columns(&columns);
assert!(droppable.is_empty());

let columns = vec![
Column {
name: "hi".to_string(),
column_type: ConcreteDataType::uint64_datatype(),
options: vec![],
},
Column {
name: "foo".to_string(),
column_type: ConcreteDataType::uint64_datatype(),
options: vec![],
},
];
let droppable = droppable_columns(&columns);
assert_eq!(droppable.len(), 2);
}
}
37 changes: 37 additions & 0 deletions tests-fuzz/src/ir/alter_expr.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,37 @@
// 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 common_query::AddColumnLocation;
use derive_builder::Builder;

use crate::ir::Column;

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

#[derive(Debug, Clone)]
pub enum AlterTableOperation {
/// `ADD [ COLUMN ] <column_def> [location]`
AddColumn {
column: Column,
location: Option<AddColumnLocation>,
},
/// `DROP COLUMN <name>`
DropColumn { name: String },
/// `RENAME <new_table_name>`
RenameTable { new_table_name: String },
}
2 changes: 0 additions & 2 deletions tests-fuzz/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,8 +12,6 @@
// See the License for the specific language governing permissions and
// limitations under the License.

#![feature(extract_if)]

pub(crate) mod context;
pub(crate) mod error;
pub(crate) mod executor;
Expand Down

0 comments on commit 568c717

Please sign in to comment.