Add optional format for explain (#621)

* Add format for explain

* Add comment

Author: Dandandan

src/ast/mod.rs

@@ -1294,6 +1294,8 @@ pub enum Statement {
         verbose: bool,
         /// A SQL query that specifies what to explain
         statement: Box<Statement>,
+        /// Optional output format of explain
+        format: Option<AnalyzeFormat>,
     },
     /// SAVEPOINT -- define a new savepoint within the current transaction
     Savepoint { name: Ident },
@@ -1344,6 +1346,7 @@ impl fmt::Display for Statement {
                 verbose,
                 analyze,
                 statement,
+                format,
             } => {
                 if *describe_alias {
                     write!(f, "DESCRIBE ")?;
@@ -1359,6 +1362,10 @@ impl fmt::Display for Statement {
                     write!(f, "VERBOSE ")?;
                 }
 
+                if let Some(format) = format {
+                    write!(f, "FORMAT {} ", format)?;
+                }
+
                 write!(f, "{}", statement)
             }
             Statement::Query(s) => write!(f, "{}", s),
@@ -2489,6 +2496,24 @@ pub struct Function {
     pub special: bool,
 }
 
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+pub enum AnalyzeFormat {
+    TEXT,
+    GRAPHVIZ,
+    JSON,
+}
+
+impl fmt::Display for AnalyzeFormat {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        f.write_str(match self {
+            AnalyzeFormat::TEXT => "TEXT",
+            AnalyzeFormat::GRAPHVIZ => "GRAPHVIZ",
+            AnalyzeFormat::JSON => "JSON",
+        })
+    }
+}
+
 impl fmt::Display for Function {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         if self.special {

src/keywords.rs

@@ -256,6 +256,7 @@ define_keywords!(
     GLOBAL,
     GRANT,
     GRANTED,
+    GRAPHVIZ,
     GROUP,
     GROUPING,
     GROUPS,
@@ -288,6 +289,7 @@ define_keywords!(
     ISOYEAR,
     JAR,
     JOIN,
+    JSON,
     JSONFILE,
     JULIAN,
     KEY,

src/parser.rs

@@ -2023,6 +2023,18 @@ impl<'a> Parser<'a> {
         }
     }
 
+    pub fn parse_analyze_format(&mut self) -> Result<AnalyzeFormat, ParserError> {
+        match self.next_token() {
+            Token::Word(w) => match w.keyword {
+                Keyword::TEXT => Ok(AnalyzeFormat::TEXT),
+                Keyword::GRAPHVIZ => Ok(AnalyzeFormat::GRAPHVIZ),
+                Keyword::JSON => Ok(AnalyzeFormat::JSON),
+                _ => self.expected("fileformat", Token::Word(w)),
+            },
+            unexpected => self.expected("fileformat", unexpected),
+        }
+    }
+
     pub fn parse_create_view(&mut self, or_replace: bool) -> Result<Statement, ParserError> {
         let materialized = self.parse_keyword(Keyword::MATERIALIZED);
         self.expect_keyword(Keyword::VIEW)?;
@@ -3432,13 +3444,18 @@ impl<'a> Parser<'a> {
     pub fn parse_explain(&mut self, describe_alias: bool) -> Result<Statement, ParserError> {
         let analyze = self.parse_keyword(Keyword::ANALYZE);
         let verbose = self.parse_keyword(Keyword::VERBOSE);
+        let mut format = None;
+        if self.parse_keyword(Keyword::FORMAT) {
+            format = Some(self.parse_analyze_format()?);
+        }
 
         if let Some(statement) = self.maybe_parse(|parser| parser.parse_statement()) {
             Ok(Statement::Explain {
                 describe_alias,
                 analyze,
                 verbose,
                 statement: Box::new(statement),
+                format,
             })
         } else {
             let table_name = self.parse_object_name()?;

tests/sqlparser_common.rs

@@ -2657,16 +2657,23 @@ fn parse_scalar_function_in_projection() {
     }
 }
 
-fn run_explain_analyze(query: &str, expected_verbose: bool, expected_analyze: bool) {
+fn run_explain_analyze(
+    query: &str,
+    expected_verbose: bool,
+    expected_analyze: bool,
+    expected_format: Option<AnalyzeFormat>,
+) {
     match verified_stmt(query) {
         Statement::Explain {
             describe_alias: _,
             analyze,
             verbose,
             statement,
+            format,
         } => {
             assert_eq!(verbose, expected_verbose);
             assert_eq!(analyze, expected_analyze);
+            assert_eq!(format, expected_format);
             assert_eq!("SELECT sqrt(id) FROM foo", statement.to_string());
         }
         _ => panic!("Unexpected Statement, must be Explain"),
@@ -2693,15 +2700,47 @@ fn parse_explain_table() {
 #[test]
 fn parse_explain_analyze_with_simple_select() {
     // Describe is an alias for EXPLAIN
-    run_explain_analyze("DESCRIBE SELECT sqrt(id) FROM foo", false, false);
+    run_explain_analyze("DESCRIBE SELECT sqrt(id) FROM foo", false, false, None);
 
-    run_explain_analyze("EXPLAIN SELECT sqrt(id) FROM foo", false, false);
-    run_explain_analyze("EXPLAIN VERBOSE SELECT sqrt(id) FROM foo", true, false);
-    run_explain_analyze("EXPLAIN ANALYZE SELECT sqrt(id) FROM foo", false, true);
+    run_explain_analyze("EXPLAIN SELECT sqrt(id) FROM foo", false, false, None);
+    run_explain_analyze(
+        "EXPLAIN VERBOSE SELECT sqrt(id) FROM foo",
+        true,
+        false,
+        None,
+    );
+    run_explain_analyze(
+        "EXPLAIN ANALYZE SELECT sqrt(id) FROM foo",
+        false,
+        true,
+        None,
+    );
     run_explain_analyze(
         "EXPLAIN ANALYZE VERBOSE SELECT sqrt(id) FROM foo",
         true,
         true,
+        None,
+    );
+
+    run_explain_analyze(
+        "EXPLAIN ANALYZE FORMAT GRAPHVIZ SELECT sqrt(id) FROM foo",
+        false,
+        true,
+        Some(AnalyzeFormat::GRAPHVIZ),
+    );
+
+    run_explain_analyze(
+        "EXPLAIN ANALYZE VERBOSE FORMAT JSON SELECT sqrt(id) FROM foo",
+        true,
+        true,
+        Some(AnalyzeFormat::JSON),
+    );
+
+    run_explain_analyze(
+        "EXPLAIN VERBOSE FORMAT TEXT SELECT sqrt(id) FROM foo",
+        true,
+        false,
+        Some(AnalyzeFormat::TEXT),
     );
 }