Snowflake: Fix ScalingPolicy and WarehouseType Refs in WarehouseObjec…

…tProperties and use ObjectReferenceSegment in AlterWarehouseStatementSegment (#5264) Co-authored-by: pascal.moreau <pascal.moreau@yousign.com>
sqlfluff · Oct 2, 2023 · 5680394 · 5680394
1 parent 72dde7d
commit 5680394
Show file tree

Hide file tree

Showing 5 changed files with 229 additions and 47 deletions.
diff --git a/src/sqlfluff/dialects/dialect_snowflake.py b/src/sqlfluff/dialects/dialect_snowflake.py
@@ -152,6 +152,14 @@
     ["CSV", "JSON", "AVRO", "ORC" "PARQUET", "XML"],
 )
 
+snowflake_dialect.sets("warehouse_types").clear()
+snowflake_dialect.sets("warehouse_types").update(
+    [
+        "STANDARD",
+        "SNOWPARK-OPTIMIZED",
+    ],
+)
+
 snowflake_dialect.sets("warehouse_sizes").clear()
 snowflake_dialect.sets("warehouse_sizes").update(
     [
@@ -177,6 +185,15 @@
     ],
 )
 
+snowflake_dialect.sets("warehouse_scaling_policies").clear()
+snowflake_dialect.sets("warehouse_scaling_policies").update(
+    [
+        "STANDARD",
+        "ECONOMY",
+    ],
+)
+
+
 snowflake_dialect.add(
     # In snowflake, these are case sensitive even though they're not quoted
     # so they need a different `name` and `type` so they're not picked up
@@ -221,6 +238,22 @@
     ),
     # We use a RegexParser instead of keywords as some (those with dashes) require
     # quotes:
+    WarehouseType=OneOf(
+        MultiStringParser(
+            [
+                type
+                for type in snowflake_dialect.sets("warehouse_types")
+                if "-" not in type
+            ],
+            CodeSegment,
+            type="warehouse_size",
+        ),
+        MultiStringParser(
+            [f"'{type}'" for type in snowflake_dialect.sets("warehouse_types")],
+            CodeSegment,
+            type="warehouse_size",
+        ),
+    ),
     WarehouseSize=OneOf(
         MultiStringParser(
             [
@@ -252,6 +285,23 @@
             type="compression_type",
         ),
     ),
+    ScalingPolicy=OneOf(
+        MultiStringParser(
+            snowflake_dialect.sets("warehouse_scaling_policies"),
+            KeywordSegment,
+            type="scaling_policy",
+        ),
+        MultiStringParser(
+            [
+                f"'{scaling_policy}'"
+                for scaling_policy in snowflake_dialect.sets(
+                    "warehouse_scaling_policies"
+                )
+            ],
+            KeywordSegment,
+            type="scaling_policy",
+        ),
+    ),
     ValidationModeOptionSegment=RegexParser(
         r"'?RETURN_(?:\d+_ROWS|ERRORS|ALL_ERRORS)'?",
         CodeSegment,
@@ -1949,7 +1999,7 @@ class AlterWarehouseStatementSegment(BaseSegment):
         Sequence("IF", "EXISTS", optional=True),
         OneOf(
             Sequence(
-                Ref("NakedIdentifierSegment", optional=True),
+                Ref("ObjectReferenceSegment", optional=True),
                 OneOf(
                     "SUSPEND",
                     Sequence(
@@ -1959,21 +2009,21 @@ class AlterWarehouseStatementSegment(BaseSegment):
                 ),
             ),
             Sequence(
-                Ref("NakedIdentifierSegment", optional=True),
+                Ref("ObjectReferenceSegment", optional=True),
                 Sequence(
                     "ABORT",
                     "ALL",
                     "QUERIES",
                 ),
             ),
             Sequence(
-                Ref("NakedIdentifierSegment"),
+                Ref("ObjectReferenceSegment"),
                 "RENAME",
                 "TO",
-                Ref("NakedIdentifierSegment"),
+                Ref("ObjectReferenceSegment"),
             ),
             Sequence(
-                Ref("NakedIdentifierSegment", optional=True),
+                Ref("ObjectReferenceSegment", optional=True),
                 "SET",
                 OneOf(
                     AnyNumberOf(
@@ -1985,7 +2035,7 @@ class AlterWarehouseStatementSegment(BaseSegment):
                 ),
             ),
             Sequence(
-                Ref("NakedIdentifierSegment"),
+                Ref("ObjectReferenceSegment"),
                 "UNSET",
                 OneOf(
                     Delimited(Ref("NakedIdentifierSegment")),
@@ -2995,7 +3045,7 @@ class WarehouseObjectPropertiesSegment(BaseSegment):
         Sequence(
             "WAREHOUSE_TYPE",
             Ref("EqualsSegment"),
-            "STANDARD",
+            Ref("WarehouseType"),
         ),
         Sequence(
             "WAREHOUSE_SIZE",
@@ -3020,10 +3070,7 @@ class WarehouseObjectPropertiesSegment(BaseSegment):
         Sequence(
             "SCALING_POLICY",
             Ref("EqualsSegment"),
-            OneOf(
-                "STANDARD",
-                "ECONOMY",
-            ),
+            Ref("ScalingPolicy"),
         ),
         Sequence(
             "AUTO_SUSPEND",

diff --git a/test/fixtures/dialects/snowflake/alter_warehouse.sql b/test/fixtures/dialects/snowflake/alter_warehouse.sql
@@ -4,7 +4,13 @@ alter warehouse LOAD_WH set warehouse_size = XXLARGE;
 alter warehouse LOAD_WH set WAIT_FOR_COMPLETION = TRUE;
 alter warehouse LOAD_WH set MAX_CLUSTER_COUNT = 5;
 alter warehouse LOAD_WH set MIN_CLUSTER_COUNT = 1;
+
 alter warehouse LOAD_WH set SCALING_POLICY = STANDARD;
+alter warehouse LOAD_WH set SCALING_POLICY = 'STANDARD';
+alter warehouse LOAD_WH set SCALING_POLICY = ECONOMY;
+alter warehouse LOAD_WH set SCALING_POLICY = 'ECONOMY';
+
+
 alter warehouse LOAD_WH set AUTO_SUSPEND = 1;
 alter warehouse LOAD_WH set AUTO_RESUME = FALSE;
 alter warehouse LOAD_WH set RESOURCE_MONITOR = monitor_name;
@@ -23,4 +29,10 @@ alter warehouse LOAD_WH UNSET WAREHOUSE_SIZE;
 alter warehouse LOAD_WH UNSET WAREHOUSE_SIZE, WAIT_FOR_COMPLETION;
 
 ALTER WAREHOUSE SET WAREHOUSE_SIZE='X-LARGE';
-alter warehouse set warehouse_size=medium
+alter warehouse set warehouse_size=medium;
+
+alter warehouse LOAD_WH set WAREHOUSE_TYPE = STANDARD;
+alter warehouse LOAD_WH set WAREHOUSE_TYPE = 'SNOWPARK-OPTIMIZED';
+
+ALTER WAREHOUSE IDENTIFIER($var_wh) SET WAREHOUSE_TYPE = STANDARD;
+