Skip to content

Commit

Permalink
Add support for constraint characteristics clause (#1099)
Browse files Browse the repository at this point in the history
  • Loading branch information
dimfeld authored Jan 24, 2024
1 parent 1fb9f3e commit c86508b
Show file tree
Hide file tree
Showing 7 changed files with 480 additions and 22 deletions.
130 changes: 119 additions & 11 deletions src/ast/ddl.rs
Original file line number Diff line number Diff line change
Expand Up @@ -385,6 +385,7 @@ pub enum TableConstraint {
columns: Vec<Ident>,
/// Whether this is a `PRIMARY KEY` or just a `UNIQUE` constraint
is_primary: bool,
characteristics: Option<ConstraintCharacteristics>,
},
/// A referential integrity constraint (`[ CONSTRAINT <name> ] FOREIGN KEY (<columns>)
/// REFERENCES <foreign_table> (<referred_columns>)
Expand All @@ -398,6 +399,7 @@ pub enum TableConstraint {
referred_columns: Vec<Ident>,
on_delete: Option<ReferentialAction>,
on_update: Option<ReferentialAction>,
characteristics: Option<ConstraintCharacteristics>,
},
/// `[ CONSTRAINT <name> ] CHECK (<expr>)`
Check {
Expand Down Expand Up @@ -454,20 +456,30 @@ impl fmt::Display for TableConstraint {
name,
columns,
is_primary,
} => write!(
f,
"{}{} ({})",
display_constraint_name(name),
if *is_primary { "PRIMARY KEY" } else { "UNIQUE" },
display_comma_separated(columns)
),
characteristics,
} => {
write!(
f,
"{}{} ({})",
display_constraint_name(name),
if *is_primary { "PRIMARY KEY" } else { "UNIQUE" },
display_comma_separated(columns)
)?;

if let Some(characteristics) = characteristics {
write!(f, " {}", characteristics)?;
}

Ok(())
}
TableConstraint::ForeignKey {
name,
columns,
foreign_table,
referred_columns,
on_delete,
on_update,
characteristics,
} => {
write!(
f,
Expand All @@ -483,6 +495,9 @@ impl fmt::Display for TableConstraint {
if let Some(action) = on_update {
write!(f, " ON UPDATE {action}")?;
}
if let Some(characteristics) = characteristics {
write!(f, " {}", characteristics)?;
}
Ok(())
}
TableConstraint::Check { name, expr } => {
Expand Down Expand Up @@ -713,20 +728,24 @@ pub enum ColumnOption {
NotNull,
/// `DEFAULT <restricted-expr>`
Default(Expr),
/// `{ PRIMARY KEY | UNIQUE }`
/// `{ PRIMARY KEY | UNIQUE } [<constraint_characteristics>]`
Unique {
is_primary: bool,
characteristics: Option<ConstraintCharacteristics>,
},
/// A referential integrity constraint (`[FOREIGN KEY REFERENCES
/// <foreign_table> (<referred_columns>)
/// { [ON DELETE <referential_action>] [ON UPDATE <referential_action>] |
/// [ON UPDATE <referential_action>] [ON DELETE <referential_action>]
/// }`).
/// }
/// [<constraint_characteristics>]
/// `).
ForeignKey {
foreign_table: ObjectName,
referred_columns: Vec<Ident>,
on_delete: Option<ReferentialAction>,
on_update: Option<ReferentialAction>,
characteristics: Option<ConstraintCharacteristics>,
},
/// `CHECK (<expr>)`
Check(Expr),
Expand Down Expand Up @@ -764,14 +783,22 @@ impl fmt::Display for ColumnOption {
Null => write!(f, "NULL"),
NotNull => write!(f, "NOT NULL"),
Default(expr) => write!(f, "DEFAULT {expr}"),
Unique { is_primary } => {
write!(f, "{}", if *is_primary { "PRIMARY KEY" } else { "UNIQUE" })
Unique {
is_primary,
characteristics,
} => {
write!(f, "{}", if *is_primary { "PRIMARY KEY" } else { "UNIQUE" })?;
if let Some(characteristics) = characteristics {
write!(f, " {}", characteristics)?;
}
Ok(())
}
ForeignKey {
foreign_table,
referred_columns,
on_delete,
on_update,
characteristics,
} => {
write!(f, "REFERENCES {foreign_table}")?;
if !referred_columns.is_empty() {
Expand All @@ -783,6 +810,9 @@ impl fmt::Display for ColumnOption {
if let Some(action) = on_update {
write!(f, " ON UPDATE {action}")?;
}
if let Some(characteristics) = characteristics {
write!(f, " {}", characteristics)?;
}
Ok(())
}
Check(expr) => write!(f, "CHECK ({expr})"),
Expand Down Expand Up @@ -874,6 +904,84 @@ fn display_constraint_name(name: &'_ Option<Ident>) -> impl fmt::Display + '_ {
ConstraintName(name)
}

/// `<constraint_characteristics> = [ DEFERRABLE | NOT DEFERRABLE ] [ INITIALLY DEFERRED | INITIALLY IMMEDIATE ] [ ENFORCED | NOT ENFORCED ]`
///
/// Used in UNIQUE and foreign key constraints. The individual settings may occur in any order.
#[derive(Debug, Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
pub struct ConstraintCharacteristics {
/// `[ DEFERRABLE | NOT DEFERRABLE ]`
pub deferrable: Option<bool>,
/// `[ INITIALLY DEFERRED | INITIALLY IMMEDIATE ]`
pub initially: Option<DeferrableInitial>,
/// `[ ENFORCED | NOT ENFORCED ]`
pub enforced: Option<bool>,
}

#[derive(Debug, Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
pub enum DeferrableInitial {
/// `INITIALLY IMMEDIATE`
Immediate,
/// `INITIALLY DEFERRED`
Deferred,
}

impl ConstraintCharacteristics {
fn deferrable_text(&self) -> Option<&'static str> {
self.deferrable.map(|deferrable| {
if deferrable {
"DEFERRABLE"
} else {
"NOT DEFERRABLE"
}
})
}

fn initially_immediate_text(&self) -> Option<&'static str> {
self.initially
.map(|initially_immediate| match initially_immediate {
DeferrableInitial::Immediate => "INITIALLY IMMEDIATE",
DeferrableInitial::Deferred => "INITIALLY DEFERRED",
})
}

fn enforced_text(&self) -> Option<&'static str> {
self.enforced.map(
|enforced| {
if enforced {
"ENFORCED"
} else {
"NOT ENFORCED"
}
},
)
}
}

impl fmt::Display for ConstraintCharacteristics {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
let deferrable = self.deferrable_text();
let initially_immediate = self.initially_immediate_text();
let enforced = self.enforced_text();

match (deferrable, initially_immediate, enforced) {
(None, None, None) => Ok(()),
(None, None, Some(enforced)) => write!(f, "{enforced}"),
(None, Some(initial), None) => write!(f, "{initial}"),
(None, Some(initial), Some(enforced)) => write!(f, "{initial} {enforced}"),
(Some(deferrable), None, None) => write!(f, "{deferrable}"),
(Some(deferrable), None, Some(enforced)) => write!(f, "{deferrable} {enforced}"),
(Some(deferrable), Some(initial), None) => write!(f, "{deferrable} {initial}"),
(Some(deferrable), Some(initial), Some(enforced)) => {
write!(f, "{deferrable} {initial} {enforced}")
}
}
}
}

/// `<referential_action> =
/// { RESTRICT | CASCADE | SET NULL | NO ACTION | SET DEFAULT }`
///
Expand Down
5 changes: 3 additions & 2 deletions src/ast/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -32,8 +32,9 @@ pub use self::data_type::{
pub use self::dcl::{AlterRoleOperation, ResetConfig, RoleOption, SetConfigValue};
pub use self::ddl::{
AlterColumnOperation, AlterIndexOperation, AlterTableOperation, ColumnDef, ColumnOption,
ColumnOptionDef, GeneratedAs, GeneratedExpressionMode, IndexType, KeyOrIndexDisplay, Partition,
ProcedureParam, ReferentialAction, TableConstraint, UserDefinedTypeCompositeAttributeDef,
ColumnOptionDef, ConstraintCharacteristics, DeferrableInitial, GeneratedAs,
GeneratedExpressionMode, IndexType, KeyOrIndexDisplay, Partition, ProcedureParam,
ReferentialAction, TableConstraint, UserDefinedTypeCompositeAttributeDef,
UserDefinedTypeRepresentation, ViewColumnDef,
};
pub use self::operator::{BinaryOperator, UnaryOperator};
Expand Down
3 changes: 3 additions & 0 deletions src/keywords.rs
Original file line number Diff line number Diff line change
Expand Up @@ -210,6 +210,7 @@ define_keywords!(
DECIMAL,
DECLARE,
DEFAULT,
DEFERRABLE,
DEFERRED,
DELAYED,
DELETE,
Expand Down Expand Up @@ -250,6 +251,7 @@ define_keywords!(
ENDPOINT,
END_FRAME,
END_PARTITION,
ENFORCED,
ENGINE,
ENUM,
EPOCH,
Expand Down Expand Up @@ -343,6 +345,7 @@ define_keywords!(
INDEX,
INDICATOR,
INHERIT,
INITIALLY,
INNER,
INOUT,
INPUTFORMAT,
Expand Down
62 changes: 60 additions & 2 deletions src/parser/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4478,9 +4478,17 @@ impl<'a> Parser<'a> {
} else if self.parse_keyword(Keyword::DEFAULT) {
Ok(Some(ColumnOption::Default(self.parse_expr()?)))
} else if self.parse_keywords(&[Keyword::PRIMARY, Keyword::KEY]) {
Ok(Some(ColumnOption::Unique { is_primary: true }))
let characteristics = self.parse_constraint_characteristics()?;
Ok(Some(ColumnOption::Unique {
is_primary: true,
characteristics,
}))
} else if self.parse_keyword(Keyword::UNIQUE) {
Ok(Some(ColumnOption::Unique { is_primary: false }))
let characteristics = self.parse_constraint_characteristics()?;
Ok(Some(ColumnOption::Unique {
is_primary: false,
characteristics,
}))
} else if self.parse_keyword(Keyword::REFERENCES) {
let foreign_table = self.parse_object_name(false)?;
// PostgreSQL allows omitting the column list and
Expand All @@ -4499,11 +4507,14 @@ impl<'a> Parser<'a> {
break;
}
}
let characteristics = self.parse_constraint_characteristics()?;

Ok(Some(ColumnOption::ForeignKey {
foreign_table,
referred_columns,
on_delete,
on_update,
characteristics,
}))
} else if self.parse_keyword(Keyword::CHECK) {
self.expect_token(&Token::LParen)?;
Expand Down Expand Up @@ -4658,6 +4669,47 @@ impl<'a> Parser<'a> {
}
}

pub fn parse_constraint_characteristics(
&mut self,
) -> Result<Option<ConstraintCharacteristics>, ParserError> {
let mut cc = ConstraintCharacteristics {
deferrable: None,
initially: None,
enforced: None,
};

loop {
if cc.deferrable.is_none() && self.parse_keywords(&[Keyword::NOT, Keyword::DEFERRABLE])
{
cc.deferrable = Some(false);
} else if cc.deferrable.is_none() && self.parse_keyword(Keyword::DEFERRABLE) {
cc.deferrable = Some(true);
} else if cc.initially.is_none() && self.parse_keyword(Keyword::INITIALLY) {
if self.parse_keyword(Keyword::DEFERRED) {
cc.initially = Some(DeferrableInitial::Deferred);
} else if self.parse_keyword(Keyword::IMMEDIATE) {
cc.initially = Some(DeferrableInitial::Immediate);
} else {
self.expected("one of DEFERRED or IMMEDIATE", self.peek_token())?;
}
} else if cc.enforced.is_none() && self.parse_keyword(Keyword::ENFORCED) {
cc.enforced = Some(true);
} else if cc.enforced.is_none()
&& self.parse_keywords(&[Keyword::NOT, Keyword::ENFORCED])
{
cc.enforced = Some(false);
} else {
break;
}
}

if cc.deferrable.is_some() || cc.initially.is_some() || cc.enforced.is_some() {
Ok(Some(cc))
} else {
Ok(None)
}
}

pub fn parse_optional_table_constraint(
&mut self,
) -> Result<Option<TableConstraint>, ParserError> {
Expand All @@ -4681,10 +4733,12 @@ impl<'a> Parser<'a> {
.or(name);

let columns = self.parse_parenthesized_column_list(Mandatory, false)?;
let characteristics = self.parse_constraint_characteristics()?;
Ok(Some(TableConstraint::Unique {
name,
columns,
is_primary,
characteristics,
}))
}
Token::Word(w) if w.keyword == Keyword::FOREIGN => {
Expand All @@ -4706,13 +4760,17 @@ impl<'a> Parser<'a> {
break;
}
}

let characteristics = self.parse_constraint_characteristics()?;

Ok(Some(TableConstraint::ForeignKey {
name,
columns,
foreign_table,
referred_columns,
on_delete,
on_update,
characteristics,
}))
}
Token::Word(w) if w.keyword == Keyword::CHECK => {
Expand Down
Loading

0 comments on commit c86508b

Please sign in to comment.