refactor: take prepare and execute parser as submodule

src/sql/src/parser.rs

@@ -180,56 +180,25 @@ impl<'a> ParserContext<'a> {
         sql: &'a str,
         dialect: &dyn Dialect,
     ) -> Result<(String, String)> {
-        let mut parser = Parser::new(dialect)
+        let parser = Parser::new(dialect)
             .with_options(ParserOptions::new().with_trailing_commas(true))
             .try_with_sql(sql)
             .context(SyntaxSnafu)?;
 
-        parser
-            .expect_keyword(Keyword::PREPARE)
-            .context(SyntaxSnafu)?;
-        let stmt_name = parser.parse_identifier(false).context(SyntaxSnafu)?;
-        parser.expect_keyword(Keyword::FROM).context(SyntaxSnafu)?;
-        let next_token = parser.peek_token();
-        let stmt = match next_token.token {
-            Token::SingleQuotedString(s) => {
-                let _ = parser.next_token();
-                s
-            }
-            _ => parser
-                .expected("string literal", next_token)
-                .context(SyntaxSnafu)?,
-        };
-        Ok((stmt_name.value, stmt))
+        ParserContext { parser, sql }.parse_mysql_prepare()
     }
 
     /// Parses MySQL style 'EXECUTE stmt_name USING param_list' into a stmt_name string and a list of parameters.
     pub fn parse_mysql_execute_stmt(
         sql: &'a str,
         dialect: &dyn Dialect,
     ) -> Result<(String, Vec<Expr>)> {
-        let mut parser = Parser::new(dialect)
+        let parser = Parser::new(dialect)
             .with_options(ParserOptions::new().with_trailing_commas(true))
             .try_with_sql(sql)
             .context(SyntaxSnafu)?;
 
-        parser
-            .expect_keyword(Keyword::EXECUTE)
-            .context(SyntaxSnafu)?;
-        let stmt_name = parser.parse_identifier(false).context(SyntaxSnafu)?;
-        if parser.parse_keyword(Keyword::USING) {
-            let param_list = parser
-                .parse_comma_separated(Parser::parse_expr)
-                .context(SyntaxSnafu)?;
-            if !parser.consume_token(&Token::EOF) {
-                parser
-                    .expected("end of statement", parser.peek_token())
-                    .context(SyntaxSnafu)?;
-            }
-            Ok((stmt_name.value, param_list))
-        } else {
-            Ok((stmt_name.value, vec![]))
-        }
+        ParserContext { parser, sql }.parse_mysql_execute()
     }
 
     /// Raises an "unsupported statement" error.
@@ -314,6 +283,7 @@ impl<'a> ParserContext<'a> {
 mod tests {
 
     use datatypes::prelude::ConcreteDataType;
+    use sqlparser::dialect::MySqlDialect;
 
     use super::*;
     use crate::dialect::GreptimeDbDialect;
@@ -408,4 +378,46 @@ mod tests {
         assert_eq!(object_name.0.len(), 1);
         assert_eq!(object_name.to_string(), table_name.to_ascii_lowercase());
     }
+
+    #[test]
+    pub fn test_parse_mysql_prepare_stmt() {
+        let sql = "PREPARE stmt1 FROM 'SELECT * FROM t1 WHERE id = ?';";
+        let (stmt_name, stmt) =
+            ParserContext::parse_mysql_prepare_stmt(sql, &MySqlDialect {}).unwrap();
+        assert_eq!(stmt_name, "stmt1");
+        assert_eq!(stmt, "SELECT * FROM t1 WHERE id = ?");
+
+        let sql = "PREPARE stmt2 FROM \"SELECT * FROM t1 WHERE id = ?\"";
+        let (stmt_name, stmt) =
+            ParserContext::parse_mysql_prepare_stmt(sql, &MySqlDialect {}).unwrap();
+        assert_eq!(stmt_name, "stmt2");
+        assert_eq!(stmt, "SELECT * FROM t1 WHERE id = ?");
+    }
+
+    #[test]
+    pub fn test_parse_mysql_execute_stmt() {
+        let sql = "EXECUTE stmt1 USING 1, 'hello';";
+        let (stmt_name, params) =
+            ParserContext::parse_mysql_execute_stmt(sql, &GreptimeDbDialect {}).unwrap();
+        assert_eq!(stmt_name, "stmt1");
+        assert_eq!(params.len(), 2);
+        assert_eq!(params[0].to_string(), "1");
+        assert_eq!(params[1].to_string(), "'hello'");
+
+        let sql = "EXECUTE stmt2;";
+        let (stmt_name, params) =
+            ParserContext::parse_mysql_execute_stmt(sql, &GreptimeDbDialect {}).unwrap();
+        assert_eq!(stmt_name, "stmt2");
+        assert_eq!(params.len(), 0);
+
+        let sql = "EXECUTE stmt3 USING 231, 'hello', \"2003-03-1\", NULL, ;";
+        let (stmt_name, params) =
+            ParserContext::parse_mysql_execute_stmt(sql, &GreptimeDbDialect {}).unwrap();
+        assert_eq!(stmt_name, "stmt3");
+        assert_eq!(params.len(), 4);
+        assert_eq!(params[0].to_string(), "231");
+        assert_eq!(params[1].to_string(), "'hello'");
+        assert_eq!(params[2].to_string(), "\"2003-03-1\"");
+        assert_eq!(params[3].to_string(), "NULL");
+    }
 }

src/sql/src/parsers.rs

@@ -19,8 +19,10 @@ pub(crate) mod delete_parser;
 pub(crate) mod describe_parser;
 pub(crate) mod drop_parser;
 pub(crate) mod error;
+pub(crate) mod execute_parser;
 pub(crate) mod explain_parser;
 pub(crate) mod insert_parser;
+pub(crate) mod prepare_parser;
 pub(crate) mod query_parser;
 pub(crate) mod set_var_parser;
 pub(crate) mod show_parser;

src/sql/src/parsers/execute_parser.rs

@@ -0,0 +1,41 @@
+// 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 snafu::ResultExt;
+use sqlparser::ast::Expr;
+use sqlparser::keywords::Keyword;
+use sqlparser::parser::Parser;
+
+use crate::error::{Result, SyntaxSnafu};
+use crate::parser::ParserContext;
+
+impl<'a> ParserContext<'a> {
+    /// Parses MySQL style 'EXECUTE stmt_name USING param_list' into a stmt_name string and a list of parameters.
+    /// Only use for MySQL. for PostgreSQL, use `sqlparser::parser::Parser::parse_execute` instead.
+    pub(crate) fn parse_mysql_execute(&mut self) -> Result<(String, Vec<Expr>)> {
+        self.parser
+            .expect_keyword(Keyword::EXECUTE)
+            .context(SyntaxSnafu)?;
+        let stmt_name = self.parser.parse_identifier(false).context(SyntaxSnafu)?;
+        if self.parser.parse_keyword(Keyword::USING) {
+            let param_list = self
+                .parser
+                .parse_comma_separated(Parser::parse_expr)
+                .context(SyntaxSnafu)?;
+            Ok((stmt_name.value, param_list))
+        } else {
+            Ok((stmt_name.value, vec![]))
+        }
+    }
+}

src/sql/src/parsers/prepare_parser.rs

@@ -0,0 +1,46 @@
+// 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 snafu::ResultExt;
+use sqlparser::keywords::Keyword;
+use sqlparser::tokenizer::Token;
+
+use crate::error::{Result, SyntaxSnafu};
+use crate::parser::ParserContext;
+
+impl<'a> ParserContext<'a> {
+    /// Parses MySQL style 'PREPARE stmt_name FROM stmt' into a (stmt_name, stmt) tuple.
+    /// Only use for MySQL. for PostgreSQL, use `sqlparser::parser::Parser::parse_prepare` instead.
+    pub(crate) fn parse_mysql_prepare(&mut self) -> Result<(String, String)> {
+        self.parser
+            .expect_keyword(Keyword::PREPARE)
+            .context(SyntaxSnafu)?;
+        let stmt_name = self.parser.parse_identifier(false).context(SyntaxSnafu)?;
+        self.parser
+            .expect_keyword(Keyword::FROM)
+            .context(SyntaxSnafu)?;
+        let next_token = self.parser.peek_token();
+        let stmt = match next_token.token {
+            Token::SingleQuotedString(s) | Token::DoubleQuotedString(s) => {
+                let _ = self.parser.next_token();
+                s
+            }
+            _ => self
+                .parser
+                .expected("string literal", next_token)
+                .context(SyntaxSnafu)?,
+        };
+        Ok((stmt_name.value, stmt))
+    }
+}