feat: support different USE statement syntaxes (#1387)

src/ast/dcl.rs

@@ -193,3 +193,30 @@ impl fmt::Display for AlterRoleOperation {
         }
     }
 }
+
+/// A `USE` (`Statement::Use`) operation
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum Use {
+    Catalog(ObjectName),   // e.g. `USE CATALOG foo.bar`
+    Schema(ObjectName),    // e.g. `USE SCHEMA foo.bar`
+    Database(ObjectName),  // e.g. `USE DATABASE foo.bar`
+    Warehouse(ObjectName), // e.g. `USE WAREHOUSE foo.bar`
+    Object(ObjectName),    // e.g. `USE foo.bar`
+    Default,               // e.g. `USE DEFAULT`
+}
+
+impl fmt::Display for Use {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        f.write_str("USE ")?;
+        match self {
+            Use::Catalog(name) => write!(f, "CATALOG {}", name),
+            Use::Schema(name) => write!(f, "SCHEMA {}", name),
+            Use::Database(name) => write!(f, "DATABASE {}", name),
+            Use::Warehouse(name) => write!(f, "WAREHOUSE {}", name),
+            Use::Object(name) => write!(f, "{}", name),
+            Use::Default => write!(f, "DEFAULT"),
+        }
+    }
+}

src/ast/mod.rs

@@ -31,7 +31,7 @@ pub use self::data_type::{
     ArrayElemTypeDef, CharLengthUnits, CharacterLength, DataType, ExactNumberInfo,
     StructBracketKind, TimezoneInfo,
 };
-pub use self::dcl::{AlterRoleOperation, ResetConfig, RoleOption, SetConfigValue};
+pub use self::dcl::{AlterRoleOperation, ResetConfig, RoleOption, SetConfigValue, Use};
 pub use self::ddl::{
     AlterColumnOperation, AlterIndexOperation, AlterTableOperation, ColumnDef, ColumnOption,
     ColumnOptionDef, ConstraintCharacteristics, Deduplicate, DeferrableInitial, GeneratedAs,
@@ -2515,11 +2515,9 @@ pub enum Statement {
     /// Note: this is a MySQL-specific statement.
     ShowCollation { filter: Option<ShowStatementFilter> },
     /// ```sql
-    /// USE
+    /// `USE ...`
     /// ```
-    ///
-    /// Note: This is a MySQL-specific statement.
-    Use { db_name: Ident },
+    Use(Use),
     /// ```sql
     /// START  [ TRANSACTION | WORK ] | START TRANSACTION } ...
     /// ```
@@ -4125,10 +4123,7 @@ impl fmt::Display for Statement {
                 }
                 Ok(())
             }
-            Statement::Use { db_name } => {
-                write!(f, "USE {db_name}")?;
-                Ok(())
-            }
+            Statement::Use(use_expr) => use_expr.fmt(f),
             Statement::ShowCollation { filter } => {
                 write!(f, "SHOW COLLATION")?;
                 if let Some(filter) = filter {

src/keywords.rs

@@ -137,6 +137,7 @@ define_keywords!(
     CASCADED,
     CASE,
     CAST,
+    CATALOG,
     CEIL,
     CEILING,
     CENTURY,
@@ -804,6 +805,7 @@ define_keywords!(
     VIEW,
     VIRTUAL,
     VOLATILE,
+    WAREHOUSE,
     WEEK,
     WHEN,
     WHENEVER,

src/parser/mod.rs

@@ -9264,8 +9264,31 @@ impl<'a> Parser<'a> {
     }
 
     pub fn parse_use(&mut self) -> Result<Statement, ParserError> {
-        let db_name = self.parse_identifier(false)?;
-        Ok(Statement::Use { db_name })
+        // Determine which keywords are recognized by the current dialect
+        let parsed_keyword = if dialect_of!(self is HiveDialect) {
+            // HiveDialect accepts USE DEFAULT; statement without any db specified
+            if self.parse_keyword(Keyword::DEFAULT) {
+                return Ok(Statement::Use(Use::Default));
+            }
+            None // HiveDialect doesn't expect any other specific keyword after `USE`
+        } else if dialect_of!(self is DatabricksDialect) {
+            self.parse_one_of_keywords(&[Keyword::CATALOG, Keyword::DATABASE, Keyword::SCHEMA])
+        } else if dialect_of!(self is SnowflakeDialect) {
+            self.parse_one_of_keywords(&[Keyword::DATABASE, Keyword::SCHEMA, Keyword::WAREHOUSE])
+        } else {
+            None // No specific keywords for other dialects, including GenericDialect
+        };
+
+        let obj_name = self.parse_object_name(false)?;
+        let result = match parsed_keyword {
+            Some(Keyword::CATALOG) => Use::Catalog(obj_name),
+            Some(Keyword::DATABASE) => Use::Database(obj_name),
+            Some(Keyword::SCHEMA) => Use::Schema(obj_name),
+            Some(Keyword::WAREHOUSE) => Use::Warehouse(obj_name),
+            _ => Use::Object(obj_name),
+        };
+
+        Ok(Statement::Use(result))
     }
 
     pub fn parse_table_and_joins(&mut self) -> Result<TableWithJoins, ParserError> {