From a9db6ed1395066c500ba2a8eddb4f5cca736a96a Mon Sep 17 00:00:00 2001
From: Alex Qyoun-ae <4062971+MazterQyou@users.noreply.github.com>
Date: Thu, 11 Aug 2022 14:49:20 +0400
Subject: [PATCH] Support `USE db` (#565)

---
 src/ast/mod.rs           |  8 ++++++++
 src/keywords.rs          |  1 +
 src/parser.rs            |  6 ++++++
 tests/sqlparser_mysql.rs | 10 ++++++++++
 4 files changed, 25 insertions(+)

diff --git a/src/ast/mod.rs b/src/ast/mod.rs
index e45ebfcd..97b647bc 100644
--- a/src/ast/mod.rs
+++ b/src/ast/mod.rs
@@ -1016,6 +1016,10 @@ pub enum Statement {
         table_name: ObjectName,
         filter: Option<ShowStatementFilter>,
     },
+    /// USE
+    ///
+    /// Note: This is a MySQL-specific statement.
+    Use { db_name: Ident },
     /// `{ BEGIN [ TRANSACTION | WORK ] | START TRANSACTION } ...`
     StartTransaction { modes: Vec<TransactionMode> },
     /// `SET TRANSACTION ...`
@@ -1825,6 +1829,10 @@ impl fmt::Display for Statement {
                 }
                 Ok(())
             }
+            Statement::Use { db_name } => {
+                write!(f, "USE {}", db_name)?;
+                Ok(())
+            }
             Statement::StartTransaction { modes } => {
                 write!(f, "START TRANSACTION")?;
                 if !modes.is_empty() {
diff --git a/src/keywords.rs b/src/keywords.rs
index af8bf2f5..7eb6b420 100644
--- a/src/keywords.rs
+++ b/src/keywords.rs
@@ -535,6 +535,7 @@ define_keywords!(
     UPDATE,
     UPPER,
     USAGE,
+    USE,
     USER,
     USING,
     UUID,
diff --git a/src/parser.rs b/src/parser.rs
index e0324182..c59c88e8 100644
--- a/src/parser.rs
+++ b/src/parser.rs
@@ -177,6 +177,7 @@ impl<'a> Parser<'a> {
                 Keyword::CLOSE => Ok(self.parse_close()?),
                 Keyword::SET => Ok(self.parse_set()?),
                 Keyword::SHOW => Ok(self.parse_show()?),
+                Keyword::USE => Ok(self.parse_use()?),
                 Keyword::GRANT => Ok(self.parse_grant()?),
                 Keyword::REVOKE => Ok(self.parse_revoke()?),
                 Keyword::START => Ok(self.parse_start_transaction()?),
@@ -3776,6 +3777,11 @@ impl<'a> Parser<'a> {
         }
     }
 
+    pub fn parse_use(&mut self) -> Result<Statement, ParserError> {
+        let db_name = self.parse_identifier()?;
+        Ok(Statement::Use { db_name })
+    }
+
     pub fn parse_table_and_joins(&mut self) -> Result<TableWithJoins, ParserError> {
         let relation = self.parse_table_factor()?;
         // Note that for keywords to be properly handled here, they need to be
diff --git a/tests/sqlparser_mysql.rs b/tests/sqlparser_mysql.rs
index 9ac775bd..ab28b879 100644
--- a/tests/sqlparser_mysql.rs
+++ b/tests/sqlparser_mysql.rs
@@ -138,6 +138,16 @@ fn parse_show_create() {
     }
 }
 
+#[test]
+fn parse_use() {
+    assert_eq!(
+        mysql_and_generic().verified_stmt("USE mydb"),
+        Statement::Use {
+            db_name: Ident::new("mydb")
+        }
+    );
+}
+
 #[test]
 fn parse_create_table_auto_increment() {
     let sql = "CREATE TABLE foo (bar INT PRIMARY KEY AUTO_INCREMENT)";