From 673a4bd4ef2b6095e32129d8ca5fd2e52cf475a6 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Thu, 25 Jan 2024 17:00:42 +0900 Subject: [PATCH] feat: add pg create alter table expr translator (#3206) * feat: add pg create table expr translator * feat: add pg alter table expr translator * refactor: refactor MappedGenerator * chore: apply suggestions from CR --- tests-fuzz/src/fake.rs | 38 +++-- tests-fuzz/src/generator/create_expr.rs | 4 +- tests-fuzz/src/translator.rs | 3 +- .../src/translator/{greptime.rs => mysql.rs} | 0 .../{greptime => mysql}/alter_expr.rs | 0 .../{greptime => mysql}/create_expr.rs | 0 tests-fuzz/src/translator/postgres.rs | 25 +++ .../src/translator/postgres/alter_expr.rs | 154 ++++++++++++++++++ .../src/translator/postgres/create_expr.rs | 95 +++++++++++ 9 files changed, 306 insertions(+), 13 deletions(-) rename tests-fuzz/src/translator/{greptime.rs => mysql.rs} (100%) rename tests-fuzz/src/translator/{greptime => mysql}/alter_expr.rs (100%) rename tests-fuzz/src/translator/{greptime => mysql}/create_expr.rs (100%) create mode 100644 tests-fuzz/src/translator/postgres.rs create mode 100644 tests-fuzz/src/translator/postgres/alter_expr.rs create mode 100644 tests-fuzz/src/translator/postgres/create_expr.rs diff --git a/tests-fuzz/src/fake.rs b/tests-fuzz/src/fake.rs index 7462ead912..fd3ac336ad 100644 --- a/tests-fuzz/src/fake.rs +++ b/tests-fuzz/src/fake.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::collections::HashSet; +use std::marker::PhantomData; use lazy_static::lazy_static; use rand::seq::{IteratorRandom, SliceRandom}; @@ -50,11 +51,16 @@ faker_impl_from_values!(Word, LOREM_WORDS); pub struct WordGenerator; impl_random!(String, WordGenerator, LOREM_WORDS); -pub type WordMapFn = Box String>; - -pub struct MapWordGenerator { - base: WordGenerator, - map: WordMapFn, +pub struct MappedGenerator +where + T: Random, + F: Fn(&mut R, V) -> V, + R: Rng, +{ + base: T, + map: F, + _r: PhantomData, + _v: PhantomData, } pub fn random_capitalize_map(rng: &mut R, s: String) -> String { @@ -154,17 +160,29 @@ pub fn merge_two_word_map_fn( } } -impl MapWordGenerator { - pub fn new(map: WordMapFn) -> Self { +impl MappedGenerator +where + T: Random, + F: Fn(&mut R, V) -> V, + R: Rng, +{ + pub fn new(base: T, map: F) -> Self { Self { - base: WordGenerator, + base, map, + _r: Default::default(), + _v: Default::default(), } } } -impl Random for MapWordGenerator { - fn choose(&self, rng: &mut R, amount: usize) -> Vec { +impl Random for MappedGenerator +where + T: Random, + F: Fn(&mut R, V) -> V, + R: Rng, +{ + fn choose(&self, rng: &mut R, amount: usize) -> Vec { self.base .choose(rng, amount) .into_iter() diff --git a/tests-fuzz/src/generator/create_expr.rs b/tests-fuzz/src/generator/create_expr.rs index 3914d3b375..ae506e6501 100644 --- a/tests-fuzz/src/generator/create_expr.rs +++ b/tests-fuzz/src/generator/create_expr.rs @@ -20,7 +20,7 @@ use snafu::{ensure, ResultExt}; use super::Generator; use crate::error::{self, Error, Result}; -use crate::fake::{random_capitalize_map, MapWordGenerator}; +use crate::fake::{random_capitalize_map, MappedGenerator, WordGenerator}; use crate::generator::{ColumnOptionGenerator, ConcreteDataTypeGenerator, Random}; use crate::ir::create_expr::CreateTableExprBuilder; use crate::ir::{ @@ -58,7 +58,7 @@ impl Default for CreateTableExprGenerator { if_not_exists: false, partition: 0, name: String::new(), - name_generator: Box::new(MapWordGenerator::new(Box::new(random_capitalize_map))), + name_generator: Box::new(MappedGenerator::new(WordGenerator, random_capitalize_map)), ts_column_type_generator: Box::new(TsColumnTypeGenerator), column_type_generator: Box::new(ColumnTypeGenerator), partible_column_type_generator: Box::new(PartibleColumnTypeGenerator), diff --git a/tests-fuzz/src/translator.rs b/tests-fuzz/src/translator.rs index 2ec4b63954..1745aa9336 100644 --- a/tests-fuzz/src/translator.rs +++ b/tests-fuzz/src/translator.rs @@ -12,7 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. -pub mod greptime; +pub mod mysql; +pub mod postgres; use std::fmt; diff --git a/tests-fuzz/src/translator/greptime.rs b/tests-fuzz/src/translator/mysql.rs similarity index 100% rename from tests-fuzz/src/translator/greptime.rs rename to tests-fuzz/src/translator/mysql.rs diff --git a/tests-fuzz/src/translator/greptime/alter_expr.rs b/tests-fuzz/src/translator/mysql/alter_expr.rs similarity index 100% rename from tests-fuzz/src/translator/greptime/alter_expr.rs rename to tests-fuzz/src/translator/mysql/alter_expr.rs diff --git a/tests-fuzz/src/translator/greptime/create_expr.rs b/tests-fuzz/src/translator/mysql/create_expr.rs similarity index 100% rename from tests-fuzz/src/translator/greptime/create_expr.rs rename to tests-fuzz/src/translator/mysql/create_expr.rs diff --git a/tests-fuzz/src/translator/postgres.rs b/tests-fuzz/src/translator/postgres.rs new file mode 100644 index 0000000000..25e0ff6ed4 --- /dev/null +++ b/tests-fuzz/src/translator/postgres.rs @@ -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 sqlparser::ast::DataType as SqlDataType; + +pub mod alter_expr; +pub mod create_expr; + +pub fn sql_data_type_to_postgres_data_type(data_type: SqlDataType) -> String { + match data_type { + SqlDataType::Double => "DOUBLE PRECISION".to_string(), + _ => data_type.to_string(), + } +} diff --git a/tests-fuzz/src/translator/postgres/alter_expr.rs b/tests-fuzz/src/translator/postgres/alter_expr.rs new file mode 100644 index 0000000000..6be23e72d0 --- /dev/null +++ b/tests-fuzz/src/translator/postgres/alter_expr.rs @@ -0,0 +1,154 @@ +// 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::data_type::ConcreteDataType; +use sql::statements::concrete_data_type_to_sql_data_type; + +use crate::error::{Error, Result}; +use crate::ir::alter_expr::AlterTableOperation; +use crate::ir::create_expr::ColumnOption; +use crate::ir::{AlterTableExpr, Column}; +use crate::translator::postgres::sql_data_type_to_postgres_data_type; +use crate::translator::DslTranslator; + +pub struct AlterTableExprTranslator; + +impl DslTranslator for AlterTableExprTranslator { + type Error = Error; + + fn translate(&self, input: &AlterTableExpr) -> Result { + Ok(match &input.alter_options { + AlterTableOperation::AddColumn { column, .. } => { + Self::format_add_column(&input.name, column) + } + AlterTableOperation::DropColumn { name } => Self::format_drop(&input.name, name), + AlterTableOperation::RenameTable { new_table_name } => { + Self::format_rename(&input.name, new_table_name) + } + }) + } +} + +impl AlterTableExprTranslator { + fn format_drop(name: &str, column: &str) -> String { + format!("ALTER TABLE {name} DROP COLUMN {column};") + } + + fn format_rename(name: &str, new_name: &str) -> String { + format!("ALTER TABLE {name} RENAME TO {new_name};") + } + + fn format_add_column(name: &str, column: &Column) -> String { + format!( + "{};", + vec![format!( + "ALTER TABLE {name} ADD COLUMN {}", + Self::format_column(column) + ),] + .into_iter() + .filter(|s| !s.is_empty()) + .collect::>() + .join(" ") + ) + } + + fn format_column(column: &Column) -> String { + vec![ + column.name.to_string(), + Self::format_column_type(&column.column_type), + Self::format_column_options(&column.options), + ] + .into_iter() + .filter(|s| !s.is_empty()) + .collect::>() + .join(" ") + } + + fn format_column_type(column_type: &ConcreteDataType) -> String { + sql_data_type_to_postgres_data_type( + // Safety: We don't use the `Dictionary` type + concrete_data_type_to_sql_data_type(column_type).unwrap(), + ) + } + + fn acceptable_column_option(option: &ColumnOption) -> bool { + matches!( + option, + ColumnOption::Null + | ColumnOption::NotNull + | ColumnOption::DefaultValue(_) + | ColumnOption::DefaultFn(_) + ) + } + + fn format_column_options(options: &[ColumnOption]) -> String { + options + .iter() + .filter(|opt| Self::acceptable_column_option(opt)) + .map(|option| option.to_string()) + .collect::>() + .join(" ") + } +} + +#[cfg(test)] +mod tests { + use common_query::AddColumnLocation; + use datatypes::data_type::ConcreteDataType; + + use super::AlterTableExprTranslator; + use crate::ir::alter_expr::AlterTableOperation; + use crate::ir::create_expr::ColumnOption; + use crate::ir::{AlterTableExpr, Column}; + use crate::translator::DslTranslator; + + #[test] + fn test_alter_table_expr() { + let alter_expr = AlterTableExpr { + name: "test".to_string(), + alter_options: AlterTableOperation::AddColumn { + column: Column { + name: "host".to_string(), + column_type: ConcreteDataType::string_datatype(), + options: vec![ColumnOption::PrimaryKey], + }, + location: Some(AddColumnLocation::First), + }, + }; + + let output = AlterTableExprTranslator.translate(&alter_expr).unwrap(); + // Ignores the location and primary key option. + assert_eq!("ALTER TABLE test ADD COLUMN host STRING;", output); + + let alter_expr = AlterTableExpr { + name: "test".to_string(), + alter_options: AlterTableOperation::RenameTable { + new_table_name: "foo".to_string(), + }, + }; + + let output = AlterTableExprTranslator.translate(&alter_expr).unwrap(); + assert_eq!("ALTER TABLE test RENAME TO foo;", output); + + let alter_expr = AlterTableExpr { + name: "test".to_string(), + alter_options: AlterTableOperation::DropColumn { + name: "foo".to_string(), + }, + }; + + let output = AlterTableExprTranslator.translate(&alter_expr).unwrap(); + assert_eq!("ALTER TABLE test DROP COLUMN foo;", output); + } +} diff --git a/tests-fuzz/src/translator/postgres/create_expr.rs b/tests-fuzz/src/translator/postgres/create_expr.rs new file mode 100644 index 0000000000..d44bd9383a --- /dev/null +++ b/tests-fuzz/src/translator/postgres/create_expr.rs @@ -0,0 +1,95 @@ +// 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::data_type::ConcreteDataType; +use sql::statements::concrete_data_type_to_sql_data_type; + +use crate::error::{Error, Result}; +use crate::ir::create_expr::ColumnOption; +use crate::ir::{Column, CreateTableExpr}; +use crate::translator::postgres::sql_data_type_to_postgres_data_type; +use crate::translator::DslTranslator; + +pub struct CreateTableExprTranslator; + +impl DslTranslator for CreateTableExprTranslator { + type Error = Error; + + fn translate(&self, input: &CreateTableExpr) -> Result { + Ok(format!( + "CREATE TABLE{}{}(\n{}\n);", + Self::create_if_not_exists(input), + input.name, + Self::format_columns(input), + )) + } +} + +impl CreateTableExprTranslator { + fn create_if_not_exists(input: &CreateTableExpr) -> &str { + if input.if_not_exists { + " IF NOT EXISTS " + } else { + " " + } + } + + fn format_columns(input: &CreateTableExpr) -> String { + let mut output = + Vec::with_capacity(input.columns.len() + (!input.primary_keys.is_empty()) as usize); + for column in &input.columns { + output.push(Self::format_column(column)); + } + output.join(",\n") + } + + fn format_column(column: &Column) -> String { + vec![ + column.name.to_string(), + Self::format_column_type(&column.column_type), + Self::format_column_options(&column.options), + ] + .into_iter() + .filter(|s| !s.is_empty()) + .collect::>() + .join(" ") + } + + fn format_column_type(column_type: &ConcreteDataType) -> String { + sql_data_type_to_postgres_data_type( + // Safety: We don't use the `Dictionary` type + concrete_data_type_to_sql_data_type(column_type).unwrap(), + ) + } + + fn acceptable_column_option(option: &ColumnOption) -> bool { + matches!( + option, + ColumnOption::Null + | ColumnOption::NotNull + | ColumnOption::DefaultValue(_) + | ColumnOption::DefaultFn(_) + ) + } + + fn format_column_options(options: &[ColumnOption]) -> String { + let mut output = Vec::with_capacity(options.len()); + for option in options { + if Self::acceptable_column_option(option) { + output.push(option.to_string()); + } + } + output.join(" ") + } +}