Change USE Statement to use Enum

Author: kacpermuda

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,13 +2515,9 @@ pub enum Statement {
     /// Note: this is a MySQL-specific statement.
     ShowCollation { filter: Option<ShowStatementFilter> },
     /// ```sql
-    /// USE [DATABASE|SCHEMA|CATALOG|...] [<db_name>.<schema_name>|<db_name>|<schema_name>]
+    /// `USE ...`
     /// ```
-    Use {
-        db_name: Option<Ident>,
-        schema_name: Option<Ident>,
-        keyword: Option<String>,
-    },
+    Use(Use),
     /// ```sql
     /// START  [ TRANSACTION | WORK ] | START TRANSACTION } ...
     /// ```
@@ -4127,28 +4123,7 @@ impl fmt::Display for Statement {
                 }
                 Ok(())
             }
-            Statement::Use {
-                db_name,
-                schema_name,
-                keyword,
-            } => {
-                write!(f, "USE")?;
-
-                if let Some(kw) = keyword.as_ref() {
-                    write!(f, " {}", kw)?;
-                }
-
-                if let Some(db_name) = db_name {
-                    write!(f, " {}", db_name)?;
-                    if let Some(schema_name) = schema_name {
-                        write!(f, ".{}", schema_name)?;
-                    }
-                } else if let Some(schema_name) = schema_name {
-                    write!(f, " {}", schema_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

@@ -804,6 +804,7 @@ define_keywords!(
     VIEW,
     VIRTUAL,
     VOLATILE,
+    WAREHOUSE,
     WEEK,
     WHEN,
     WHENEVER,

src/parser/mod.rs

@@ -9225,63 +9225,31 @@ impl<'a> Parser<'a> {
     }
 
     pub fn parse_use(&mut self) -> Result<Statement, ParserError> {
-        // What should be treated as keyword in given dialect
-        let allowed_keywords = if dialect_of!(self is HiveDialect) {
-            vec![Keyword::DEFAULT]
+        // 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) {
-            vec![Keyword::CATALOG, Keyword::DATABASE, Keyword::SCHEMA]
+            self.parse_one_of_keywords(&[Keyword::CATALOG, Keyword::DATABASE, Keyword::SCHEMA])
         } else if dialect_of!(self is SnowflakeDialect) {
-            vec![Keyword::DATABASE, Keyword::SCHEMA]
+            self.parse_one_of_keywords(&[Keyword::DATABASE, Keyword::SCHEMA, Keyword::WAREHOUSE])
         } else {
-            vec![]
+            None // No specific keywords for other dialects, including GenericDialect
         };
-        let parsed_keyword = self.parse_one_of_keywords(&allowed_keywords);
 
-        // Hive dialect accepts USE DEFAULT; statement without any db specified
-        if dialect_of!(self is HiveDialect) && parsed_keyword == Some(Keyword::DEFAULT) {
-            return Ok(Statement::Use {
-                db_name: None,
-                schema_name: None,
-                keyword: Some("DEFAULT".to_string()),
-            });
-        }
-
-        // Parse the object name, which might be a single identifier or fully qualified name (e.g., x.y)
-        let parts = self.parse_object_name(false)?.0;
-        let (db_name, schema_name) = match parts.len() {
-            1 => {
-                // Single identifier found
-                if dialect_of!(self is DatabricksDialect) {
-                    if parsed_keyword == Some(Keyword::CATALOG) {
-                        // Databricks: CATALOG keyword provided, treat as database name
-                        (Some(parts[0].clone()), None)
-                    } else {
-                        // Databricks: DATABASE, SCHEMA or no keyword provided, treat as schema name
-                        (None, Some(parts[0].clone()))
-                    }
-                } else if dialect_of!(self is SnowflakeDialect)
-                    && parsed_keyword == Some(Keyword::SCHEMA)
-                {
-                    // Snowflake: SCHEMA keyword provided, treat as schema name
-                    (None, Some(parts[0].clone()))
-                } else {
-                    // Other dialects: treat as database name by default
-                    (Some(parts[0].clone()), None)
-                }
-            }
-            2 => (Some(parts[0].clone()), Some(parts[1].clone())),
-            _ => {
-                return Err(ParserError::ParserError(
-                    "Invalid format in the USE statement".to_string(),
-                ))
-            }
+        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 {
-            db_name,
-            schema_name,
-            keyword: parsed_keyword.map(|kw| format!("{:?}", kw)),
-        })
+        Ok(Statement::Use(result))
     }
 
     pub fn parse_table_and_joins(&mut self) -> Result<TableWithJoins, ParserError> {

tests/sqlparser_clickhouse.rs

@@ -1162,38 +1162,35 @@ fn test_prewhere() {
 
 #[test]
 fn parse_use() {
-    assert_eq!(
-        clickhouse().verified_stmt("USE mydb"),
-        Statement::Use {
-            db_name: Some(Ident::new("mydb")),
-            schema_name: None,
-            keyword: None
-        }
-    );
-    assert_eq!(
-        clickhouse().verified_stmt("USE DATABASE"),
-        Statement::Use {
-            db_name: Some(Ident::new("DATABASE")),
-            schema_name: None,
-            keyword: None
-        }
-    );
-    assert_eq!(
-        clickhouse().verified_stmt("USE SCHEMA"),
-        Statement::Use {
-            db_name: Some(Ident::new("SCHEMA")),
-            schema_name: None,
-            keyword: None
-        }
-    );
-    assert_eq!(
-        clickhouse().verified_stmt("USE CATALOG"),
-        Statement::Use {
-            db_name: Some(Ident::new("CATALOG")),
-            schema_name: None,
-            keyword: None
+    let valid_object_names = [
+        "mydb",
+        "SCHEMA",
+        "DATABASE",
+        "CATALOG",
+        "WAREHOUSE",
+        "DEFAULT",
+    ];
+    let quote_styles = ['"', '`'];
+
+    for object_name in &valid_object_names {
+        // Test single identifier without quotes
+        assert_eq!(
+            clickhouse().verified_stmt(&format!("USE {}", object_name)),
+            Statement::Use(Use::Object(ObjectName(vec![Ident::new(
+                object_name.to_string()
+            )])))
+        );
+        for &quote in &quote_styles {
+            // Test single identifier with different type of quotes
+            assert_eq!(
+                clickhouse().verified_stmt(&format!("USE {0}{1}{0}", quote, object_name)),
+                Statement::Use(Use::Object(ObjectName(vec![Ident::with_quote(
+                    quote,
+                    object_name.to_string(),
+                )])))
+            );
         }
-    );
+    }
 }
 
 #[test]

tests/sqlparser_databricks.rs

@@ -192,47 +192,69 @@ fn test_values_clause() {
 
 #[test]
 fn parse_use() {
-    assert_eq!(
-        databricks().verified_stmt("USE my_schema"),
-        Statement::Use {
-            db_name: None,
-            schema_name: Some(Ident::new("my_schema")),
-            keyword: None
+    let valid_object_names = ["mydb", "WAREHOUSE", "DEFAULT"];
+    let quote_styles = ['"', '`'];
+
+    for object_name in &valid_object_names {
+        // Test single identifier without quotes
+        assert_eq!(
+            databricks().verified_stmt(&format!("USE {}", object_name)),
+            Statement::Use(Use::Object(ObjectName(vec![Ident::new(
+                object_name.to_string()
+            )])))
+        );
+        for &quote in &quote_styles {
+            // Test single identifier with different type of quotes
+            assert_eq!(
+                databricks().verified_stmt(&format!("USE {0}{1}{0}", quote, object_name)),
+                Statement::Use(Use::Object(ObjectName(vec![Ident::with_quote(
+                    quote,
+                    object_name.to_string(),
+                )])))
+            );
         }
-    );
+    }
+
+    for &quote in &quote_styles {
+        // Test single identifier with keyword and different type of quotes
+        assert_eq!(
+            databricks().verified_stmt(&format!("USE CATALOG {0}my_catalog{0}", quote)),
+            Statement::Use(Use::Catalog(ObjectName(vec![Ident::with_quote(
+                quote,
+                "my_catalog".to_string(),
+            )])))
+        );
+        assert_eq!(
+            databricks().verified_stmt(&format!("USE DATABASE {0}my_database{0}", quote)),
+            Statement::Use(Use::Database(ObjectName(vec![Ident::with_quote(
+                quote,
+                "my_database".to_string(),
+            )])))
+        );
+        assert_eq!(
+            databricks().verified_stmt(&format!("USE SCHEMA {0}my_schema{0}", quote)),
+            Statement::Use(Use::Schema(ObjectName(vec![Ident::with_quote(
+                quote,
+                "my_schema".to_string(),
+            )])))
+        );
+    }
+
+    // Test single identifier with keyword and no quotes
     assert_eq!(
         databricks().verified_stmt("USE CATALOG my_catalog"),
-        Statement::Use {
-            db_name: Some(Ident::new("my_catalog")),
-            schema_name: None,
-            keyword: Some("CATALOG".to_string())
-        }
-    );
-    assert_eq!(
-        databricks().verified_stmt("USE CATALOG 'my_catalog'"),
-        Statement::Use {
-            db_name: Some(Ident::with_quote('\'', "my_catalog")),
-            schema_name: None,
-            keyword: Some("CATALOG".to_string())
-        }
+        Statement::Use(Use::Catalog(ObjectName(vec![Ident::new("my_catalog")])))
     );
     assert_eq!(
         databricks().verified_stmt("USE DATABASE my_schema"),
-        Statement::Use {
-            db_name: None,
-            schema_name: Some(Ident::new("my_schema")),
-            keyword: Some("DATABASE".to_string())
-        }
+        Statement::Use(Use::Database(ObjectName(vec![Ident::new("my_schema")])))
     );
     assert_eq!(
         databricks().verified_stmt("USE SCHEMA my_schema"),
-        Statement::Use {
-            db_name: None,
-            schema_name: Some(Ident::new("my_schema")),
-            keyword: Some("SCHEMA".to_string())
-        }
+        Statement::Use(Use::Schema(ObjectName(vec![Ident::new("my_schema")])))
     );
 
+    // Test invalid syntax - missing identifier
     let invalid_cases = ["USE SCHEMA", "USE DATABASE", "USE CATALOG"];
     for sql in &invalid_cases {
         assert_eq!(