mirrors/datafusion-sqlparse
1
0
Fork 0
mirror of https://github.com/apache/datafusion-sqlparser-rs.git synced 2025-08-23 23:44:07 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

Support CREATE/DROP SECRET for duckdb dialect (#1208)

Browse source 
Co-authored-by: Jichao Sun <4977515+JichaoS@users.noreply.github.com>
This commit is contained in:
Andrew Lamb 2024-04-09 16:21:08 -04:00 • committed by GitHub
parent 8f67d1a713
commit 241da85d67
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 486 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

169
src/ast/mod.rs
View file

@ -2040,6 +2040,19 @@ pub enum Statement {
authorization_owner: Option<ObjectName>,
},
/// ```sql
/// CREATE SECRET
/// ```
/// See [duckdb](https://duckdb.org/docs/sql/statements/create_secret.html)
CreateSecret {
or_replace: bool,
temporary: Option<bool>,
if_not_exists: bool,
name: Option<Ident>,
storage_specifier: Option<Ident>,
secret_type: Ident,
options: Vec<SecretOption>,
},
/// ```sql
/// ALTER TABLE
/// ```
AlterTable {
@ -2088,6 +2101,31 @@ pub enum Statement {
/// true if the syntax is 'ATTACH DATABASE', false if it's just 'ATTACH'
database: bool,
},
/// (DuckDB-specific)
/// ```sql
/// ATTACH 'sqlite_file.db' AS sqlite_db (READ_ONLY, TYPE SQLITE);
/// ```
/// See <https://duckdb.org/docs/sql/statements/attach.html>
AttachDuckDBDatabase {
if_not_exists: bool,
/// true if the syntax is 'ATTACH DATABASE', false if it's just 'ATTACH'
database: bool,
/// An expression that indicates the path to the database file
database_path: Ident,
database_alias: Option<Ident>,
attach_options: Vec<AttachDuckDBDatabaseOption>,
},
/// (DuckDB-specific)
/// ```sql
/// DETACH db_alias;
/// ```
/// See <https://duckdb.org/docs/sql/statements/attach.html>
DetachDuckDBDatabase {
if_exists: bool,
/// true if the syntax is 'DETACH DATABASE', false if it's just 'DETACH'
database: bool,
database_alias: Ident,
},
/// ```sql
/// DROP [TABLE, VIEW, ...]
/// ```
@ -2121,6 +2159,15 @@ pub enum Statement {
option: Option<ReferentialAction>,
},
/// ```sql
/// DROP SECRET
/// ```
DropSecret {
if_exists: bool,
temporary: Option<bool>,
name: Ident,
storage_specifier: Option<Ident>,
},
/// ```sql
/// DECLARE
/// ```
/// Declare Cursor Variables
@ -2772,6 +2819,40 @@ impl fmt::Display for Statement {
let keyword = if *database { "DATABASE " } else { "" };
write!(f, "ATTACH {keyword}{database_file_name} AS {schema_name}")
}
Statement::AttachDuckDBDatabase {
if_not_exists,
database,
database_path,
database_alias,
attach_options,
} => {
write!(
f,
"ATTACH{database}{if_not_exists} {database_path}",
database = if *database { " DATABASE" } else { "" },
if_not_exists = if *if_not_exists { " IF NOT EXISTS" } else { "" },
)?;
if let Some(alias) = database_alias {
write!(f, " AS {alias}")?;
}
if !attach_options.is_empty() {
write!(f, " ({})", display_comma_separated(attach_options))?;
}
Ok(())
}
Statement::DetachDuckDBDatabase {
if_exists,
database,
database_alias,
} => {
write!(
f,
"DETACH{database}{if_exists} {database_alias}",
database = if *database { " DATABASE" } else { "" },
if_exists = if *if_exists { " IF EXISTS" } else { "" },
)?;
Ok(())
}
Statement::Analyze {
table_name,
partitions,
@ -3556,6 +3637,41 @@ impl fmt::Display for Statement {
}
Ok(())
}
Statement::CreateSecret {
or_replace,
temporary,
if_not_exists,
name,
storage_specifier,
secret_type,
options,
} => {
write!(
f,
"CREATE {or_replace}",
or_replace = if *or_replace { "OR REPLACE " } else { "" },
)?;
if let Some(t) = temporary {
write!(f, "{}", if *t { "TEMPORARY " } else { "PERSISTENT " })?;
}
write!(
f,
"SECRET {if_not_exists}",
if_not_exists = if *if_not_exists { "IF NOT EXISTS " } else { "" },
)?;
if let Some(n) = name {
write!(f, "{n} ")?;
};
if let Some(s) = storage_specifier {
write!(f, "IN {s} ")?;
}
write!(f, "( TYPE {secret_type}",)?;
if !options.is_empty() {
write!(f, ", {o}", o = display_comma_separated(options))?;
}
write!(f, " )")?;
Ok(())
}
Statement::AlterTable {
name,
if_exists,
@ -3636,6 +3752,26 @@ impl fmt::Display for Statement {
}
Ok(())
}
Statement::DropSecret {
if_exists,
temporary,
name,
storage_specifier,
} => {
write!(f, "DROP ")?;
if let Some(t) = temporary {
write!(f, "{}", if *t { "TEMPORARY " } else { "PERSISTENT " })?;
}
write!(
f,
"SECRET {if_exists}{name}",
if_exists = if *if_exists { "IF EXISTS " } else { "" },
)?;
if let Some(s) = storage_specifier {
write!(f, " FROM {s}")?;
}
Ok(())
}
Statement::Discard { object_type } => {
write!(f, "DISCARD {object_type}")?;
Ok(())
@ -5070,6 +5206,39 @@ impl fmt::Display for SqlOption {
}
}
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
pub struct SecretOption {
pub key: Ident,
pub value: Ident,
}
impl fmt::Display for SecretOption {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "{} {}", self.key, self.value)
}
}
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
pub enum AttachDuckDBDatabaseOption {
ReadOnly(Option<bool>),
Type(Ident),
}
impl fmt::Display for AttachDuckDBDatabaseOption {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match self {
AttachDuckDBDatabaseOption::ReadOnly(Some(true)) => write!(f, "READ_ONLY true"),
AttachDuckDBDatabaseOption::ReadOnly(Some(false)) => write!(f, "READ_ONLY false"),
AttachDuckDBDatabaseOption::ReadOnly(None) => write!(f, "READ_ONLY"),
AttachDuckDBDatabaseOption::Type(t) => write!(f, "TYPE {}", t),
}
}
}
#[derive(Debug, Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]

4
src/keywords.rs
View file

@ -224,6 +224,7 @@ define_keywords!(
DEREF,
DESC,
DESCRIBE,
DETACH,
DETAIL,
DETERMINISTIC,
DIRECTORY,
@ -514,6 +515,7 @@ define_keywords!(
PERCENTILE_DISC,
PERCENT_RANK,
PERIOD,
PERSISTENT,
PIVOT,
PLACING,
PLANS,
@ -543,6 +545,7 @@ define_keywords!(
RCFILE,
READ,
READS,
READ_ONLY,
REAL,
RECURSIVE,
REF,
@ -601,6 +604,7 @@ define_keywords!(
SCROLL,
SEARCH,
SECOND,
SECRET,
SECURITY,
SELECT,
SEMI,

174
src/parser/mod.rs
View file

@ -473,7 +473,16 @@ impl<'a> Parser<'a> {
Ok(Statement::Query(self.parse_boxed_query()?))
}
Keyword::TRUNCATE => Ok(self.parse_truncate()?),
Keyword::ATTACH => Ok(self.parse_attach_database()?),
Keyword::ATTACH => {
if dialect_of!(self is DuckDbDialect) {
Ok(self.parse_attach_duckdb_database()?)
} else {
Ok(self.parse_attach_database()?)
}
}
Keyword::DETACH if dialect_of!(self is DuckDbDialect | GenericDialect) => {
Ok(self.parse_detach_duckdb_database()?)
}
Keyword::MSCK => Ok(self.parse_msck()?),
Keyword::CREATE => Ok(self.parse_create()?),
Keyword::CACHE => Ok(self.parse_cache_table()?),
@ -666,6 +675,72 @@ impl<'a> Parser<'a> {
})
}
pub fn parse_attach_duckdb_database_options(
&mut self,
) -> Result<Vec<AttachDuckDBDatabaseOption>, ParserError> {
if !self.consume_token(&Token::LParen) {
return Ok(vec![]);
}
let mut options = vec![];
loop {
if self.parse_keyword(Keyword::READ_ONLY) {
let boolean = if self.parse_keyword(Keyword::TRUE) {
Some(true)
} else if self.parse_keyword(Keyword::FALSE) {
Some(false)
} else {
None
};
options.push(AttachDuckDBDatabaseOption::ReadOnly(boolean));
} else if self.parse_keyword(Keyword::TYPE) {
let ident = self.parse_identifier(false)?;
options.push(AttachDuckDBDatabaseOption::Type(ident));
} else {
return self.expected("expected one of: ), READ_ONLY, TYPE", self.peek_token());
};
if self.consume_token(&Token::RParen) {
return Ok(options);
} else if self.consume_token(&Token::Comma) {
continue;
} else {
return self.expected("expected one of: ')', ','", self.peek_token());
}
}
}
pub fn parse_attach_duckdb_database(&mut self) -> Result<Statement, ParserError> {
let database = self.parse_keyword(Keyword::DATABASE);
let if_not_exists = self.parse_keywords(&[Keyword::IF, Keyword::NOT, Keyword::EXISTS]);
let database_path = self.parse_identifier(false)?;
let database_alias = if self.parse_keyword(Keyword::AS) {
Some(self.parse_identifier(false)?)
} else {
None
};
let attach_options = self.parse_attach_duckdb_database_options()?;
Ok(Statement::AttachDuckDBDatabase {
if_not_exists,
database,
database_path,
database_alias,
attach_options,
})
}
pub fn parse_detach_duckdb_database(&mut self) -> Result<Statement, ParserError> {
let database = self.parse_keyword(Keyword::DATABASE);
let if_exists = self.parse_keywords(&[Keyword::IF, Keyword::EXISTS]);
let database_alias = self.parse_identifier(false)?;
Ok(Statement::DetachDuckDBDatabase {
if_exists,
database,
database_alias,
})
}
pub fn parse_attach_database(&mut self) -> Result<Statement, ParserError> {
let database = self.parse_keyword(Keyword::DATABASE);
let database_file_name = self.parse_expr()?;
@ -3075,6 +3150,8 @@ impl<'a> Parser<'a> {
let temporary = self
.parse_one_of_keywords(&[Keyword::TEMP, Keyword::TEMPORARY])
.is_some();
let persistent = dialect_of!(self is DuckDbDialect)
&& self.parse_one_of_keywords(&[Keyword::PERSISTENT]).is_some();
if self.parse_keyword(Keyword::TABLE) {
self.parse_create_table(or_replace, temporary, global, transient)
} else if self.parse_keyword(Keyword::MATERIALIZED) || self.parse_keyword(Keyword::VIEW) {
@ -3086,6 +3163,8 @@ impl<'a> Parser<'a> {
self.parse_create_function(or_replace, temporary)
} else if self.parse_keyword(Keyword::MACRO) {
self.parse_create_macro(or_replace, temporary)
} else if self.parse_keyword(Keyword::SECRET) {
self.parse_create_secret(or_replace, temporary, persistent)
} else if or_replace {
self.expected(
"[EXTERNAL] TABLE or [MATERIALIZED] VIEW or FUNCTION after CREATE OR REPLACE",
@ -3116,6 +3195,65 @@ impl<'a> Parser<'a> {
}
}
/// See [DuckDB Docs](https://duckdb.org/docs/sql/statements/create_secret.html) for more details.
pub fn parse_create_secret(
&mut self,
or_replace: bool,
temporary: bool,
persistent: bool,
) -> Result<Statement, ParserError> {
let if_not_exists = self.parse_keywords(&[Keyword::IF, Keyword::NOT, Keyword::EXISTS]);
let mut storage_specifier = None;
let mut name = None;
if self.peek_token() != Token::LParen {
if self.parse_keyword(Keyword::IN) {
storage_specifier = self.parse_identifier(false).ok()
} else {
name = self.parse_identifier(false).ok();
}
// Storage specifier may follow the name
if storage_specifier.is_none()
&& self.peek_token() != Token::LParen
&& self.parse_keyword(Keyword::IN)
{
storage_specifier = self.parse_identifier(false).ok();
}
}
self.expect_token(&Token::LParen)?;
self.expect_keyword(Keyword::TYPE)?;
let secret_type = self.parse_identifier(false)?;
let mut options = Vec::new();
if self.consume_token(&Token::Comma) {
options.append(&mut self.parse_comma_separated(|p| {
let key = p.parse_identifier(false)?;
let value = p.parse_identifier(false)?;
Ok(SecretOption { key, value })
})?);
}
self.expect_token(&Token::RParen)?;
let temp = match (temporary, persistent) {
(true, false) => Some(true),
(false, true) => Some(false),
(false, false) => None,
_ => self.expected("TEMPORARY or PERSISTENT", self.peek_token())?,
};
Ok(Statement::CreateSecret {
or_replace,
temporary: temp,
if_not_exists,
name,
storage_specifier,
secret_type,
options,
})
}
/// Parse a CACHE TABLE statement
pub fn parse_cache_table(&mut self) -> Result<Statement, ParserError> {
let (mut table_flag, mut options, mut has_as, mut query) = (None, vec![], false, None);
@ -3889,8 +4027,10 @@ impl<'a> Parser<'a> {
pub fn parse_drop(&mut self) -> Result<Statement, ParserError> {
// MySQL dialect supports `TEMPORARY`
let temporary = dialect_of!(self is MySqlDialect | GenericDialect)
let temporary = dialect_of!(self is MySqlDialect | GenericDialect | DuckDbDialect)
&& self.parse_keyword(Keyword::TEMPORARY);
let persistent = dialect_of!(self is DuckDbDialect)
&& self.parse_one_of_keywords(&[Keyword::PERSISTENT]).is_some();
let object_type = if self.parse_keyword(Keyword::TABLE) {
ObjectType::Table
@ -3908,6 +4048,8 @@ impl<'a> Parser<'a> {
ObjectType::Stage
} else if self.parse_keyword(Keyword::FUNCTION) {
return self.parse_drop_function();
} else if self.parse_keyword(Keyword::SECRET) {
return self.parse_drop_secret(temporary, persistent);
} else {
return self.expected(
"TABLE, VIEW, INDEX, ROLE, SCHEMA, FUNCTION, STAGE or SEQUENCE after DROP",
@ -3980,6 +4122,34 @@ impl<'a> Parser<'a> {
Ok(DropFunctionDesc { name, args })
}
/// See [DuckDB Docs](https://duckdb.org/docs/sql/statements/create_secret.html) for more details.
fn parse_drop_secret(
&mut self,
temporary: bool,
persistent: bool,
) -> Result<Statement, ParserError> {
let if_exists = self.parse_keywords(&[Keyword::IF, Keyword::EXISTS]);
let name = self.parse_identifier(false)?;
let storage_specifier = if self.parse_keyword(Keyword::FROM) {
self.parse_identifier(false).ok()
} else {
None
};
let temp = match (temporary, persistent) {
(true, false) => Some(true),
(false, true) => Some(false),
(false, false) => None,
_ => self.expected("TEMPORARY or PERSISTENT", self.peek_token())?,
};
Ok(Statement::DropSecret {
if_exists,
temporary: temp,
name,
storage_specifier,
})
}
/// Parse a `DECLARE` statement.
///
/// ```sql

141
tests/sqlparser_duckdb.rs
View file

@ -334,6 +334,147 @@ fn test_duckdb_struct_literal() {
);
}
#[test]
fn test_create_secret() {
let sql = r#"CREATE OR REPLACE PERSISTENT SECRET IF NOT EXISTS name IN storage ( TYPE type, key1 value1, key2 value2 )"#;
let stmt = duckdb().verified_stmt(sql);
assert_eq!(
Statement::CreateSecret {
or_replace: true,
temporary: Some(false),
if_not_exists: true,
name: Some(Ident::new("name")),
storage_specifier: Some(Ident::new("storage")),
secret_type: Ident::new("type"),
options: vec![
SecretOption {
key: Ident::new("key1"),
value: Ident::new("value1"),
},
SecretOption {
key: Ident::new("key2"),
value: Ident::new("value2"),
}
]
},
stmt
);
}
#[test]
fn test_create_secret_simple() {
let sql = r#"CREATE SECRET ( TYPE type )"#;
let stmt = duckdb().verified_stmt(sql);
assert_eq!(
Statement::CreateSecret {
or_replace: false,
temporary: None,
if_not_exists: false,
name: None,
storage_specifier: None,
secret_type: Ident::new("type"),
options: vec![]
},
stmt
);
}
#[test]
fn test_drop_secret() {
let sql = r#"DROP PERSISTENT SECRET IF EXISTS secret FROM storage"#;
let stmt = duckdb().verified_stmt(sql);
assert_eq!(
Statement::DropSecret {
if_exists: true,
temporary: Some(false),
name: Ident::new("secret"),
storage_specifier: Some(Ident::new("storage"))
},
stmt
);
}
#[test]
fn test_drop_secret_simple() {
let sql = r#"DROP SECRET secret"#;
let stmt = duckdb().verified_stmt(sql);
assert_eq!(
Statement::DropSecret {
if_exists: false,
temporary: None,
name: Ident::new("secret"),
storage_specifier: None
},
stmt
);
}
#[test]
fn test_attach_database() {
let sql = r#"ATTACH DATABASE IF NOT EXISTS 'sqlite_file.db' AS sqlite_db (READ_ONLY false, TYPE SQLITE)"#;
let stmt = duckdb().verified_stmt(sql);
assert_eq!(
Statement::AttachDuckDBDatabase {
if_not_exists: true,
database: true,
database_path: Ident::with_quote('\'', "sqlite_file.db"),
database_alias: Some(Ident::new("sqlite_db")),
attach_options: vec![
AttachDuckDBDatabaseOption::ReadOnly(Some(false)),
AttachDuckDBDatabaseOption::Type(Ident::new("SQLITE")),
]
},
stmt
);
}
#[test]
fn test_attach_database_simple() {
let sql = r#"ATTACH 'postgres://user.name:pass-word@some.url.com:5432/postgres'"#;
let stmt = duckdb().verified_stmt(sql);
assert_eq!(
Statement::AttachDuckDBDatabase {
if_not_exists: false,
database: false,
database_path: Ident::with_quote(
'\'',
"postgres://user.name:pass-word@some.url.com:5432/postgres"
),
database_alias: None,
attach_options: vec![]
},
stmt
);
}
#[test]
fn test_detach_database() {
let sql = r#"DETACH DATABASE IF EXISTS db_name"#;
let stmt = duckdb().verified_stmt(sql);
assert_eq!(
Statement::DetachDuckDBDatabase {
if_exists: true,
database: true,
database_alias: Ident::new("db_name"),
},
stmt
);
}
#[test]
fn test_detach_database_simple() {
let sql = r#"DETACH db_name"#;
let stmt = duckdb().verified_stmt(sql);
assert_eq!(
Statement::DetachDuckDBDatabase {
if_exists: false,
database: false,
database_alias: Ident::new("db_name"),
},
stmt
);
}
#[test]
fn test_duckdb_named_argument_function_with_assignment_operator() {
let sql = "SELECT FUN(a := '1', b := '2') FROM foo";