MySQL: Support EXPLAIN ANALYZE format variants (#1945)

Author: yoavcloud

src/ast/mod.rs

@@ -4126,7 +4126,7 @@ pub enum Statement {
         /// A SQL query that specifies what to explain
         statement: Box<Statement>,
         /// Optional output format of explain
-        format: Option<AnalyzeFormat>,
+        format: Option<AnalyzeFormatKind>,
         /// Postgres style utility options, `(analyze, verbose true)`
         options: Option<Vec<UtilityOption>>,
     },
@@ -4494,7 +4494,7 @@ impl fmt::Display for Statement {
                 }
 
                 if let Some(format) = format {
-                    write!(f, "FORMAT {format} ")?;
+                    write!(f, "{format} ")?;
                 }
 
                 if let Some(options) = options {
@@ -7641,13 +7641,34 @@ impl fmt::Display for DuplicateTreatment {
     }
 }
 
+#[derive(Debug, Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum AnalyzeFormatKind {
+    /// e.g. `EXPLAIN ANALYZE FORMAT JSON SELECT * FROM tbl`
+    Keyword(AnalyzeFormat),
+    /// e.g. `EXPLAIN ANALYZE FORMAT=JSON SELECT * FROM tbl`
+    Assignment(AnalyzeFormat),
+}
+
+impl fmt::Display for AnalyzeFormatKind {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        match self {
+            AnalyzeFormatKind::Keyword(format) => write!(f, "FORMAT {format}"),
+            AnalyzeFormatKind::Assignment(format) => write!(f, "FORMAT={format}"),
+        }
+    }
+}
+
 #[derive(Debug, Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
 pub enum AnalyzeFormat {
     TEXT,
     GRAPHVIZ,
     JSON,
+    TRADITIONAL,
+    TREE,
 }
 
 impl fmt::Display for AnalyzeFormat {
@@ -7656,6 +7677,8 @@ impl fmt::Display for AnalyzeFormat {
             AnalyzeFormat::TEXT => "TEXT",
             AnalyzeFormat::GRAPHVIZ => "GRAPHVIZ",
             AnalyzeFormat::JSON => "JSON",
+            AnalyzeFormat::TRADITIONAL => "TRADITIONAL",
+            AnalyzeFormat::TREE => "TREE",
         })
     }
 }

src/parser/mod.rs

@@ -5674,6 +5674,14 @@ impl<'a> Parser<'a> {
         }
     }
 
+    fn parse_analyze_format_kind(&mut self) -> Result<AnalyzeFormatKind, ParserError> {
+        if self.consume_token(&Token::Eq) {
+            Ok(AnalyzeFormatKind::Assignment(self.parse_analyze_format()?))
+        } else {
+            Ok(AnalyzeFormatKind::Keyword(self.parse_analyze_format()?))
+        }
+    }
+
     pub fn parse_analyze_format(&mut self) -> Result<AnalyzeFormat, ParserError> {
         let next_token = self.next_token();
         match &next_token.token {
@@ -11074,7 +11082,7 @@ impl<'a> Parser<'a> {
             analyze = self.parse_keyword(Keyword::ANALYZE);
             verbose = self.parse_keyword(Keyword::VERBOSE);
             if self.parse_keyword(Keyword::FORMAT) {
-                format = Some(self.parse_analyze_format()?);
+                format = Some(self.parse_analyze_format_kind()?);
             }
         }
 

tests/sqlparser_common.rs

@@ -5161,7 +5161,7 @@ fn run_explain_analyze(
     query: &str,
     expected_verbose: bool,
     expected_analyze: bool,
-    expected_format: Option<AnalyzeFormat>,
+    expected_format: Option<AnalyzeFormatKind>,
     expected_options: Option<Vec<UtilityOption>>,
 ) {
     match dialect.verified_stmt(query) {
@@ -5272,7 +5272,7 @@ fn parse_explain_analyze_with_simple_select() {
         "EXPLAIN ANALYZE FORMAT GRAPHVIZ SELECT sqrt(id) FROM foo",
         false,
         true,
-        Some(AnalyzeFormat::GRAPHVIZ),
+        Some(AnalyzeFormatKind::Keyword(AnalyzeFormat::GRAPHVIZ)),
         None,
     );
 
@@ -5281,7 +5281,16 @@ fn parse_explain_analyze_with_simple_select() {
         "EXPLAIN ANALYZE VERBOSE FORMAT JSON SELECT sqrt(id) FROM foo",
         true,
         true,
-        Some(AnalyzeFormat::JSON),
+        Some(AnalyzeFormatKind::Keyword(AnalyzeFormat::JSON)),
+        None,
+    );
+
+    run_explain_analyze(
+        all_dialects(),
+        "EXPLAIN ANALYZE VERBOSE FORMAT=JSON SELECT sqrt(id) FROM foo",
+        true,
+        true,
+        Some(AnalyzeFormatKind::Assignment(AnalyzeFormat::JSON)),
         None,
     );
 
@@ -5290,7 +5299,7 @@ fn parse_explain_analyze_with_simple_select() {
         "EXPLAIN VERBOSE FORMAT TEXT SELECT sqrt(id) FROM foo",
         true,
         false,
-        Some(AnalyzeFormat::TEXT),
+        Some(AnalyzeFormatKind::Keyword(AnalyzeFormat::TEXT)),
         None,
     );
 }