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

Support parsing constraints in CREATE TABLE

Browse source 
<table element> ::= ... | <table constraint definition> | ...
https://jakewheat.github.io/sql-overview/sql-2011-foundation-grammar.html#table-element-list
This commit is contained in:
Nickolay Ponomarev 2019-05-11 23:51:30 +03:00
parent c69a1881c7
commit aab0c36443
4 changed files with 123 additions and 65 deletions
Download patch file Download diff file Expand all files Collapse all files

34
src/sqlast/mod.rs
View file

@ -396,6 +396,7 @@ pub enum SQLStatement {
name: SQLObjectName, name: SQLObjectName,
/// Optional schema /// Optional schema
columns: Vec<SQLColumnDef>, columns: Vec<SQLColumnDef>,
constraints: Vec<TableConstraint>,
external: bool, external: bool,
file_format: Option<FileFormat>, file_format: Option<FileFormat>,
location: Option<String>, location: Option<String>,
@ -503,21 +504,30 @@ impl ToString for SQLStatement {
SQLStatement::SQLCreateTable { SQLStatement::SQLCreateTable {
name, name,
columns, columns,
constraints,
external, external,
file_format, file_format,
location, location,
} if *external => format!( } => {
"CREATE EXTERNAL TABLE {} ({}) STORED AS {} LOCATION '{}'", let mut s = format!(
name.to_string(), "CREATE {}TABLE {} ({}",
comma_separated_string(columns), if *external { "EXTERNAL " } else { "" },
file_format.as_ref().unwrap().to_string(), name.to_string(),
location.as_ref().unwrap() comma_separated_string(columns)
), );
SQLStatement::SQLCreateTable { name, columns, .. } => format!( if !constraints.is_empty() {
"CREATE TABLE {} ({})", s += &format!(", {}", comma_separated_string(constraints));
name.to_string(), }
comma_separated_string(columns) s += ")";
), if *external {
s += &format!(
" STORED AS {} LOCATION '{}'",
file_format.as_ref().unwrap().to_string(),
location.as_ref().unwrap()
);
}
s
}
SQLStatement::SQLAlterTable { name, operation } => { SQLStatement::SQLAlterTable { name, operation } => {
format!("ALTER TABLE {} {}", name.to_string(), operation.to_string()) format!("ALTER TABLE {} {}", name.to_string(), operation.to_string())
} }

124
src/sqlparser.rs
View file

@ -769,7 +769,7 @@ impl Parser {
pub fn parse_create_external_table(&mut self) -> Result<SQLStatement, ParserError> { pub fn parse_create_external_table(&mut self) -> Result<SQLStatement, ParserError> {
self.expect_keyword("TABLE")?; self.expect_keyword("TABLE")?;
let table_name = self.parse_object_name()?; let table_name = self.parse_object_name()?;
let columns = self.parse_columns()?; let (columns, constraints) = self.parse_columns()?;
self.expect_keyword("STORED")?; self.expect_keyword("STORED")?;
self.expect_keyword("AS")?; self.expect_keyword("AS")?;
let file_format = self.parse_identifier()?.parse::<FileFormat>()?; let file_format = self.parse_identifier()?.parse::<FileFormat>()?;
@ -780,6 +780,7 @@ impl Parser {
Ok(SQLStatement::SQLCreateTable { Ok(SQLStatement::SQLCreateTable {
name: table_name, name: table_name,
columns, columns,
constraints,
external: true, external: true,
file_format: Some(file_format), file_format: Some(file_format),
location: Some(location), location: Some(location),
@ -845,74 +846,78 @@ impl Parser {
pub fn parse_create_table(&mut self) -> Result<SQLStatement, ParserError> { pub fn parse_create_table(&mut self) -> Result<SQLStatement, ParserError> {
let table_name = self.parse_object_name()?; let table_name = self.parse_object_name()?;
// parse optional column list (schema) // parse optional column list (schema)
let columns = self.parse_columns()?; let (columns, constraints) = self.parse_columns()?;
Ok(SQLStatement::SQLCreateTable { Ok(SQLStatement::SQLCreateTable {
name: table_name, name: table_name,
columns, columns,
constraints,
external: false, external: false,
file_format: None, file_format: None,
location: None, location: None,
}) })
} }
fn parse_columns(&mut self) -> Result<Vec<SQLColumnDef>, ParserError> { fn parse_columns(&mut self) -> Result<(Vec<SQLColumnDef>, Vec<TableConstraint>), ParserError> {
let mut columns = vec![]; let mut columns = vec![];
let mut constraints = vec![];
if !self.consume_token(&Token::LParen) { if !self.consume_token(&Token::LParen) {
return Ok(columns); return Ok((columns, constraints));
} }
loop { loop {
match self.next_token() { if let Some(constraint) = self.parse_optional_table_constraint()? {
Some(Token::SQLWord(column_name)) => { constraints.push(constraint);
let data_type = self.parse_data_type()?; } else if let Some(Token::SQLWord(column_name)) = self.peek_token() {
let is_primary = self.parse_keywords(vec!["PRIMARY", "KEY"]); self.next_token();
let is_unique = self.parse_keyword("UNIQUE"); let data_type = self.parse_data_type()?;
let default = if self.parse_keyword("DEFAULT") { let is_primary = self.parse_keywords(vec!["PRIMARY", "KEY"]);
let expr = self.parse_default_expr(0)?; let is_unique = self.parse_keyword("UNIQUE");
Some(expr) let default = if self.parse_keyword("DEFAULT") {
} else { let expr = self.parse_default_expr(0)?;
None Some(expr)
}; } else {
let allow_null = if self.parse_keywords(vec!["NOT", "NULL"]) { None
false };
} else { let allow_null = if self.parse_keywords(vec!["NOT", "NULL"]) {
let _ = self.parse_keyword("NULL"); false
true } else {
}; let _ = self.parse_keyword("NULL");
debug!("default: {:?}", default); true
};
debug!("default: {:?}", default);
columns.push(SQLColumnDef { columns.push(SQLColumnDef {
name: column_name.as_sql_ident(), name: column_name.as_sql_ident(),
data_type, data_type,
allow_null, allow_null,
is_primary, is_primary,
is_unique, is_unique,
default, default,
}); });
match self.next_token() { } else {
Some(Token::Comma) => {} return self.expected("column name or constraint definition", self.peek_token());
Some(Token::RParen) => { }
break; match self.next_token() {
} Some(Token::Comma) => {}
other => { Some(Token::RParen) => {
return parser_err!(format!( break;
"Expected ',' or ')' after column definition but found {:?}",
other
));
}
}
} }
unexpected => { other => {
return parser_err!(format!("Expected column name, got {:?}", unexpected)); return parser_err!(format!(
"Expected ',' or ')' after column definition but found {:?}",
other
));
} }
} }
} }
Ok(columns) Ok((columns, constraints))
} }
pub fn parse_table_constraint(&mut self) -> Result<TableConstraint, ParserError> { pub fn parse_optional_table_constraint(
&mut self,
) -> Result<Option<TableConstraint>, ParserError> {
let name = if self.parse_keyword("CONSTRAINT") { let name = if self.parse_keyword("CONSTRAINT") {
Some(self.parse_identifier()?) Some(self.parse_identifier()?)
} else { } else {
@ -925,11 +930,11 @@ impl Parser {
self.expect_keyword("KEY")?; self.expect_keyword("KEY")?;
} }
let columns = self.parse_parenthesized_column_list(Mandatory)?; let columns = self.parse_parenthesized_column_list(Mandatory)?;
Ok(TableConstraint::Unique { Ok(Some(TableConstraint::Unique {
name, name,
columns, columns,
is_primary, is_primary,
}) }))
} }
Some(Token::SQLWord(ref k)) if k.keyword == "FOREIGN" => { Some(Token::SQLWord(ref k)) if k.keyword == "FOREIGN" => {
self.expect_keyword("KEY")?; self.expect_keyword("KEY")?;
@ -937,20 +942,29 @@ impl Parser {
self.expect_keyword("REFERENCES")?; self.expect_keyword("REFERENCES")?;
let foreign_table = self.parse_object_name()?; let foreign_table = self.parse_object_name()?;
let referred_columns = self.parse_parenthesized_column_list(Mandatory)?; let referred_columns = self.parse_parenthesized_column_list(Mandatory)?;
Ok(TableConstraint::ForeignKey { Ok(Some(TableConstraint::ForeignKey {
name, name,
columns, columns,
foreign_table, foreign_table,
referred_columns, referred_columns,
}) }))
} }
Some(Token::SQLWord(ref k)) if k.keyword == "CHECK" => { Some(Token::SQLWord(ref k)) if k.keyword == "CHECK" => {
self.expect_token(&Token::LParen)?; self.expect_token(&Token::LParen)?;
let expr = Box::new(self.parse_expr()?); let expr = Box::new(self.parse_expr()?);
self.expect_token(&Token::RParen)?; self.expect_token(&Token::RParen)?;
Ok(TableConstraint::Check { name, expr }) Ok(Some(TableConstraint::Check { name, expr }))
}
unexpected => {
if name.is_some() {
self.expected("PRIMARY, UNIQUE, FOREIGN, or CHECK", unexpected)
} else {
if unexpected.is_some() {
self.prev_token();
}
Ok(None)
}
} }
_ => self.expected("PRIMARY, UNIQUE, or FOREIGN", self.peek_token()),
} }
} }
@ -959,7 +973,11 @@ impl Parser {
let _ = self.parse_keyword("ONLY"); let _ = self.parse_keyword("ONLY");
let table_name = self.parse_object_name()?; let table_name = self.parse_object_name()?;
let operation = if self.parse_keyword("ADD") { let operation = if self.parse_keyword("ADD") {
AlterOperation::AddConstraint(self.parse_table_constraint()?) if let Some(constraint) = self.parse_optional_table_constraint()? {
AlterOperation::AddConstraint(constraint)
} else {
return self.expected("a constraint in ALTER TABLE .. ADD", self.peek_token());
}
} else { } else {
return self.expected("ADD after ALTER TABLE", self.peek_token()); return self.expected("ADD after ALTER TABLE", self.peek_token());
}; };

24
tests/sqlparser_common.rs
View file

@ -660,12 +660,14 @@ fn parse_create_table() {
SQLStatement::SQLCreateTable { SQLStatement::SQLCreateTable {
name, name,
columns, columns,
constraints,
external: false, external: false,
file_format: None, file_format: None,
location: None, location: None,
} => { } => {
assert_eq!("uk_cities", name.to_string()); assert_eq!("uk_cities", name.to_string());
assert_eq!(3, columns.len()); assert_eq!(3, columns.len());
assert!(constraints.is_empty());
let c_name = &columns[0]; let c_name = &columns[0];
assert_eq!("name", c_name.name); assert_eq!("name", c_name.name);
@ -705,12 +707,14 @@ fn parse_create_external_table() {
SQLStatement::SQLCreateTable { SQLStatement::SQLCreateTable {
name, name,
columns, columns,
constraints,
external, external,
file_format, file_format,
location, location,
} => { } => {
assert_eq!("uk_cities", name.to_string()); assert_eq!("uk_cities", name.to_string());
assert_eq!(3, columns.len()); assert_eq!(3, columns.len());
assert!(constraints.is_empty());
let c_name = &columns[0]; let c_name = &columns[0];
assert_eq!("name", c_name.name); assert_eq!("name", c_name.name);
@ -761,9 +765,29 @@ fn parse_alter_table_constraints() {
} }
_ => unreachable!(), _ => unreachable!(),
} }
verified_stmt(&format!("CREATE TABLE foo (id int, {})", constraint_text));
} }
} }
#[test]
fn parse_bad_constraint() {
let res = parse_sql_statements("ALTER TABLE tab ADD");
assert_eq!(
ParserError::ParserError(
"Expected a constraint in ALTER TABLE .. ADD, found: EOF".to_string()
),
res.unwrap_err()
);
let res = parse_sql_statements("CREATE TABLE tab (foo int,");
assert_eq!(
ParserError::ParserError(
"Expected column name or constraint definition, found: EOF".to_string()
),
res.unwrap_err()
);
}
#[test] #[test]
fn parse_scalar_function_in_projection() { fn parse_scalar_function_in_projection() {
let sql = "SELECT sqrt(id) FROM foo"; let sql = "SELECT sqrt(id) FROM foo";

6
tests/sqlparser_postgres.rs
View file

@ -23,12 +23,14 @@ fn parse_create_table_with_defaults() {
SQLStatement::SQLCreateTable { SQLStatement::SQLCreateTable {
name, name,
columns, columns,
constraints,
external: false, external: false,
file_format: None, file_format: None,
location: None, location: None,
} => { } => {
assert_eq!("public.customer", name.to_string()); assert_eq!("public.customer", name.to_string());
assert_eq!(10, columns.len()); assert_eq!(10, columns.len());
assert!(constraints.is_empty());
let c_name = &columns[0]; let c_name = &columns[0];
assert_eq!("customer_id", c_name.name); assert_eq!("customer_id", c_name.name);
@ -69,11 +71,13 @@ fn parse_create_table_from_pg_dump() {
SQLStatement::SQLCreateTable { SQLStatement::SQLCreateTable {
name, name,
columns, columns,
constraints,
external: false, external: false,
file_format: None, file_format: None,
location: None, location: None,
} => { } => {
assert_eq!("public.customer", name.to_string()); assert_eq!("public.customer", name.to_string());
assert!(constraints.is_empty());
let c_customer_id = &columns[0]; let c_customer_id = &columns[0];
assert_eq!("customer_id", c_customer_id.name); assert_eq!("customer_id", c_customer_id.name);
@ -130,11 +134,13 @@ fn parse_create_table_with_inherit() {
SQLStatement::SQLCreateTable { SQLStatement::SQLCreateTable {
name, name,
columns, columns,
constraints,
external: false, external: false,
file_format: None, file_format: None,
location: None, location: None,
} => { } => {
assert_eq!("bazaar.settings", name.to_string()); assert_eq!("bazaar.settings", name.to_string());
assert!(constraints.is_empty());
let c_name = &columns[0]; let c_name = &columns[0];
assert_eq!("settings_id", c_name.name); assert_eq!("settings_id", c_name.name);