ISSUE-1147: Add support for MATERIALIZED CTEs

In Postgres, the [NOT] MATERIALIZED keywords can be applied when
creating CTEs. Add support for them.

Author: ReppCodes

src/ast/mod.rs

@@ -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,

src/ast/query.rs

@@ -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,
+    /// 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.
@@ -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}")?;
         }

src/parser/mod.rs

@@ -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)?;
@@ -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)?;
@@ -6691,6 +6708,7 @@ impl<'a> Parser<'a> {
                 alias,
                 query,
                 from: None,
+                materialized: is_materialized,
             }
         };
         if self.parse_keyword(Keyword::FROM) {

tests/sqlparser_common.rs

@@ -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);
 }

tests/sqlparser_postgres.rs

@@ -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);
+}