Skip to content

ISSUE-1147: Add support for MATERIALIZED CTEs #1148

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Feb 29, 2024
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 7 additions & 7 deletions src/ast/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -39,13 +39,13 @@ pub use self::ddl::{
};
pub use self::operator::{BinaryOperator, UnaryOperator};
pub use self::query::{
Cte, Distinct, ExceptSelectItem, ExcludeSelectItem, Fetch, ForClause, ForJson, ForXml,
GroupByExpr, IdentWithAlias, Join, JoinConstraint, JoinOperator, JsonTableColumn,
JsonTableColumnErrorHandling, LateralView, LockClause, LockType, NamedWindowDefinition,
NonBlock, Offset, OffsetRows, OrderByExpr, Query, RenameSelectItem, ReplaceSelectElement,
ReplaceSelectItem, Select, SelectInto, SelectItem, SetExpr, SetOperator, SetQuantifier, Table,
TableAlias, TableFactor, TableVersion, TableWithJoins, Top, TopQuantity, Values,
WildcardAdditionalOptions, With,
Cte, CteAsMaterialized, Distinct, ExceptSelectItem, ExcludeSelectItem, Fetch, ForClause,
ForJson, ForXml, GroupByExpr, IdentWithAlias, Join, JoinConstraint, JoinOperator,
JsonTableColumn, JsonTableColumnErrorHandling, LateralView, LockClause, LockType,
NamedWindowDefinition, NonBlock, Offset, OffsetRows, OrderByExpr, Query, RenameSelectItem,
ReplaceSelectElement, ReplaceSelectItem, Select, SelectInto, SelectItem, SetExpr, SetOperator,
SetQuantifier, Table, TableAlias, TableFactor, TableVersion, TableWithJoins, Top, TopQuantity,
Values, WildcardAdditionalOptions, With,
};
pub use self::value::{
escape_quoted_string, DateTimeField, DollarQuotedString, TrimWhereField, Value,
Expand Down
42 changes: 40 additions & 2 deletions src/ast/query.rs
Original file line number Diff line number Diff line change
Expand Up @@ -376,7 +376,36 @@ impl fmt::Display for With {
}
}

/// A single CTE (used after `WITH`): `alias [(col1, col2, ...)] AS ( query )`
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
pub enum CteAsMaterialized {
/// The WITH statement does not specify MATERIALIZED behavior
Default,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Most other places in sqlparser-rs use Option<..> to represent the case when a clause is not present.

Since I had this branch checked out locally to resolve a merge conflict, I went ahead and made that change as well in 6df536a

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah for sure!

This route means we should remove the Default enum member though because now it's not useful, and if it does get added at some point it'll print with extra whitespace.

Mind if I add a commit to the PR doing so?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sure -- go ahead. I think I did remove CteAsMaterialized::Default in 6df536a though

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah nevermind, so you did. I misread it. Thanks!

/// The WITH statement specifies AS MATERIALIZED behavior
Materialized,
/// The WITH statement specifies AS NOT MATERIALIZED behavior
NotMaterialized,
}

impl fmt::Display for CteAsMaterialized {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match *self {
CteAsMaterialized::Default => {
write!(f, "")?;
}
CteAsMaterialized::Materialized => {
write!(f, "MATERIALIZED")?;
}
CteAsMaterialized::NotMaterialized => {
write!(f, "NOT MATERIALIZED")?;
}
};
Ok(())
}
}

/// A single CTE (used after `WITH`): `<alias> [(col1, col2, ...)] AS <materialized> ( <query> )`
/// The names in the column list before `AS`, when specified, replace the names
/// of the columns returned by the query. The parser does not validate that the
/// number of columns in the query matches the number of columns in the query.
Expand All @@ -387,11 +416,20 @@ pub struct Cte {
pub alias: TableAlias,
pub query: Box<Query>,
pub from: Option<Ident>,
pub materialized: CteAsMaterialized,
}

impl fmt::Display for Cte {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "{} AS ({})", self.alias, self.query)?;
if matches!(self.materialized, CteAsMaterialized::Default) {
write!(f, "{} AS ({})", self.alias, self.query)?;
} else {
write!(
f,
"{} AS {} ({})",
self.alias, self.materialized, self.query
)?;
}
if let Some(ref fr) = self.from {
write!(f, " FROM {fr}")?;
}
Expand Down
18 changes: 18 additions & 0 deletions src/parser/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6668,6 +6668,14 @@ impl<'a> Parser<'a> {
let name = self.parse_identifier(false)?;

let mut cte = if self.parse_keyword(Keyword::AS) {
let mut is_materialized = CteAsMaterialized::Default;
if dialect_of!(self is PostgreSqlDialect) {
if self.parse_keyword(Keyword::MATERIALIZED) {
is_materialized = CteAsMaterialized::Materialized;
} else if self.parse_keywords(&[Keyword::NOT, Keyword::MATERIALIZED]) {
is_materialized = CteAsMaterialized::NotMaterialized;
}
}
self.expect_token(&Token::LParen)?;
let query = Box::new(self.parse_query()?);
self.expect_token(&Token::RParen)?;
Expand All @@ -6679,10 +6687,19 @@ impl<'a> Parser<'a> {
alias,
query,
from: None,
materialized: is_materialized,
}
} else {
let columns = self.parse_parenthesized_column_list(Optional, false)?;
self.expect_keyword(Keyword::AS)?;
let mut is_materialized = CteAsMaterialized::Default;
if dialect_of!(self is PostgreSqlDialect) {
if self.parse_keyword(Keyword::MATERIALIZED) {
is_materialized = CteAsMaterialized::Materialized;
} else if self.parse_keywords(&[Keyword::NOT, Keyword::MATERIALIZED]) {
is_materialized = CteAsMaterialized::NotMaterialized;
}
}
self.expect_token(&Token::LParen)?;
let query = Box::new(self.parse_query()?);
self.expect_token(&Token::RParen)?;
Expand All @@ -6691,6 +6708,7 @@ impl<'a> Parser<'a> {
alias,
query,
from: None,
materialized: is_materialized,
}
};
if self.parse_keyword(Keyword::FROM) {
Expand Down
1 change: 1 addition & 0 deletions tests/sqlparser_common.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5586,6 +5586,7 @@ fn parse_recursive_cte() {
},
query: Box::new(cte_query),
from: None,
materialized: CteAsMaterialized::Default,
};
assert_eq!(with.cte_tables.first().unwrap(), &expected);
}
Expand Down
9 changes: 9 additions & 0 deletions tests/sqlparser_postgres.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3834,3 +3834,12 @@ fn parse_array_agg() {
let sql4 = "SELECT ARRAY_AGG(my_schema.sections_tbl.*) AS sections FROM sections_tbl";
pg().verified_stmt(sql4);
}

#[test]
fn parse_mat_cte() {
let sql = r#"WITH cte AS MATERIALIZED (SELECT id FROM accounts) SELECT id FROM cte"#;
pg().verified_stmt(sql);

let sql2 = r#"WITH cte AS NOT MATERIALIZED (SELECT id FROM accounts) SELECT id FROM cte"#;
pg().verified_stmt(sql2);
}