feat: PREPARE and EXECUTE statement from mysql client (#4125)

* feat: prepare stmt in mysql client

* feat: execute stmt in mysql client

* fix: handle parameters properly

* refactor: use existing funcs to convert expr to scalar value

* refactor: use uuid strings as stmt_key for queries from COM_PREPARE packet

* refactor: take prepare and execute parser as submodule

* test: add unit test for converting expr to scalar value

* feat: deallocate stmt in mysql client

* chore: comments and duplicates

---------

Co-authored-by: dennis zhuang <killme2008@gmail.com>

src/sql/src/parser.rs

@@ -175,6 +175,27 @@ impl<'a> ParserContext<'a> {
         }
     }
 
+    /// Parses MySQL style 'PREPARE stmt_name FROM stmt' into a (stmt_name, stmt) tuple.
+    pub fn parse_mysql_prepare_stmt(
+        sql: &'a str,
+        dialect: &dyn Dialect,
+    ) -> Result<(String, String)> {
+        ParserContext::new(dialect, 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>)> {
+        ParserContext::new(dialect, sql)?.parse_mysql_execute()
+    }
+
+    /// Parses MySQL style 'DEALLOCATE stmt_name' into a stmt_name string.
+    pub fn parse_mysql_deallocate_stmt(sql: &'a str, dialect: &dyn Dialect) -> Result<String> {
+        ParserContext::new(dialect, sql)?.parse_deallocate()
+    }
+
     /// Raises an "unsupported statement" error.
     pub fn unsupported<T>(&self, keyword: String) -> Result<T> {
         error::UnsupportedSnafu {
@@ -257,6 +278,7 @@ impl<'a> ParserContext<'a> {
 mod tests {
 
     use datatypes::prelude::ConcreteDataType;
+    use sqlparser::dialect::MySqlDialect;
 
     use super::*;
     use crate::dialect::GreptimeDbDialect;
@@ -351,4 +373,57 @@ 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");
+    }
+
+    #[test]
+    pub fn test_parse_mysql_deallocate_stmt() {
+        let sql = "DEALLOCATE stmt1;";
+        let stmt_name = ParserContext::parse_mysql_deallocate_stmt(sql, &MySqlDialect {}).unwrap();
+        assert_eq!(stmt_name, "stmt1");
+
+        let sql = "DEALLOCATE stmt2";
+        let stmt_name = ParserContext::parse_mysql_deallocate_stmt(sql, &MySqlDialect {}).unwrap();
+        assert_eq!(stmt_name, "stmt2");
+    }
 }

src/sql/src/parsers.rs

@@ -15,12 +15,15 @@
 mod alter_parser;
 pub(crate) mod copy_parser;
 pub(crate) mod create_parser;
+pub(crate) mod deallocate_parser;
 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/deallocate_parser.rs

@@ -0,0 +1,30 @@
+// 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 crate::error::{Result, SyntaxSnafu};
+use crate::parser::ParserContext;
+
+impl<'a> ParserContext<'a> {
+    /// Parses MySQL style 'PREPARE stmt_name' into a stmt_name string.
+    pub(crate) fn parse_deallocate(&mut self) -> Result<String> {
+        self.parser
+            .expect_keyword(Keyword::DEALLOCATE)
+            .context(SyntaxSnafu)?;
+        let stmt_name = self.parser.parse_identifier(false).context(SyntaxSnafu)?;
+        Ok(stmt_name.value)
+    }
+}

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