feat: add alter table sql smith impl (datafuselabs#13231)

andylokandy · Nov 27, 2023 · db5d4a2 · db5d4a2
1 parent 14e3b1e
commit db5d4a2
Show file tree

Hide file tree

Showing 3 changed files with 194 additions and 1 deletion.
diff --git a/src/tests/sqlsmith/src/runner.rs b/src/tests/sqlsmith/src/runner.rs
@@ -136,12 +136,29 @@ impl Runner {
         let table_stmts = generator.gen_base_tables();
         let tables = self.create_base_table(table_stmts).await.unwrap();
         let conn = self.client.get_conn().await.unwrap();
+        let row_count = 50;
 
         for table in &tables {
-            let insert_stmt = generator.gen_insert(table, 50);
+            let insert_stmt = generator.gen_insert(table, row_count);
             let insert_sql = insert_stmt.to_string();
             tracing::info!("insert_sql: {}", insert_sql);
             conn.exec(&insert_sql).await.unwrap();
+            let alter_stmt_opt = generator.gen_alter(table, row_count);
+            if let Some((alter_stmt, insert_stmt_opt)) = alter_stmt_opt {
+                let alter_sql = alter_stmt.to_string();
+                tracing::info!("alter_sql: {}", alter_sql);
+                if let Err(err) = conn.exec(&alter_sql).await {
+                    tracing::info!("alter_sql err: {}", err);
+                    continue;
+                }
+                if let Some(insert_stmt) = insert_stmt_opt {
+                    let insert_sql = insert_stmt.to_string();
+                    tracing::info!("after alter insert_sql: {}", insert_sql);
+                    if let Err(err) = conn.exec(&insert_sql).await {
+                        tracing::info!("after alter insert_sql err: {}", err);
+                    }
+                }
+            }
         }
         generator.tables = tables;
 

diff --git a/src/tests/sqlsmith/src/sql_gen/ddl.rs b/src/tests/sqlsmith/src/sql_gen/ddl.rs
@@ -24,6 +24,7 @@ use common_ast::ast::NullableConstraint;
 use common_ast::ast::TypeName;
 use rand::Rng;
 
+use super::Table;
 use crate::sql_gen::SqlGenerator;
 
 const BASE_TABLE_NAMES: [&str; 4] = ["t1", "t2", "t3", "t4"];
@@ -149,6 +150,38 @@ impl<'a, R: Rng> SqlGenerator<'a, R> {
         }
     }
 
+    pub fn gen_data_type_name(&mut self) -> (TypeName, Option<NullableConstraint>) {
+        let i = self.rng.gen_range(0..=36);
+        if i <= 17 {
+            (
+                SIMPLE_COLUMN_TYPES[i].clone(),
+                Some(NullableConstraint::NotNull),
+            )
+        } else if i <= 35 {
+            (
+                SIMPLE_COLUMN_TYPES[i - 18].clone(),
+                Some(NullableConstraint::Null),
+            )
+        } else {
+            // TODO: add nested data types
+            let depth = self.rng.gen_range(1..=3);
+            (self.gen_nested_type(depth), None)
+        }
+    }
+
+    pub fn gen_new_column(&mut self, table: &Table) -> ColumnDefinition {
+        let field_num = table.schema.num_fields();
+        let new_column_name = Identifier::from_name(format!("cc{:?}", field_num));
+        let (data_type, nullable_constraint) = self.gen_data_type_name();
+        ColumnDefinition {
+            name: new_column_name,
+            data_type,
+            expr: None,
+            comment: None,
+            nullable_constraint,
+        }
+    }
+
     fn gen_table_source(&mut self) -> CreateTableSource {
         let mut column_defs = Vec::with_capacity(40);
 

diff --git a/src/tests/sqlsmith/src/sql_gen/dml.rs b/src/tests/sqlsmith/src/sql_gen/dml.rs
@@ -13,12 +13,20 @@
 // limitations under the License.
 
 use chrono_tz::Tz;
+use common_ast::ast::AddColumnOption;
+use common_ast::ast::AlterTableAction;
+use common_ast::ast::AlterTableStmt;
+use common_ast::ast::ColumnDefinition;
 use common_ast::ast::Identifier;
 use common_ast::ast::InsertSource;
 use common_ast::ast::InsertStmt;
+use common_ast::ast::NullableConstraint;
+use common_ast::ast::TableReference;
+use common_exception::Span;
 use common_expression::types::DataType;
 use common_expression::Column;
 use common_expression::ScalarRef;
+use common_expression::TableField;
 use common_formats::field_encoder::FieldEncoderRowBased;
 use common_formats::field_encoder::FieldEncoderValues;
 use common_formats::CommonSettings;
@@ -27,6 +35,7 @@ use common_io::constants::INF_BYTES_LOWER;
 use common_io::constants::NAN_BYTES_LOWER;
 use common_io::constants::NULL_BYTES_UPPER;
 use common_io::constants::TRUE_BYTES_LOWER;
+use common_sql::resolve_type_name;
 use itertools::join;
 use rand::Rng;
 use roaring::RoaringTreemap;
@@ -59,6 +68,140 @@ impl<'a, R: Rng> SqlGenerator<'a, R> {
         }
     }
 
+    fn is_column_not_null(column: &ColumnDefinition) -> bool {
+        match column.nullable_constraint {
+            Some(NullableConstraint::NotNull) => true,
+            Some(NullableConstraint::Null) => false,
+            None => true,
+        }
+    }
+
+    fn random_select_field(&mut self, table: &Table) -> TableField {
+        let field_index = self.rng.gen_range(0..=table.schema.num_fields());
+        table.schema.fields[field_index].clone()
+    }
+
+    // generate alter table statement, and insert statement of new column(if any)
+    pub(crate) fn gen_alter(
+        &mut self,
+        table: &Table,
+        row_count: usize,
+    ) -> Option<(AlterTableStmt, Option<InsertStmt>)> {
+        if self.rng.gen_bool(0.3) {
+            return None;
+        }
+        let (action, new_column) = match self.rng.gen_range(0..=4) {
+            0 => {
+                let new_table = format!("{}_{}", table.name, self.rng.gen_range(0..10));
+                (
+                    AlterTableAction::RenameTable {
+                        new_table: Identifier::from_name(new_table),
+                    },
+                    None,
+                )
+            }
+            1 => {
+                let column = self.gen_new_column(table);
+                let option = match self.rng.gen_range(0..=2) {
+                    0 => AddColumnOption::End,
+                    1 => AddColumnOption::First,
+                    2 => {
+                        let field = self.random_select_field(table);
+                        let column = Identifier::from_name(field.name);
+                        AddColumnOption::After(column)
+                    }
+                    _ => unreachable!(),
+                };
+                (
+                    AlterTableAction::AddColumn {
+                        column: column.clone(),
+                        option,
+                    },
+                    Some(column),
+                )
+            }
+            2 => {
+                let field = self.random_select_field(table);
+                let old_column = Identifier::from_name(field.name);
+                let new_column = self.gen_new_column(table).name;
+                (
+                    AlterTableAction::RenameColumn {
+                        old_column,
+                        new_column,
+                    },
+                    None,
+                )
+            }
+            3 => {
+                let field = self.random_select_field(table);
+                let name = Identifier::from_name(field.name);
+                let (data_type, nullable_constraint) = self.gen_data_type_name();
+                let new_column = ColumnDefinition {
+                    name,
+                    data_type,
+                    expr: None,
+                    comment: None,
+                    nullable_constraint,
+                };
+                (
+                    AlterTableAction::ModifyColumn {
+                        action: common_ast::ast::ModifyColumnAction::SetDataType(vec![
+                            new_column.clone(),
+                        ]),
+                    },
+                    Some(new_column),
+                )
+            }
+            4 => {
+                let field = self.random_select_field(table);
+                let column = Identifier::from_name(field.name);
+                (AlterTableAction::DropColumn { column }, None)
+            }
+            _ => unreachable!(),
+        };
+
+        let insert_stmt_opt = if let Some(new_column) = new_column {
+            let not_null = Self::is_column_not_null(&new_column);
+            let table_name = Identifier::from_name(table.name.clone());
+            let columns = vec![new_column.name.clone()];
+            let data_type = resolve_type_name(&new_column.data_type, not_null).unwrap();
+            let data_types = vec![(&data_type).into()];
+            let source = self.gen_insert_source(&data_types, row_count);
+
+            Some(InsertStmt {
+                // TODO
+                hints: None,
+                catalog: None,
+                database: None,
+                table: table_name,
+                columns,
+                source,
+                overwrite: false,
+            })
+        } else {
+            None
+        };
+
+        let table_reference = TableReference::Table {
+            span: Span::default(),
+            catalog: None,
+            database: None,
+            table: Identifier::from_name(table.name.clone()),
+            alias: None,
+            travel_point: None,
+            pivot: None,
+            unpivot: None,
+        };
+        Some((
+            AlterTableStmt {
+                if_exists: true,
+                action,
+                table_reference,
+            },
+            insert_stmt_opt,
+        ))
+    }
+
     fn gen_insert_source(&mut self, data_types: &[DataType], row_count: usize) -> InsertSource {
         match self.rng.gen_range(0..=9) {
             0..=9 => {