Skip to content

Commit 0b148fc

Browse files
nickolaynickolay
committed
Support parsing constraints in CREATE TABLE
<table element> ::= ... | <table constraint definition> | ... https://jakewheat.github.io/sql-overview/sql-2011-foundation-grammar.html#table-element-list
1 parent 1150e17 commit 0b148fc

File tree

4 files changed

+123
-65
lines changed

4 files changed

+123
-65
lines changed

‎src/sqlast/mod.rs

Lines changed: 22 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -390,6 +390,7 @@ pub enum SQLStatement {
390390
name: SQLObjectName,
391391
/// Optional schema
392392
columns: Vec<SQLColumnDef>,
393+
constraints: Vec<TableConstraint>,
393394
external: bool,
394395
file_format: Option<FileFormat>,
395396
location: Option<String>,
@@ -493,21 +494,30 @@ impl ToString for SQLStatement {
493494
SQLStatement::SQLCreateTable {
494495
name,
495496
columns,
497+
constraints,
496498
external,
497499
file_format,
498500
location,
499-
} if *external => format!(
500-
"CREATE EXTERNAL TABLE {} ({}) STORED AS {} LOCATION '{}'",
501-
name.to_string(),
502-
comma_separated_string(columns),
503-
file_format.as_ref().unwrap().to_string(),
504-
location.as_ref().unwrap()
505-
),
506-
SQLStatement::SQLCreateTable { name, columns, .. } => format!(
507-
"CREATE TABLE {} ({})",
508-
name.to_string(),
509-
comma_separated_string(columns)
510-
),
501+
} => {
502+
let mut s = format!(
503+
"CREATE {}TABLE {} ({}",
504+
if *external { "EXTERNAL " } else { "" },
505+
name.to_string(),
506+
comma_separated_string(columns)
507+
);
508+
if !constraints.is_empty() {
509+
s += &format!(", {}", comma_separated_string(constraints));
510+
}
511+
s += ")";
512+
if *external {
513+
s += &format!(
514+
" STORED AS {} LOCATION '{}'",
515+
file_format.as_ref().unwrap().to_string(),
516+
location.as_ref().unwrap()
517+
);
518+
}
519+
s
520+
}
511521
SQLStatement::SQLAlterTable { name, operation } => {
512522
format!("ALTER TABLE {} {}", name.to_string(), operation.to_string())
513523
}

‎src/sqlparser.rs

Lines changed: 71 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -703,7 +703,7 @@ impl Parser {
703703
pub fn parse_create_external_table(&mut self) -> Result<SQLStatement, ParserError> {
704704
self.expect_keyword("TABLE")?;
705705
let table_name = self.parse_object_name()?;
706-
let columns = self.parse_columns()?;
706+
let (columns, constraints) = self.parse_columns()?;
707707
self.expect_keyword("STORED")?;
708708
self.expect_keyword("AS")?;
709709
let file_format = self.parse_identifier()?.parse::<FileFormat>()?;
@@ -714,6 +714,7 @@ impl Parser {
714714
Ok(SQLStatement::SQLCreateTable {
715715
name: table_name,
716716
columns,
717+
constraints,
717718
external: true,
718719
file_format: Some(file_format),
719720
location: Some(location),
@@ -742,74 +743,78 @@ impl Parser {
742743
pub fn parse_create_table(&mut self) -> Result<SQLStatement, ParserError> {
743744
let table_name = self.parse_object_name()?;
744745
// parse optional column list (schema)
745-
let columns = self.parse_columns()?;
746+
let (columns, constraints) = self.parse_columns()?;
746747

747748
Ok(SQLStatement::SQLCreateTable {
748749
name: table_name,
749750
columns,
751+
constraints,
750752
external: false,
751753
file_format: None,
752754
location: None,
753755
})
754756
}
755757

756-
fn parse_columns(&mut self) -> Result<Vec<SQLColumnDef>, ParserError> {
758+
fn parse_columns(&mut self) -> Result<(Vec<SQLColumnDef>, Vec<TableConstraint>), ParserError> {
757759
let mut columns = vec![];
760+
let mut constraints = vec![];
758761
if !self.consume_token(&Token::LParen) {
759-
return Ok(columns);
762+
return Ok((columns, constraints));
760763
}
761764

762765
loop {
766+
if let Some(constraint) = self.parse_optional_table_constraint()? {
767+
constraints.push(constraint);
768+
} else if let Some(Token::SQLWord(column_name)) = self.peek_token() {
769+
self.next_token();
770+
let data_type = self.parse_data_type()?;
771+
let is_primary = self.parse_keywords(vec!["PRIMARY", "KEY"]);
772+
let is_unique = self.parse_keyword("UNIQUE");
773+
let default = if self.parse_keyword("DEFAULT") {
774+
let expr = self.parse_default_expr(0)?;
775+
Some(expr)
776+
} else {
777+
None
778+
};
779+
let allow_null = if self.parse_keywords(vec!["NOT", "NULL"]) {
780+
false
781+
} else {
782+
let _ = self.parse_keyword("NULL");
783+
true
784+
};
785+
debug!("default: {:?}", default);
786+
787+
columns.push(SQLColumnDef {
788+
name: column_name.as_sql_ident(),
789+
data_type,
790+
allow_null,
791+
is_primary,
792+
is_unique,
793+
default,
794+
});
795+
} else {
796+
return self.expected("column name or constraint definition", self.peek_token());
797+
}
763798
match self.next_token() {
764-
Some(Token::SQLWord(column_name)) => {
765-
let data_type = self.parse_data_type()?;
766-
let is_primary = self.parse_keywords(vec!["PRIMARY", "KEY"]);
767-
let is_unique = self.parse_keyword("UNIQUE");
768-
let default = if self.parse_keyword("DEFAULT") {
769-
let expr = self.parse_default_expr(0)?;
770-
Some(expr)
771-
} else {
772-
None
773-
};
774-
let allow_null = if self.parse_keywords(vec!["NOT", "NULL"]) {
775-
false
776-
} else {
777-
let _ = self.parse_keyword("NULL");
778-
true
779-
};
780-
debug!("default: {:?}", default);
781-
782-
columns.push(SQLColumnDef {
783-
name: column_name.as_sql_ident(),
784-
data_type,
785-
allow_null,
786-
is_primary,
787-
is_unique,
788-
default,
789-
});
790-
match self.next_token() {
791-
Some(Token::Comma) => {}
792-
Some(Token::RParen) => {
793-
break;
794-
}
795-
other => {
796-
return parser_err!(format!(
797-
"Expected ',' or ')' after column definition but found {:?}",
798-
other
799-
));
800-
}
801-
}
799+
Some(Token::Comma) => {}
800+
Some(Token::RParen) => {
801+
break;
802802
}
803-
unexpected => {
804-
return parser_err!(format!("Expected column name, got {:?}", unexpected));
803+
other => {
804+
return parser_err!(format!(
805+
"Expected ',' or ')' after column definition but found {:?}",
806+
other
807+
));
805808
}
806809
}
807810
}
808811

809-
Ok(columns)
812+
Ok((columns, constraints))
810813
}
811814

812-
pub fn parse_table_constraint(&mut self) -> Result<TableConstraint, ParserError> {
815+
pub fn parse_optional_table_constraint(
816+
&mut self,
817+
) -> Result<Option<TableConstraint>, ParserError> {
813818
let name = if self.parse_keyword("CONSTRAINT") {
814819
Some(self.parse_identifier()?)
815820
} else {
@@ -822,32 +827,41 @@ impl Parser {
822827
self.expect_keyword("KEY")?;
823828
}
824829
let columns = self.parse_parenthesized_column_list(Mandatory)?;
825-
Ok(TableConstraint::Unique {
830+
Ok(Some(TableConstraint::Unique {
826831
name,
827832
columns,
828833
is_primary,
829-
})
834+
}))
830835
}
831836
Some(Token::SQLWord(ref k)) if k.keyword == "FOREIGN" => {
832837
self.expect_keyword("KEY")?;
833838
let columns = self.parse_parenthesized_column_list(Mandatory)?;
834839
self.expect_keyword("REFERENCES")?;
835840
let foreign_table = self.parse_object_name()?;
836841
let referred_columns = self.parse_parenthesized_column_list(Mandatory)?;
837-
Ok(TableConstraint::ForeignKey {
842+
Ok(Some(TableConstraint::ForeignKey {
838843
name,
839844
columns,
840845
foreign_table,
841846
referred_columns,
842-
})
847+
}))
843848
}
844849
Some(Token::SQLWord(ref k)) if k.keyword == "CHECK" => {
845850
self.expect_token(&Token::LParen)?;
846851
let expr = Box::new(self.parse_expr()?);
847852
self.expect_token(&Token::RParen)?;
848-
Ok(TableConstraint::Check { name, expr })
853+
Ok(Some(TableConstraint::Check { name, expr }))
854+
}
855+
unexpected => {
856+
if name.is_some() {
857+
self.expected("PRIMARY, UNIQUE, FOREIGN, or CHECK", unexpected)
858+
} else {
859+
if unexpected.is_some() {
860+
self.prev_token();
861+
}
862+
Ok(None)
863+
}
849864
}
850-
_ => self.expected("PRIMARY, UNIQUE, or FOREIGN", self.peek_token()),
851865
}
852866
}
853867

@@ -868,7 +882,11 @@ impl Parser {
868882
let _ = self.parse_keyword("ONLY");
869883
let table_name = self.parse_object_name()?;
870884
let operation = if self.parse_keyword("ADD") {
871-
AlterOperation::AddConstraint(self.parse_table_constraint()?)
885+
if let Some(constraint) = self.parse_optional_table_constraint()? {
886+
AlterOperation::AddConstraint(constraint)
887+
} else {
888+
return self.expected("a constraint in ALTER TABLE .. ADD", self.peek_token());
889+
}
872890
} else {
873891
return self.expected("ADD after ALTER TABLE", self.peek_token());
874892
};

‎tests/sqlparser_common.rs

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -545,12 +545,14 @@ fn parse_create_table() {
545545
SQLStatement::SQLCreateTable {
546546
name,
547547
columns,
548+
constraints,
548549
external: false,
549550
file_format: None,
550551
location: None,
551552
} => {
552553
assert_eq!("uk_cities", name.to_string());
553554
assert_eq!(3, columns.len());
555+
assert!(constraints.is_empty());
554556

555557
let c_name = &columns[0];
556558
assert_eq!("name", c_name.name);
@@ -590,12 +592,14 @@ fn parse_create_external_table() {
590592
SQLStatement::SQLCreateTable {
591593
name,
592594
columns,
595+
constraints,
593596
external,
594597
file_format,
595598
location,
596599
} => {
597600
assert_eq!("uk_cities", name.to_string());
598601
assert_eq!(3, columns.len());
602+
assert!(constraints.is_empty());
599603

600604
let c_name = &columns[0];
601605
assert_eq!("name", c_name.name);
@@ -662,9 +666,29 @@ fn parse_alter_table_constraints() {
662666
}
663667
_ => unreachable!(),
664668
}
669+
verified_stmt(&format!("CREATE TABLE foo (id int, {})", constraint_text));
665670
}
666671
}
667672

673+
#[test]
674+
fn parse_bad_constraint() {
675+
let res = parse_sql_statements("ALTER TABLE tab ADD");
676+
assert_eq!(
677+
ParserError::ParserError(
678+
"Expected a constraint in ALTER TABLE .. ADD, found: EOF".to_string()
679+
),
680+
res.unwrap_err()
681+
);
682+
683+
let res = parse_sql_statements("CREATE TABLE tab (foo int,");
684+
assert_eq!(
685+
ParserError::ParserError(
686+
"Expected column name or constraint definition, found: EOF".to_string()
687+
),
688+
res.unwrap_err()
689+
);
690+
}
691+
668692
#[test]
669693
fn parse_scalar_function_in_projection() {
670694
let sql = "SELECT sqrt(id) FROM foo";

‎tests/sqlparser_postgres.rs

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,12 +23,14 @@ fn parse_create_table_with_defaults() {
2323
SQLStatement::SQLCreateTable {
2424
name,
2525
columns,
26+
constraints,
2627
external: false,
2728
file_format: None,
2829
location: None,
2930
} => {
3031
assert_eq!("public.customer", name.to_string());
3132
assert_eq!(10, columns.len());
33+
assert!(constraints.is_empty());
3234

3335
let c_name = &columns[0];
3436
assert_eq!("customer_id", c_name.name);
@@ -69,11 +71,13 @@ fn parse_create_table_from_pg_dump() {
6971
SQLStatement::SQLCreateTable {
7072
name,
7173
columns,
74+
constraints,
7275
external: false,
7376
file_format: None,
7477
location: None,
7578
} => {
7679
assert_eq!("public.customer", name.to_string());
80+
assert!(constraints.is_empty());
7781

7882
let c_customer_id = &columns[0];
7983
assert_eq!("customer_id", c_customer_id.name);
@@ -130,11 +134,13 @@ fn parse_create_table_with_inherit() {
130134
SQLStatement::SQLCreateTable {
131135
name,
132136
columns,
137+
constraints,
133138
external: false,
134139
file_format: None,
135140
location: None,
136141
} => {
137142
assert_eq!("bazaar.settings", name.to_string());
143+
assert!(constraints.is_empty());
138144

139145
let c_name = &columns[0];
140146
assert_eq!("settings_id", c_name.name);

0 commit comments

Comments
 (0)