Skip to content

Commit 7282ce2

Browse files
kacpermudakacpermuda
authored
feat: support different USE statement syntaxes (apache#1387)
1 parent 19e694a commit 7282ce2

File tree

11 files changed

+386
-17
lines changed

11 files changed

+386
-17
lines changed

‎src/ast/dcl.rs‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -193,3 +193,30 @@ impl fmt::Display for AlterRoleOperation {
193193
}
194194
}
195195
}
196+
197+
/// A `USE` (`Statement::Use`) operation
198+
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
199+
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
200+
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
201+
pub enum Use {
202+
Catalog(ObjectName), // e.g. `USE CATALOG foo.bar`
203+
Schema(ObjectName), // e.g. `USE SCHEMA foo.bar`
204+
Database(ObjectName), // e.g. `USE DATABASE foo.bar`
205+
Warehouse(ObjectName), // e.g. `USE WAREHOUSE foo.bar`
206+
Object(ObjectName), // e.g. `USE foo.bar`
207+
Default, // e.g. `USE DEFAULT`
208+
}
209+
210+
impl fmt::Display for Use {
211+
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
212+
f.write_str("USE ")?;
213+
match self {
214+
Use::Catalog(name) => write!(f, "CATALOG {}", name),
215+
Use::Schema(name) => write!(f, "SCHEMA {}", name),
216+
Use::Database(name) => write!(f, "DATABASE {}", name),
217+
Use::Warehouse(name) => write!(f, "WAREHOUSE {}", name),
218+
Use::Object(name) => write!(f, "{}", name),
219+
Use::Default => write!(f, "DEFAULT"),
220+
}
221+
}
222+
}

‎src/ast/mod.rs‎

Lines changed: 4 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ pub use self::data_type::{
3131
ArrayElemTypeDef, CharLengthUnits, CharacterLength, DataType, ExactNumberInfo,
3232
StructBracketKind, TimezoneInfo,
3333
};
34-
pub use self::dcl::{AlterRoleOperation, ResetConfig, RoleOption, SetConfigValue};
34+
pub use self::dcl::{AlterRoleOperation, ResetConfig, RoleOption, SetConfigValue, Use};
3535
pub use self::ddl::{
3636
AlterColumnOperation, AlterIndexOperation, AlterTableOperation, ColumnDef, ColumnOption,
3737
ColumnOptionDef, ConstraintCharacteristics, Deduplicate, DeferrableInitial, GeneratedAs,
@@ -2515,11 +2515,9 @@ pub enum Statement {
25152515
/// Note: this is a MySQL-specific statement.
25162516
ShowCollation { filter: Option<ShowStatementFilter> },
25172517
/// ```sql
2518-
/// USE
2518+
/// `USE ...`
25192519
/// ```
2520-
///
2521-
/// Note: This is a MySQL-specific statement.
2522-
Use { db_name: Ident },
2520+
Use(Use),
25232521
/// ```sql
25242522
/// START [ TRANSACTION | WORK ] | START TRANSACTION } ...
25252523
/// ```
@@ -4125,10 +4123,7 @@ impl fmt::Display for Statement {
41254123
}
41264124
Ok(())
41274125
}
4128-
Statement::Use { db_name } => {
4129-
write!(f, "USE {db_name}")?;
4130-
Ok(())
4131-
}
4126+
Statement::Use(use_expr) => use_expr.fmt(f),
41324127
Statement::ShowCollation { filter } => {
41334128
write!(f, "SHOW COLLATION")?;
41344129
if let Some(filter) = filter {

‎src/keywords.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -137,6 +137,7 @@ define_keywords!(
137137
CASCADED,
138138
CASE,
139139
CAST,
140+
CATALOG,
140141
CEIL,
141142
CEILING,
142143
CENTURY,
@@ -804,6 +805,7 @@ define_keywords!(
804805
VIEW,
805806
VIRTUAL,
806807
VOLATILE,
808+
WAREHOUSE,
807809
WEEK,
808810
WHEN,
809811
WHENEVER,

‎src/parser/mod.rs‎

Lines changed: 25 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9264,8 +9264,31 @@ impl<'a> Parser<'a> {
92649264
}
92659265

92669266
pub fn parse_use(&mut self) -> Result<Statement, ParserError> {
9267-
let db_name = self.parse_identifier(false)?;
9268-
Ok(Statement::Use { db_name })
9267+
// Determine which keywords are recognized by the current dialect
9268+
let parsed_keyword = if dialect_of!(self is HiveDialect) {
9269+
// HiveDialect accepts USE DEFAULT; statement without any db specified
9270+
if self.parse_keyword(Keyword::DEFAULT) {
9271+
return Ok(Statement::Use(Use::Default));
9272+
}
9273+
None // HiveDialect doesn't expect any other specific keyword after `USE`
9274+
} else if dialect_of!(self is DatabricksDialect) {
9275+
self.parse_one_of_keywords(&[Keyword::CATALOG, Keyword::DATABASE, Keyword::SCHEMA])
9276+
} else if dialect_of!(self is SnowflakeDialect) {
9277+
self.parse_one_of_keywords(&[Keyword::DATABASE, Keyword::SCHEMA, Keyword::WAREHOUSE])
9278+
} else {
9279+
None // No specific keywords for other dialects, including GenericDialect
9280+
};
9281+
9282+
let obj_name = self.parse_object_name(false)?;
9283+
let result = match parsed_keyword {
9284+
Some(Keyword::CATALOG) => Use::Catalog(obj_name),
9285+
Some(Keyword::DATABASE) => Use::Database(obj_name),
9286+
Some(Keyword::SCHEMA) => Use::Schema(obj_name),
9287+
Some(Keyword::WAREHOUSE) => Use::Warehouse(obj_name),
9288+
_ => Use::Object(obj_name),
9289+
};
9290+
9291+
Ok(Statement::Use(result))
92699292
}
92709293

92719294
pub fn parse_table_and_joins(&mut self) -> Result<TableWithJoins, ParserError> {

‎tests/sqlparser_clickhouse.rs‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1232,6 +1232,39 @@ fn test_prewhere() {
12321232
}
12331233
}
12341234

1235+
#[test]
1236+
fn parse_use() {
1237+
let valid_object_names = [
1238+
"mydb",
1239+
"SCHEMA",
1240+
"DATABASE",
1241+
"CATALOG",
1242+
"WAREHOUSE",
1243+
"DEFAULT",
1244+
];
1245+
let quote_styles = ['"', '`'];
1246+
1247+
for object_name in &valid_object_names {
1248+
// Test single identifier without quotes
1249+
assert_eq!(
1250+
clickhouse().verified_stmt(&format!("USE {}", object_name)),
1251+
Statement::Use(Use::Object(ObjectName(vec![Ident::new(
1252+
object_name.to_string()
1253+
)])))
1254+
);
1255+
for &quote in &quote_styles {
1256+
// Test single identifier with different type of quotes
1257+
assert_eq!(
1258+
clickhouse().verified_stmt(&format!("USE {0}{1}{0}", quote, object_name)),
1259+
Statement::Use(Use::Object(ObjectName(vec![Ident::with_quote(
1260+
quote,
1261+
object_name.to_string(),
1262+
)])))
1263+
);
1264+
}
1265+
}
1266+
}
1267+
12351268
#[test]
12361269
fn test_query_with_format_clause() {
12371270
let format_options = vec!["TabSeparated", "JSONCompact", "NULL"];

‎tests/sqlparser_databricks.rs‎

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -189,3 +189,77 @@ fn test_values_clause() {
189189
// TODO: support this example from https://docs.databricks.com/en/sql/language-manual/sql-ref-syntax-qry-select-values.html#examples
190190
// databricks().verified_query("VALUES 1, 2, 3");
191191
}
192+
193+
#[test]
194+
fn parse_use() {
195+
let valid_object_names = ["mydb", "WAREHOUSE", "DEFAULT"];
196+
let quote_styles = ['"', '`'];
197+
198+
for object_name in &valid_object_names {
199+
// Test single identifier without quotes
200+
assert_eq!(
201+
databricks().verified_stmt(&format!("USE {}", object_name)),
202+
Statement::Use(Use::Object(ObjectName(vec![Ident::new(
203+
object_name.to_string()
204+
)])))
205+
);
206+
for &quote in &quote_styles {
207+
// Test single identifier with different type of quotes
208+
assert_eq!(
209+
databricks().verified_stmt(&format!("USE {0}{1}{0}", quote, object_name)),
210+
Statement::Use(Use::Object(ObjectName(vec![Ident::with_quote(
211+
quote,
212+
object_name.to_string(),
213+
)])))
214+
);
215+
}
216+
}
217+
218+
for &quote in &quote_styles {
219+
// Test single identifier with keyword and different type of quotes
220+
assert_eq!(
221+
databricks().verified_stmt(&format!("USE CATALOG {0}my_catalog{0}", quote)),
222+
Statement::Use(Use::Catalog(ObjectName(vec![Ident::with_quote(
223+
quote,
224+
"my_catalog".to_string(),
225+
)])))
226+
);
227+
assert_eq!(
228+
databricks().verified_stmt(&format!("USE DATABASE {0}my_database{0}", quote)),
229+
Statement::Use(Use::Database(ObjectName(vec![Ident::with_quote(
230+
quote,
231+
"my_database".to_string(),
232+
)])))
233+
);
234+
assert_eq!(
235+
databricks().verified_stmt(&format!("USE SCHEMA {0}my_schema{0}", quote)),
236+
Statement::Use(Use::Schema(ObjectName(vec![Ident::with_quote(
237+
quote,
238+
"my_schema".to_string(),
239+
)])))
240+
);
241+
}
242+
243+
// Test single identifier with keyword and no quotes
244+
assert_eq!(
245+
databricks().verified_stmt("USE CATALOG my_catalog"),
246+
Statement::Use(Use::Catalog(ObjectName(vec![Ident::new("my_catalog")])))
247+
);
248+
assert_eq!(
249+
databricks().verified_stmt("USE DATABASE my_schema"),
250+
Statement::Use(Use::Database(ObjectName(vec![Ident::new("my_schema")])))
251+
);
252+
assert_eq!(
253+
databricks().verified_stmt("USE SCHEMA my_schema"),
254+
Statement::Use(Use::Schema(ObjectName(vec![Ident::new("my_schema")])))
255+
);
256+
257+
// Test invalid syntax - missing identifier
258+
let invalid_cases = ["USE SCHEMA", "USE DATABASE", "USE CATALOG"];
259+
for sql in &invalid_cases {
260+
assert_eq!(
261+
databricks().parse_sql_statements(sql).unwrap_err(),
262+
ParserError::ParserError("Expected: identifier, found: EOF".to_string()),
263+
);
264+
}
265+
}

‎tests/sqlparser_duckdb.rs‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -756,3 +756,55 @@ fn test_duckdb_union_datatype() {
756756
stmt
757757
);
758758
}
759+
760+
#[test]
761+
fn parse_use() {
762+
let valid_object_names = [
763+
"mydb",
764+
"SCHEMA",
765+
"DATABASE",
766+
"CATALOG",
767+
"WAREHOUSE",
768+
"DEFAULT",
769+
];
770+
let quote_styles = ['"', '\''];
771+
772+
for object_name in &valid_object_names {
773+
// Test single identifier without quotes
774+
assert_eq!(
775+
duckdb().verified_stmt(&format!("USE {}", object_name)),
776+
Statement::Use(Use::Object(ObjectName(vec![Ident::new(
777+
object_name.to_string()
778+
)])))
779+
);
780+
for &quote in &quote_styles {
781+
// Test single identifier with different type of quotes
782+
assert_eq!(
783+
duckdb().verified_stmt(&format!("USE {0}{1}{0}", quote, object_name)),
784+
Statement::Use(Use::Object(ObjectName(vec![Ident::with_quote(
785+
quote,
786+
object_name.to_string(),
787+
)])))
788+
);
789+
}
790+
}
791+
792+
for &quote in &quote_styles {
793+
// Test double identifier with different type of quotes
794+
assert_eq!(
795+
duckdb().verified_stmt(&format!("USE {0}CATALOG{0}.{0}my_schema{0}", quote)),
796+
Statement::Use(Use::Object(ObjectName(vec![
797+
Ident::with_quote(quote, "CATALOG"),
798+
Ident::with_quote(quote, "my_schema")
799+
])))
800+
);
801+
}
802+
// Test double identifier without quotes
803+
assert_eq!(
804+
duckdb().verified_stmt("USE mydb.my_schema"),
805+
Statement::Use(Use::Object(ObjectName(vec![
806+
Ident::new("mydb"),
807+
Ident::new("my_schema")
808+
])))
809+
);
810+
}

‎tests/sqlparser_hive.rs‎

Lines changed: 31 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818
use sqlparser::ast::{
1919
CreateFunctionBody, CreateFunctionUsing, Expr, Function, FunctionArgumentList,
2020
FunctionArguments, Ident, ObjectName, OneOrManyWithParens, SelectItem, Statement, TableFactor,
21-
UnaryOperator, Value,
21+
UnaryOperator, Use, Value,
2222
};
2323
use sqlparser::dialect::{GenericDialect, HiveDialect, MsSqlDialect};
2424
use sqlparser::parser::ParserError;
@@ -401,6 +401,36 @@ fn parse_delimited_identifiers() {
401401
//TODO verified_stmt(r#"UPDATE foo SET "bar" = 5"#);
402402
}
403403

404+
#[test]
405+
fn parse_use() {
406+
let valid_object_names = ["mydb", "SCHEMA", "DATABASE", "CATALOG", "WAREHOUSE"];
407+
let quote_styles = ['\'', '"', '`'];
408+
for object_name in &valid_object_names {
409+
// Test single identifier without quotes
410+
assert_eq!(
411+
hive().verified_stmt(&format!("USE {}", object_name)),
412+
Statement::Use(Use::Object(ObjectName(vec![Ident::new(
413+
object_name.to_string()
414+
)])))
415+
);
416+
for &quote in &quote_styles {
417+
// Test single identifier with different type of quotes
418+
assert_eq!(
419+
hive().verified_stmt(&format!("USE {}{}{}", quote, object_name, quote)),
420+
Statement::Use(Use::Object(ObjectName(vec![Ident::with_quote(
421+
quote,
422+
object_name.to_string(),
423+
)])))
424+
);
425+
}
426+
}
427+
// Test DEFAULT keyword that is special case in Hive
428+
assert_eq!(
429+
hive().verified_stmt("USE DEFAULT"),
430+
Statement::Use(Use::Default)
431+
);
432+
}
433+
404434
fn hive() -> TestedDialects {
405435
TestedDialects {
406436
dialects: vec![Box::new(HiveDialect {})],

‎tests/sqlparser_mssql.rs‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -621,6 +621,38 @@ fn parse_mssql_declare() {
621621
);
622622
}
623623

624+
#[test]
625+
fn parse_use() {
626+
let valid_object_names = [
627+
"mydb",
628+
"SCHEMA",
629+
"DATABASE",
630+
"CATALOG",
631+
"WAREHOUSE",
632+
"DEFAULT",
633+
];
634+
let quote_styles = ['\'', '"'];
635+
for object_name in &valid_object_names {
636+
// Test single identifier without quotes
637+
assert_eq!(
638+
ms().verified_stmt(&format!("USE {}", object_name)),
639+
Statement::Use(Use::Object(ObjectName(vec![Ident::new(
640+
object_name.to_string()
641+
)])))
642+
);
643+
for &quote in &quote_styles {
644+
// Test single identifier with different type of quotes
645+
assert_eq!(
646+
ms().verified_stmt(&format!("USE {}{}{}", quote, object_name, quote)),
647+
Statement::Use(Use::Object(ObjectName(vec![Ident::with_quote(
648+
quote,
649+
object_name.to_string(),
650+
)])))
651+
);
652+
}
653+
}
654+
}
655+
624656
fn ms() -> TestedDialects {
625657
TestedDialects {
626658
dialects: vec![Box::new(MsSqlDialect {})],

0 commit comments

Comments
 (0)