From 78a92a9cd93d3abf7a61d4daade3f6c2f2fb0b07 Mon Sep 17 00:00:00 2001
From: Derick Lung'aho <83162+derickl@users.noreply.github.com>
Date: Sun, 12 Dec 2021 21:39:47 +0300
Subject: [PATCH] Allow for more value types in ALTER TABLE ALTER COLUMN SET
 DEFAULT statement (#2101)

---
 src/sqlfluff/dialects/dialect_postgres.py     |  13 +-
 .../postgres/postgres_alter_table.sql         |  20 ++
 .../postgres/postgres_alter_table.yml         | 328 +++++++++++++++++-
 3 files changed, 351 insertions(+), 10 deletions(-)

diff --git a/src/sqlfluff/dialects/dialect_postgres.py b/src/sqlfluff/dialects/dialect_postgres.py
index 374d4065ab7..ef8ba24d2a1 100644
--- a/src/sqlfluff/dialects/dialect_postgres.py
+++ b/src/sqlfluff/dialects/dialect_postgres.py
@@ -1272,7 +1272,18 @@ class AlterTableActionSegment(BaseSegment):
                     Sequence("COLLATE", Ref("QuotedLiteralSegment"), optional=True),
                     Sequence("USING", OneOf(Ref("ExpressionSegment")), optional=True),
                 ),
-                Sequence("SET", "DEFAULT", Ref("ExpressionSegment")),
+                Sequence(
+                    "SET",
+                    "DEFAULT",
+                    OneOf(
+                        OneOf(
+                            Ref("LiteralGrammar"),
+                            Ref("FunctionSegment"),
+                            Ref("BareFunctionSegment"),
+                            Ref("ExpressionSegment"),
+                        )
+                    ),
+                ),
                 Sequence("DROP", "DEFAULT"),
                 Sequence(OneOf("SET", "DROP", optional=True), "NOT", "NULL"),
                 Sequence("DROP", "EXPRESSION", Sequence("IF", "EXISTS", optional=True)),
diff --git a/test/fixtures/dialects/postgres/postgres_alter_table.sql b/test/fixtures/dialects/postgres/postgres_alter_table.sql
index 279ddacd0e6..165b3328da9 100644
--- a/test/fixtures/dialects/postgres/postgres_alter_table.sql
+++ b/test/fixtures/dialects/postgres/postgres_alter_table.sql
@@ -25,6 +25,25 @@ ALTER TABLE foo
         timestamp 'epoch' + foo_timestamp * interval '1 second',
     ALTER COLUMN foo_timestamp SET DEFAULT now();
 
+ALTER TABLE mytable ALTER date_column SET DEFAULT NOW();
+ALTER TABLE mytable ALTER int_column SET DEFAULT 1;
+ALTER TABLE mytable ALTER text_column SET DEFAULT 'value';
+ALTER TABLE mytable ALTER bool_column SET DEFAULT false;
+ALTER TABLE mytable ALTER other_column SET DEFAULT other_value;
+ALTER TABLE mytable ALTER other_column SET DEFAULT CURRENT_TIMESTAMP;
+ALTER TABLE mytable ALTER other_column SET DEFAULT a_function(a_parameter);
+ALTER TABLE mytable ALTER other_column SET DEFAULT a_function('a_parameter');
+ALTER TABLE mytable ALTER other_column DROP DEFAULT;
+ALTER TABLE IF EXISTS mytable ALTER date_column SET DEFAULT NOW();
+ALTER TABLE IF EXISTS mytable ALTER int_column SET DEFAULT 1;
+ALTER TABLE IF EXISTS mytable ALTER text_column SET DEFAULT 'value';
+ALTER TABLE IF EXISTS mytable ALTER bool_column SET DEFAULT false;
+ALTER TABLE IF EXISTS mytable ALTER other_column SET DEFAULT other_value;
+ALTER TABLE IF EXISTS mytable ALTER other_column SET DEFAULT CURRENT_TIMESTAMP;
+ALTER TABLE IF EXISTS mytable ALTER other_column SET DEFAULT a_function(a_parameter);
+ALTER TABLE IF EXISTS mytable ALTER other_column SET DEFAULT a_function('a_parameter');
+ALTER TABLE IF EXISTS mytable ALTER other_column DROP DEFAULT;
+
 ALTER TABLE distributors RENAME COLUMN address TO city;
 
 ALTER TABLE distributors RENAME TO suppliers;
@@ -109,3 +128,4 @@ ALTER TABLE IF EXISTS ONLY public.test OWNER TO CURRENT_ROLE;
 ALTER TABLE public.history ALTER COLUMN id ADD GENERATED ALWAYS AS IDENTITY (
     SEQUENCE NAME public.history_id_seq
 );
+
diff --git a/test/fixtures/dialects/postgres/postgres_alter_table.yml b/test/fixtures/dialects/postgres/postgres_alter_table.yml
index 3fdc7c88aae..19739198585 100644
--- a/test/fixtures/dialects/postgres/postgres_alter_table.yml
+++ b/test/fixtures/dialects/postgres/postgres_alter_table.yml
@@ -3,7 +3,7 @@
 # computed by SQLFluff when running the tests. Please run
 # `python test/generate_parse_fixture_yml.py`  to generate them after adding or
 # altering SQL files.
-_hash: 93dec596055d5d295adb944c2c4cf29a0e1e219d3e18bc7d62381bc724e2c1c9
+_hash: 4b4d674351c381ca18130cf851463e859022bf7d854f8fe92c1a343857014cd1
 file:
 - statement:
     alter_table_statement:
@@ -79,8 +79,7 @@ file:
           identifier: status
       - keyword: SET
       - keyword: default
-      - expression:
-          literal: "'current'"
+      - literal: "'current'"
 - statement_terminator: ;
 - statement:
     alter_table_statement:
@@ -215,13 +214,324 @@ file:
           identifier: foo_timestamp
       - keyword: SET
       - keyword: DEFAULT
+      - function:
+          function_name:
+            function_name_identifier: now
+          bracketed:
+            start_bracket: (
+            end_bracket: )
+- statement_terminator: ;
+- statement:
+    alter_table_statement:
+    - keyword: ALTER
+    - keyword: TABLE
+    - table_reference:
+        identifier: mytable
+    - alter_table_action_segment:
+      - keyword: ALTER
+      - column_reference:
+          identifier: date_column
+      - keyword: SET
+      - keyword: DEFAULT
+      - function:
+          function_name:
+            function_name_identifier: NOW
+          bracketed:
+            start_bracket: (
+            end_bracket: )
+- statement_terminator: ;
+- statement:
+    alter_table_statement:
+    - keyword: ALTER
+    - keyword: TABLE
+    - table_reference:
+        identifier: mytable
+    - alter_table_action_segment:
+      - keyword: ALTER
+      - column_reference:
+          identifier: int_column
+      - keyword: SET
+      - keyword: DEFAULT
+      - literal: '1'
+- statement_terminator: ;
+- statement:
+    alter_table_statement:
+    - keyword: ALTER
+    - keyword: TABLE
+    - table_reference:
+        identifier: mytable
+    - alter_table_action_segment:
+      - keyword: ALTER
+      - column_reference:
+          identifier: text_column
+      - keyword: SET
+      - keyword: DEFAULT
+      - literal: "'value'"
+- statement_terminator: ;
+- statement:
+    alter_table_statement:
+    - keyword: ALTER
+    - keyword: TABLE
+    - table_reference:
+        identifier: mytable
+    - alter_table_action_segment:
+      - keyword: ALTER
+      - column_reference:
+          identifier: bool_column
+      - keyword: SET
+      - keyword: DEFAULT
+      - literal: 'false'
+- statement_terminator: ;
+- statement:
+    alter_table_statement:
+    - keyword: ALTER
+    - keyword: TABLE
+    - table_reference:
+        identifier: mytable
+    - alter_table_action_segment:
+      - keyword: ALTER
+      - column_reference:
+          identifier: other_column
+      - keyword: SET
+      - keyword: DEFAULT
       - expression:
-          function:
-            function_name:
-              function_name_identifier: now
-            bracketed:
-              start_bracket: (
-              end_bracket: )
+          column_reference:
+            identifier: other_value
+- statement_terminator: ;
+- statement:
+    alter_table_statement:
+    - keyword: ALTER
+    - keyword: TABLE
+    - table_reference:
+        identifier: mytable
+    - alter_table_action_segment:
+      - keyword: ALTER
+      - column_reference:
+          identifier: other_column
+      - keyword: SET
+      - keyword: DEFAULT
+      - bare_function: CURRENT_TIMESTAMP
+- statement_terminator: ;
+- statement:
+    alter_table_statement:
+    - keyword: ALTER
+    - keyword: TABLE
+    - table_reference:
+        identifier: mytable
+    - alter_table_action_segment:
+      - keyword: ALTER
+      - column_reference:
+          identifier: other_column
+      - keyword: SET
+      - keyword: DEFAULT
+      - function:
+          function_name:
+            function_name_identifier: a_function
+          bracketed:
+            start_bracket: (
+            expression:
+              column_reference:
+                identifier: a_parameter
+            end_bracket: )
+- statement_terminator: ;
+- statement:
+    alter_table_statement:
+    - keyword: ALTER
+    - keyword: TABLE
+    - table_reference:
+        identifier: mytable
+    - alter_table_action_segment:
+      - keyword: ALTER
+      - column_reference:
+          identifier: other_column
+      - keyword: SET
+      - keyword: DEFAULT
+      - function:
+          function_name:
+            function_name_identifier: a_function
+          bracketed:
+            start_bracket: (
+            expression:
+              literal: "'a_parameter'"
+            end_bracket: )
+- statement_terminator: ;
+- statement:
+    alter_table_statement:
+    - keyword: ALTER
+    - keyword: TABLE
+    - table_reference:
+        identifier: mytable
+    - alter_table_action_segment:
+      - keyword: ALTER
+      - column_reference:
+          identifier: other_column
+      - keyword: DROP
+      - keyword: DEFAULT
+- statement_terminator: ;
+- statement:
+    alter_table_statement:
+    - keyword: ALTER
+    - keyword: TABLE
+    - keyword: IF
+    - keyword: EXISTS
+    - table_reference:
+        identifier: mytable
+    - alter_table_action_segment:
+      - keyword: ALTER
+      - column_reference:
+          identifier: date_column
+      - keyword: SET
+      - keyword: DEFAULT
+      - function:
+          function_name:
+            function_name_identifier: NOW
+          bracketed:
+            start_bracket: (
+            end_bracket: )
+- statement_terminator: ;
+- statement:
+    alter_table_statement:
+    - keyword: ALTER
+    - keyword: TABLE
+    - keyword: IF
+    - keyword: EXISTS
+    - table_reference:
+        identifier: mytable
+    - alter_table_action_segment:
+      - keyword: ALTER
+      - column_reference:
+          identifier: int_column
+      - keyword: SET
+      - keyword: DEFAULT
+      - literal: '1'
+- statement_terminator: ;
+- statement:
+    alter_table_statement:
+    - keyword: ALTER
+    - keyword: TABLE
+    - keyword: IF
+    - keyword: EXISTS
+    - table_reference:
+        identifier: mytable
+    - alter_table_action_segment:
+      - keyword: ALTER
+      - column_reference:
+          identifier: text_column
+      - keyword: SET
+      - keyword: DEFAULT
+      - literal: "'value'"
+- statement_terminator: ;
+- statement:
+    alter_table_statement:
+    - keyword: ALTER
+    - keyword: TABLE
+    - keyword: IF
+    - keyword: EXISTS
+    - table_reference:
+        identifier: mytable
+    - alter_table_action_segment:
+      - keyword: ALTER
+      - column_reference:
+          identifier: bool_column
+      - keyword: SET
+      - keyword: DEFAULT
+      - literal: 'false'
+- statement_terminator: ;
+- statement:
+    alter_table_statement:
+    - keyword: ALTER
+    - keyword: TABLE
+    - keyword: IF
+    - keyword: EXISTS
+    - table_reference:
+        identifier: mytable
+    - alter_table_action_segment:
+      - keyword: ALTER
+      - column_reference:
+          identifier: other_column
+      - keyword: SET
+      - keyword: DEFAULT
+      - expression:
+          column_reference:
+            identifier: other_value
+- statement_terminator: ;
+- statement:
+    alter_table_statement:
+    - keyword: ALTER
+    - keyword: TABLE
+    - keyword: IF
+    - keyword: EXISTS
+    - table_reference:
+        identifier: mytable
+    - alter_table_action_segment:
+      - keyword: ALTER
+      - column_reference:
+          identifier: other_column
+      - keyword: SET
+      - keyword: DEFAULT
+      - bare_function: CURRENT_TIMESTAMP
+- statement_terminator: ;
+- statement:
+    alter_table_statement:
+    - keyword: ALTER
+    - keyword: TABLE
+    - keyword: IF
+    - keyword: EXISTS
+    - table_reference:
+        identifier: mytable
+    - alter_table_action_segment:
+      - keyword: ALTER
+      - column_reference:
+          identifier: other_column
+      - keyword: SET
+      - keyword: DEFAULT
+      - function:
+          function_name:
+            function_name_identifier: a_function
+          bracketed:
+            start_bracket: (
+            expression:
+              column_reference:
+                identifier: a_parameter
+            end_bracket: )
+- statement_terminator: ;
+- statement:
+    alter_table_statement:
+    - keyword: ALTER
+    - keyword: TABLE
+    - keyword: IF
+    - keyword: EXISTS
+    - table_reference:
+        identifier: mytable
+    - alter_table_action_segment:
+      - keyword: ALTER
+      - column_reference:
+          identifier: other_column
+      - keyword: SET
+      - keyword: DEFAULT
+      - function:
+          function_name:
+            function_name_identifier: a_function
+          bracketed:
+            start_bracket: (
+            expression:
+              literal: "'a_parameter'"
+            end_bracket: )
+- statement_terminator: ;
+- statement:
+    alter_table_statement:
+    - keyword: ALTER
+    - keyword: TABLE
+    - keyword: IF
+    - keyword: EXISTS
+    - table_reference:
+        identifier: mytable
+    - alter_table_action_segment:
+      - keyword: ALTER
+      - column_reference:
+          identifier: other_column
+      - keyword: DROP
+      - keyword: DEFAULT
 - statement_terminator: ;
 - statement:
     alter_table_statement: