From 0835f45596edc08a6fa6f9106c140d236c23f971 Mon Sep 17 00:00:00 2001
From: Alex Qyoun-ae <4062971+MazterQyou@users.noreply.github.com>
Date: Wed, 10 Aug 2022 03:58:32 +0400
Subject: [PATCH] Support `USE db`

---
 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 312a86813..e43f01814 100644
--- a/src/ast/mod.rs
+++ b/src/ast/mod.rs
@@ -1012,6 +1012,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 ...`
@@ -1814,6 +1818,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 7643298e4..360b4b9ee 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 76190feb9..2f2c401b9 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()?),
@@ -3763,6 +3764,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 b3c02de30..158238928 100644
--- a/tests/sqlparser_mysql.rs
+++ b/tests/sqlparser_mysql.rs
@@ -140,6 +140,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)";