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

@@ -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<R> = Box<dyn Fn(&mut R, String) -> String>;
-
-pub struct MapWordGenerator<R: Rng> {
-    base: WordGenerator,
-    map: WordMapFn<R>,
+pub struct MappedGenerator<T, F, R, V>
+where
+    T: Random<V, R>,
+    F: Fn(&mut R, V) -> V,
+    R: Rng,
+{
+    base: T,
+    map: F,
+    _r: PhantomData<R>,
+    _v: PhantomData<V>,
 }
 
 pub fn random_capitalize_map<R: Rng + 'static>(rng: &mut R, s: String) -> String {
@@ -154,17 +160,29 @@ pub fn merge_two_word_map_fn<R: Rng>(
     }
 }
 
-impl<R: Rng> MapWordGenerator<R> {
-    pub fn new(map: WordMapFn<R>) -> Self {
+impl<T, F, R, V> MappedGenerator<T, F, R, V>
+where
+    T: Random<V, R>,
+    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<R: Rng> Random<String, R> for MapWordGenerator<R> {
-    fn choose(&self, rng: &mut R, amount: usize) -> Vec<String> {
+impl<T, F, R, V> Random<V, R> for MappedGenerator<T, F, R, V>
+where
+    T: Random<V, R>,
+    F: Fn(&mut R, V) -> V,
+    R: Rng,
+{
+    fn choose(&self, rng: &mut R, amount: usize) -> Vec<V> {
         self.base
             .choose(rng, amount)
             .into_iter()

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<R: Rng + 'static> Default for CreateTableExprGenerator<R> {
             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),

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;
 

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(),
+    }
+}

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<AlterTableExpr, String> for AlterTableExprTranslator {
+    type Error = Error;
+
+    fn translate(&self, input: &AlterTableExpr) -> Result<String> {
+        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::<Vec<_>>()
+            .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::<Vec<_>>()
+        .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::<Vec<_>>()
+            .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);
+    }
+}

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<CreateTableExpr, String> for CreateTableExprTranslator {
+    type Error = Error;
+
+    fn translate(&self, input: &CreateTableExpr) -> Result<String> {
+        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::<Vec<_>>()
+        .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(" ")
+    }
+}