Fix STORAGE LIFECYCLE POLICY for snowflake queries (#2264)

Author: romanoff

src/ast/ddl.rs

@@ -48,9 +48,9 @@ use crate::ast::{
     HiveFormat, HiveIOFormat, HiveRowFormat, HiveSetLocation, Ident, InitializeKind,
     MySQLColumnPosition, ObjectName, OnCommit, OneOrManyWithParens, OperateFunctionArg,
     OrderByExpr, ProjectionSelect, Query, RefreshModeKind, RowAccessPolicy, SequenceOptions,
-    Spanned, SqlOption, StorageSerializationPolicy, TableVersion, Tag, TriggerEvent,
-    TriggerExecBody, TriggerObject, TriggerPeriod, TriggerReferencing, Value, ValueWithSpan,
-    WrappedCollection,
+    Spanned, SqlOption, StorageLifecyclePolicy, StorageSerializationPolicy, TableVersion, Tag,
+    TriggerEvent, TriggerExecBody, TriggerObject, TriggerPeriod, TriggerReferencing, Value,
+    ValueWithSpan, WrappedCollection,
 };
 use crate::display_utils::{DisplayCommaSeparated, Indent, NewLine, SpaceOrNewline};
 use crate::keywords::Keyword;
@@ -3012,6 +3012,9 @@ pub struct CreateTable {
     /// Snowflake "WITH ROW ACCESS POLICY" clause
     /// <https://docs.snowflake.com/en/sql-reference/sql/create-table>
     pub with_row_access_policy: Option<RowAccessPolicy>,
+    /// Snowflake `WITH STORAGE LIFECYCLE POLICY` clause
+    /// <https://docs.snowflake.com/en/sql-reference/sql/create-table>
+    pub with_storage_lifecycle_policy: Option<StorageLifecyclePolicy>,
     /// Snowflake "WITH TAG" clause
     /// <https://docs.snowflake.com/en/sql-reference/sql/create-table>
     pub with_tags: Option<Vec<Tag>>,
@@ -3317,6 +3320,10 @@ impl fmt::Display for CreateTable {
             write!(f, " {row_access_policy}",)?;
         }
 
+        if let Some(storage_lifecycle_policy) = &self.with_storage_lifecycle_policy {
+            write!(f, " {storage_lifecycle_policy}",)?;
+        }
+
         if let Some(tag) = &self.with_tags {
             write!(f, " WITH TAG ({})", display_comma_separated(tag.as_slice()))?;
         }

src/ast/helpers/stmt_create_table.rs

@@ -28,8 +28,8 @@ use crate::ast::{
     ClusteredBy, ColumnDef, CommentDef, CreateTable, CreateTableLikeKind, CreateTableOptions,
     DistStyle, Expr, FileFormat, ForValues, HiveDistributionStyle, HiveFormat, Ident,
     InitializeKind, ObjectName, OnCommit, OneOrManyWithParens, Query, RefreshModeKind,
-    RowAccessPolicy, Statement, StorageSerializationPolicy, TableConstraint, TableVersion, Tag,
-    WrappedCollection,
+    RowAccessPolicy, Statement, StorageLifecyclePolicy, StorageSerializationPolicy,
+    TableConstraint, TableVersion, Tag, WrappedCollection,
 };
 
 use crate::parser::ParserError;
@@ -149,6 +149,8 @@ pub struct CreateTableBuilder {
     pub with_aggregation_policy: Option<ObjectName>,
     /// Optional row access policy applied to the table.
     pub with_row_access_policy: Option<RowAccessPolicy>,
+    /// Optional storage lifecycle policy applied to the table.
+    pub with_storage_lifecycle_policy: Option<StorageLifecyclePolicy>,
     /// Optional tags/labels attached to the table metadata.
     pub with_tags: Option<Vec<Tag>>,
     /// Optional base location for staged data.
@@ -227,6 +229,7 @@ impl CreateTableBuilder {
             default_ddl_collation: None,
             with_aggregation_policy: None,
             with_row_access_policy: None,
+            with_storage_lifecycle_policy: None,
             with_tags: None,
             base_location: None,
             external_volume: None,
@@ -459,6 +462,14 @@ impl CreateTableBuilder {
         self.with_row_access_policy = with_row_access_policy;
         self
     }
+    /// Attach a storage lifecycle policy to the table.
+    pub fn with_storage_lifecycle_policy(
+        mut self,
+        with_storage_lifecycle_policy: Option<StorageLifecyclePolicy>,
+    ) -> Self {
+        self.with_storage_lifecycle_policy = with_storage_lifecycle_policy;
+        self
+    }
     /// Attach tags/labels to the table metadata.
     pub fn with_tags(mut self, with_tags: Option<Vec<Tag>>) -> Self {
         self.with_tags = with_tags;
@@ -582,6 +593,7 @@ impl CreateTableBuilder {
             default_ddl_collation: self.default_ddl_collation,
             with_aggregation_policy: self.with_aggregation_policy,
             with_row_access_policy: self.with_row_access_policy,
+            with_storage_lifecycle_policy: self.with_storage_lifecycle_policy,
             with_tags: self.with_tags,
             base_location: self.base_location,
             external_volume: self.external_volume,
@@ -661,6 +673,7 @@ impl From<CreateTable> for CreateTableBuilder {
             default_ddl_collation: table.default_ddl_collation,
             with_aggregation_policy: table.with_aggregation_policy,
             with_row_access_policy: table.with_row_access_policy,
+            with_storage_lifecycle_policy: table.with_storage_lifecycle_policy,
             with_tags: table.with_tags,
             base_location: table.base_location,
             external_volume: table.external_volume,

src/ast/mod.rs

@@ -10472,6 +10472,30 @@ impl Display for RowAccessPolicy {
     }
 }
 
+/// Snowflake `[ WITH ] STORAGE LIFECYCLE POLICY <policy_name> ON ( <col_name> [ , ... ] )`
+///
+/// <https://docs.snowflake.com/en/sql-reference/sql/create-table>
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct StorageLifecyclePolicy {
+    /// The fully-qualified policy object name.
+    pub policy: ObjectName,
+    /// Column names the policy applies to.
+    pub on: Vec<Ident>,
+}
+
+impl Display for StorageLifecyclePolicy {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(
+            f,
+            "WITH STORAGE LIFECYCLE POLICY {} ON ({})",
+            self.policy,
+            display_comma_separated(self.on.as_slice())
+        )
+    }
+}
+
 /// Snowflake `WITH TAG ( tag_name = '<tag_value>', ...)`
 ///
 /// <https://docs.snowflake.com/en/sql-reference/sql/create-table>

src/ast/spans.rs

@@ -572,6 +572,7 @@ impl Spanned for CreateTable {
             default_ddl_collation: _,           // string, no span
             with_aggregation_policy: _,         // todo, Snowflake specific
             with_row_access_policy: _,          // todo, Snowflake specific
+            with_storage_lifecycle_policy: _,   // todo, Snowflake specific
             with_tags: _,                       // todo, Snowflake specific
             external_volume: _,                 // todo, Snowflake specific
             base_location: _,                   // todo, Snowflake specific

src/dialect/snowflake.rs

@@ -33,7 +33,7 @@ use crate::ast::{
     IdentityPropertyFormatKind, IdentityPropertyKind, IdentityPropertyOrder, InitializeKind,
     Insert, MultiTableInsertIntoClause, MultiTableInsertType, MultiTableInsertValue,
     MultiTableInsertValues, MultiTableInsertWhenClause, ObjectName, ObjectNamePart,
-    RefreshModeKind, RowAccessPolicy, ShowObjects, SqlOption, Statement,
+    RefreshModeKind, RowAccessPolicy, ShowObjects, SqlOption, Statement, StorageLifecyclePolicy,
     StorageSerializationPolicy, TableObject, TagsColumnOption, Value, WrappedCollection,
 };
 use crate::dialect::{Dialect, Precedence};
@@ -917,6 +917,7 @@ pub fn parse_create_table(
                 Keyword::WITH => {
                     parser.expect_one_of_keywords(&[
                         Keyword::AGGREGATION,
+                        Keyword::STORAGE,
                         Keyword::TAG,
                         Keyword::ROW,
                     ])?;
@@ -938,6 +939,19 @@ pub fn parse_create_table(
                     builder =
                         builder.with_row_access_policy(Some(RowAccessPolicy::new(policy, columns)))
                 }
+                Keyword::STORAGE => {
+                    parser.expect_keywords(&[Keyword::LIFECYCLE, Keyword::POLICY])?;
+                    let policy = parser.parse_object_name(false)?;
+                    parser.expect_keyword_is(Keyword::ON)?;
+                    parser.expect_token(&Token::LParen)?;
+                    let columns = parser.parse_comma_separated(|p| p.parse_identifier())?;
+                    parser.expect_token(&Token::RParen)?;
+
+                    builder = builder.with_storage_lifecycle_policy(Some(StorageLifecyclePolicy {
+                        policy,
+                        on: columns,
+                    }))
+                }
                 Keyword::TAG => {
                     parser.expect_token(&Token::LParen)?;
                     let tags = parser.parse_comma_separated(Parser::parse_tag)?;

src/keywords.rs

@@ -573,6 +573,7 @@ define_keywords!(
     LEFT,
     LEFTARG,
     LEVEL,
+    LIFECYCLE,
     LIKE,
     LIKE_REGEX,
     LIMIT,

tests/sqlparser_duckdb.rs

@@ -776,6 +776,7 @@ fn test_duckdb_union_datatype() {
             default_ddl_collation: Default::default(),
             with_aggregation_policy: Default::default(),
             with_row_access_policy: Default::default(),
+            with_storage_lifecycle_policy: Default::default(),
             with_tags: Default::default(),
             base_location: Default::default(),
             external_volume: Default::default(),

tests/sqlparser_mssql.rs

@@ -1994,6 +1994,7 @@ fn parse_create_table_with_valid_options() {
                 default_ddl_collation: None,
                 with_aggregation_policy: None,
                 with_row_access_policy: None,
+                with_storage_lifecycle_policy: None,
                 with_tags: None,
                 base_location: None,
                 external_volume: None,
@@ -2166,6 +2167,7 @@ fn parse_create_table_with_identity_column() {
                 default_ddl_collation: None,
                 with_aggregation_policy: None,
                 with_row_access_policy: None,
+                with_storage_lifecycle_policy: None,
                 with_tags: None,
                 base_location: None,
                 external_volume: None,

tests/sqlparser_postgres.rs

@@ -6462,6 +6462,7 @@ fn parse_trigger_related_functions() {
             default_ddl_collation: None,
             with_aggregation_policy: None,
             with_row_access_policy: None,
+            with_storage_lifecycle_policy: None,
             with_tags: None,
             base_location: None,
             external_volume: None,

tests/sqlparser_snowflake.rs

@@ -286,6 +286,32 @@ fn test_snowflake_create_table_with_row_access_policy() {
     }
 }
 
+#[test]
+fn test_snowflake_create_table_with_storage_lifecycle_policy() {
+    // WITH keyword
+    match snowflake().verified_stmt(
+        "CREATE TABLE IF NOT EXISTS my_table (a NUMBER(38, 0), b VARIANT) WITH STORAGE LIFECYCLE POLICY dba.global_settings.my_policy ON (a)",
+    ) {
+        Statement::CreateTable(CreateTable {
+            name,
+            with_storage_lifecycle_policy,
+            ..
+        }) => {
+            assert_eq!("my_table", name.to_string());
+            let policy = with_storage_lifecycle_policy.unwrap();
+            assert_eq!("dba.global_settings.my_policy", policy.policy.to_string());
+            assert_eq!(vec![Ident::new("a")], policy.on);
+        }
+        _ => unreachable!(),
+    }
+
+    // Without WITH keyword — canonicalizes to WITH form
+    snowflake().one_statement_parses_to(
+        "CREATE TABLE my_table (a NUMBER(38, 0)) STORAGE LIFECYCLE POLICY my_policy ON (a, b)",
+        "CREATE TABLE my_table (a NUMBER(38, 0)) WITH STORAGE LIFECYCLE POLICY my_policy ON (a, b)",
+    );
+}
+
 #[test]
 fn test_snowflake_create_table_with_tag() {
     match snowflake()