feat: impl Display for Statement (#3744)

* feat: impl Display for Statement * fix: add license header * fix: inline function manually * fix: redacte options * fix: check secret key and replace value * test: add test for statement display * fix: fix check * fix: inline method * fix: inline methods * fix: format * showcase how to write Display impl Signed-off-by: tison <[email protected]> * for others Signed-off-by: tison <[email protected]> * create and copy Signed-off-by: tison <[email protected]> * create rest Signed-off-by: tison <[email protected]> * fixup Signed-off-by: tison <[email protected]> * address comments Signed-off-by: tison <[email protected]> * fixup quote Signed-off-by: tison <[email protected]> --------- Signed-off-by: tison <[email protected]> Co-authored-by: tison <[email protected]>
GreptimeTeam · Apr 24, 2024 · 62037ee · 62037ee
1 parent 8d229dd
commit 62037ee
Show file tree

Hide file tree

Showing 13 changed files with 1,087 additions and 82 deletions.
diff --git a/src/common/query/src/lib.rs b/src/common/query/src/lib.rs
@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::fmt::{Debug, Formatter};
+use std::fmt::{Debug, Display, Formatter};
 use std::sync::Arc;
 
 use api::greptime_proto::v1::add_column_location::LocationType;
@@ -126,6 +126,17 @@ pub enum AddColumnLocation {
     After { column_name: String },
 }
 
+impl Display for AddColumnLocation {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        match self {
+            AddColumnLocation::First => write!(f, r#"FIRST"#),
+            AddColumnLocation::After { column_name } => {
+                write!(f, r#"AFTER {column_name}"#)
+            }
+        }
+    }
+}
+
 impl From<&AddColumnLocation> for Location {
     fn from(value: &AddColumnLocation) -> Self {
         match value {

diff --git a/src/sql/src/statements.rs b/src/sql/src/statements.rs
@@ -43,6 +43,7 @@ use datatypes::schema::constraint::{CURRENT_TIMESTAMP, CURRENT_TIMESTAMP_FN};
 use datatypes::schema::{ColumnDefaultConstraint, ColumnSchema, COMMENT_KEY};
 use datatypes::types::{cast, TimestampType};
 use datatypes::value::{OrderedF32, OrderedF64, Value};
+use itertools::Itertools;
 pub use option_map::OptionMap;
 use snafu::{ensure, OptionExt, ResultExt};
 use sqlparser::ast::{ExactNumberInfo, UnaryOperator};
@@ -58,6 +59,29 @@ use crate::error::{
     SerializeColumnDefaultConstraintSnafu, TimestampOverflowSnafu, UnsupportedDefaultValueSnafu,
 };
 
+const REDACTED_OPTIONS: [&str; 2] = ["access_key_id", "secret_access_key"];
+
+/// Convert the options into redacted and sorted key-value string. Options with key in
+/// [REDACTED_OPTIONS] will be converted into `<key> = '******'`.
+fn redact_and_sort_options(options: &OptionMap) -> Vec<String> {
+    let options = options.as_ref();
+    let mut result = Vec::with_capacity(options.len());
+    let keys = options.keys().sorted();
+    for key in keys {
+        if let Some(val) = options.get(key) {
+            let redacted = REDACTED_OPTIONS
+                .iter()
+                .any(|opt| opt.eq_ignore_ascii_case(key));
+            if redacted {
+                result.push(format!("{key} = '******'"));
+            } else {
+                result.push(format!("{key} = '{}'", val.escape_default()));
+            }
+        }
+    }
+    result
+}
+
 fn parse_string_to_value(
     column_name: &str,
     s: String,

diff --git a/src/sql/src/statements/alter.rs b/src/sql/src/statements/alter.rs
@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::fmt::{Debug, Display};
+
 use common_query::AddColumnLocation;
 use sqlparser::ast::{ColumnDef, Ident, ObjectName, TableConstraint};
 use sqlparser_derive::{Visit, VisitMut};
@@ -39,6 +41,14 @@ impl AlterTable {
     }
 }
 
+impl Display for AlterTable {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let table_name = self.table_name();
+        let alter_operation = self.alter_operation();
+        write!(f, r#"ALTER TABLE {table_name} {alter_operation}"#)
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut)]
 pub enum AlterTableOperation {
     /// `ADD <table_constraint>`
@@ -53,3 +63,100 @@ pub enum AlterTableOperation {
     /// `RENAME <new_table_name>`
     RenameTable { new_table_name: String },
 }
+
+impl Display for AlterTableOperation {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            AlterTableOperation::AddConstraint(constraint) => write!(f, r#"ADD {constraint}"#),
+            AlterTableOperation::AddColumn {
+                column_def,
+                location,
+            } => {
+                if let Some(location) = location {
+                    write!(f, r#"ADD COLUMN {column_def} {location}"#)
+                } else {
+                    write!(f, r#"ADD COLUMN {column_def}"#)
+                }
+            }
+            AlterTableOperation::DropColumn { name } => write!(f, r#"DROP COLUMN {name}"#),
+            AlterTableOperation::RenameTable { new_table_name } => {
+                write!(f, r#"RENAME {new_table_name}"#)
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::assert_matches::assert_matches;
+
+    use crate::dialect::GreptimeDbDialect;
+    use crate::parser::{ParseOptions, ParserContext};
+    use crate::statements::statement::Statement;
+
+    #[test]
+    fn test_display_alter() {
+        let sql = r"alter table monitor add column app string default 'shop' primary key;";
+        let stmts =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
+                .unwrap();
+        assert_eq!(1, stmts.len());
+        assert_matches!(&stmts[0], Statement::Alter { .. });
+
+        match &stmts[0] {
+            Statement::Alter(set) => {
+                let new_sql = format!("\n{}", set);
+                assert_eq!(
+                    r#"
+ALTER TABLE monitor ADD COLUMN app STRING DEFAULT 'shop' PRIMARY KEY"#,
+                    &new_sql
+                );
+            }
+            _ => {
+                unreachable!();
+            }
+        }
+
+        let sql = r"alter table monitor drop column load_15;";
+        let stmts =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
+                .unwrap();
+        assert_eq!(1, stmts.len());
+        assert_matches!(&stmts[0], Statement::Alter { .. });
+
+        match &stmts[0] {
+            Statement::Alter(set) => {
+                let new_sql = format!("\n{}", set);
+                assert_eq!(
+                    r#"
+ALTER TABLE monitor DROP COLUMN load_15"#,
+                    &new_sql
+                );
+            }
+            _ => {
+                unreachable!();
+            }
+        }
+
+        let sql = r"alter table monitor rename monitor_new;";
+        let stmts =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
+                .unwrap();
+        assert_eq!(1, stmts.len());
+        assert_matches!(&stmts[0], Statement::Alter { .. });
+
+        match &stmts[0] {
+            Statement::Alter(set) => {
+                let new_sql = format!("\n{}", set);
+                assert_eq!(
+                    r#"
+ALTER TABLE monitor RENAME monitor_new"#,
+                    &new_sql
+                );
+            }
+            _ => {
+                unreachable!();
+            }
+        }
+    }
+}
diff --git a/src/sql/src/statements/copy.rs b/src/sql/src/statements/copy.rs
@@ -12,29 +12,90 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::fmt::Display;
+
 use sqlparser::ast::ObjectName;
 use sqlparser_derive::{Visit, VisitMut};
 
-use crate::statements::OptionMap;
+use crate::statements::{redact_and_sort_options, OptionMap};
 
 #[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut)]
 pub enum Copy {
     CopyTable(CopyTable),
     CopyDatabase(CopyDatabase),
 }
 
+impl Display for Copy {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Copy::CopyTable(s) => s.fmt(f),
+            Copy::CopyDatabase(s) => s.fmt(f),
+        }
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut)]
 pub enum CopyTable {
     To(CopyTableArgument),
     From(CopyTableArgument),
 }
 
+impl Display for CopyTable {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "COPY ")?;
+        let (with, connection) = match self {
+            CopyTable::To(args) => {
+                write!(f, "{} TO {}", &args.table_name, &args.location)?;
+                (&args.with, &args.connection)
+            }
+            CopyTable::From(args) => {
+                write!(f, "{} FROM {}", &args.table_name, &args.location)?;
+                (&args.with, &args.connection)
+            }
+        };
+        if !with.map.is_empty() {
+            let options = redact_and_sort_options(with);
+            write!(f, " WITH ({})", options.join(", "))?;
+        }
+        if !connection.map.is_empty() {
+            let options = redact_and_sort_options(connection);
+            write!(f, " CONNECTION ({})", options.join(", "))?;
+        }
+        Ok(())
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut)]
 pub enum CopyDatabase {
     To(CopyDatabaseArgument),
     From(CopyDatabaseArgument),
 }
 
+impl Display for CopyDatabase {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "COPY DATABASE ")?;
+        let (with, connection) = match self {
+            CopyDatabase::To(args) => {
+                write!(f, "{} TO {}", &args.database_name, &args.location)?;
+                (&args.with, &args.connection)
+            }
+            CopyDatabase::From(args) => {
+                write!(f, "{} FROM {}", &args.database_name, &args.location)?;
+                (&args.with, &args.connection)
+            }
+        };
+        if !with.map.is_empty() {
+            let options = redact_and_sort_options(with);
+            write!(f, " WITH ({})", options.join(", "))?;
+        }
+        if !connection.map.is_empty() {
+            let options = redact_and_sort_options(connection);
+            write!(f, " CONNECTION ({})", options.join(", "))?;
+        }
+        Ok(())
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut)]
 pub struct CopyDatabaseArgument {
     pub database_name: ObjectName,
@@ -67,3 +128,112 @@ impl CopyTableArgument {
             .cloned()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::assert_matches::assert_matches;
+
+    use crate::dialect::GreptimeDbDialect;
+    use crate::parser::{ParseOptions, ParserContext};
+    use crate::statements::statement::Statement;
+
+    #[test]
+    fn test_display_copy_from_tb() {
+        let sql = r"copy tbl from 's3://my-bucket/data.parquet'
+            with (format = 'parquet', pattern = '.*parquet.*')
+            connection(region = 'us-west-2', secret_access_key = '12345678');";
+        let stmts =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
+                .unwrap();
+        assert_eq!(1, stmts.len());
+        assert_matches!(&stmts[0], Statement::Copy { .. });
+
+        match &stmts[0] {
+            Statement::Copy(copy) => {
+                let new_sql = format!("{}", copy);
+                assert_eq!(
+                    r#"COPY tbl FROM s3://my-bucket/data.parquet WITH (format = 'parquet', pattern = '.*parquet.*') CONNECTION (region = 'us-west-2', secret_access_key = '******')"#,
+                    &new_sql
+                );
+            }
+            _ => {
+                unreachable!();
+            }
+        }
+    }
+
+    #[test]
+    fn test_display_copy_to_tb() {
+        let sql = r"copy tbl to 's3://my-bucket/data.parquet'
+            with (format = 'parquet', pattern = '.*parquet.*')
+            connection(region = 'us-west-2', secret_access_key = '12345678');";
+        let stmts =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
+                .unwrap();
+        assert_eq!(1, stmts.len());
+        assert_matches!(&stmts[0], Statement::Copy { .. });
+
+        match &stmts[0] {
+            Statement::Copy(copy) => {
+                let new_sql = format!("{}", copy);
+                assert_eq!(
+                    r#"COPY tbl TO s3://my-bucket/data.parquet WITH (format = 'parquet', pattern = '.*parquet.*') CONNECTION (region = 'us-west-2', secret_access_key = '******')"#,
+                    &new_sql
+                );
+            }
+            _ => {
+                unreachable!();
+            }
+        }
+    }
+
+    #[test]
+    fn test_display_copy_from_db() {
+        let sql = r"copy database db1 from 's3://my-bucket/data.parquet'
+            with (format = 'parquet', pattern = '.*parquet.*')
+            connection(region = 'us-west-2', secret_access_key = '12345678');";
+        let stmts =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
+                .unwrap();
+        assert_eq!(1, stmts.len());
+        assert_matches!(&stmts[0], Statement::Copy { .. });
+
+        match &stmts[0] {
+            Statement::Copy(copy) => {
+                let new_sql = format!("{}", copy);
+                assert_eq!(
+                    r#"COPY DATABASE db1 FROM s3://my-bucket/data.parquet WITH (format = 'parquet', pattern = '.*parquet.*') CONNECTION (region = 'us-west-2', secret_access_key = '******')"#,
+                    &new_sql
+                );
+            }
+            _ => {
+                unreachable!();
+            }
+        }
+    }
+
+    #[test]
+    fn test_display_copy_to_db() {
+        let sql = r"copy database db1 to 's3://my-bucket/data.parquet'
+            with (format = 'parquet', pattern = '.*parquet.*')
+            connection(region = 'us-west-2', secret_access_key = '12345678');";
+        let stmts =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
+                .unwrap();
+        assert_eq!(1, stmts.len());
+        assert_matches!(&stmts[0], Statement::Copy { .. });
+
+        match &stmts[0] {
+            Statement::Copy(copy) => {
+                let new_sql = format!("{}", copy);
+                assert_eq!(
+                    r#"COPY DATABASE db1 TO s3://my-bucket/data.parquet WITH (format = 'parquet', pattern = '.*parquet.*') CONNECTION (region = 'us-west-2', secret_access_key = '******')"#,
+                    &new_sql
+                );
+            }
+            _ => {
+                unreachable!();
+            }
+        }
+    }
+}