support ALTER TABLE DROP EXPRESSION for virtual generated column

From b3b23e5d7fee6143521560790bf4ad14e21e8a49 Mon Sep 17 00:00:00 2001
From: jian he <jian.universality@gmail.com>
Date: Thu, 27 Mar 2025 09:49:29 +0800
Subject: [PATCH v1 1/1] support ALTER TABLE DROP EXPRESSION for virtual
 generated column

It seems fairly straightforward:
drop the generation expression in ATExecDropExpression, and instruct phase 3 to
compute the generation expression and do the table rewrite.

discussion: https://postgr.es/m/
commitfest entry:
---
 doc/src/sgml/ref/alter_table.sgml             | 14 ++--
 src/backend/commands/tablecmds.c              | 36 +++++----
 src/test/regress/expected/fast_default.out    |  3 +
 .../regress/expected/generated_virtual.out    | 81 ++++++++++++++-----
 src/test/regress/sql/fast_default.sql         |  2 +
 src/test/regress/sql/generated_virtual.sql    | 26 +++++-
 6 files changed, 114 insertions(+), 48 deletions(-)

diff --git a/doc/src/sgml/ref/alter_table.sgml b/doc/src/sgml/ref/alter_table.sgml
index 4f15b89a98f..cd68697e215 100644
--- a/doc/src/sgml/ref/alter_table.sgml
+++ b/doc/src/sgml/ref/alter_table.sgml
@@ -275,14 +275,12 @@ WITH ( MODULUS <replaceable class="parameter">numeric_literal</replaceable>, REM
     <term><literal>DROP EXPRESSION [ IF EXISTS ]</literal></term>
     <listitem>
      <para>
-      This form turns a stored generated column into a normal base column.
-      Existing data in the columns is retained, but future changes will no
-      longer apply the generation expression.
-     </para>
-
-     <para>
-      This form is currently only supported for stored generated columns (not
-      virtual ones).
+      This form turns a generated column into a normal base column.
+      For stored generated column, existing data in the columns is retained;
+      For virtual generated column, it will compute the generation expression and
+      store the value in the columns. For inheritance hierarchy or partition hierarchy,
+      the virtual generation expression is computed based on the child's own generation expression.
+      The future changes will no longer apply the generation expression.
      </para>
 
      <para>
diff --git a/src/backend/commands/tablecmds.c b/src/backend/commands/tablecmds.c
index 1202544ebd0..f1f5abf50de 100644
--- a/src/backend/commands/tablecmds.c
+++ b/src/backend/commands/tablecmds.c
@@ -494,7 +494,8 @@ static ObjectAddress ATExecDropIdentity(Relation rel, const char *colName, bool
 static ObjectAddress ATExecSetExpression(AlteredTableInfo *tab, Relation rel, const char *colName,
 										 Node *newExpr, LOCKMODE lockmode);
 static void ATPrepDropExpression(Relation rel, AlterTableCmd *cmd, bool recurse, bool recursing, LOCKMODE lockmode);
-static ObjectAddress ATExecDropExpression(Relation rel, const char *colName, bool missing_ok, LOCKMODE lockmode);
+static ObjectAddress ATExecDropExpression(AlteredTableInfo *tab, Relation rel, const char *colName,
+										  bool missing_ok, LOCKMODE lockmode);
 static ObjectAddress ATExecSetStatistics(Relation rel, const char *colName, int16 colNum,
 										 Node *newValue, LOCKMODE lockmode);
 static ObjectAddress ATExecSetOptions(Relation rel, const char *colName,
@@ -5381,7 +5382,7 @@ ATExecCmd(List **wqueue, AlteredTableInfo *tab,
 			address = ATExecSetExpression(tab, rel, cmd->name, cmd->def, lockmode);
 			break;
 		case AT_DropExpression:
-			address = ATExecDropExpression(rel, cmd->name, cmd->missing_ok, lockmode);
+			address = ATExecDropExpression(tab, rel, cmd->name, cmd->missing_ok, lockmode);
 			break;
 		case AT_SetStatistics:	/* ALTER COLUMN SET STATISTICS */
 			address = ATExecSetStatistics(rel, cmd->name, cmd->num, cmd->def, lockmode);
@@ -8658,7 +8659,8 @@ ATPrepDropExpression(Relation rel, AlterTableCmd *cmd, bool recurse, bool recurs
  * Return the address of the affected column.
  */
 static ObjectAddress
-ATExecDropExpression(Relation rel, const char *colName, bool missing_ok, LOCKMODE lockmode)
+ATExecDropExpression(AlteredTableInfo *tab, Relation rel, const char *colName,
+					 bool missing_ok, LOCKMODE lockmode)
 {
 	HeapTuple	tuple;
 	Form_pg_attribute attTup;
@@ -8684,19 +8686,6 @@ ATExecDropExpression(Relation rel, const char *colName, bool missing_ok, LOCKMOD
 				 errmsg("cannot alter system column \"%s\"",
 						colName)));
 
-	/*
-	 * TODO: This could be done, but it would need a table rewrite to
-	 * materialize the generated values.  Note that for the time being, we
-	 * still error with missing_ok, so that we don't silently leave the column
-	 * as generated.
-	 */
-	if (attTup->attgenerated == ATTRIBUTE_GENERATED_VIRTUAL)
-		ereport(ERROR,
-				(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
-				 errmsg("ALTER TABLE / DROP EXPRESSION is not supported for virtual generated columns"),
-				 errdetail("Column \"%s\" of relation \"%s\" is a virtual generated column.",
-						   colName, RelationGetRelationName(rel))));
-
 	if (!attTup->attgenerated)
 	{
 		if (!missing_ok)
@@ -8715,6 +8704,21 @@ ATExecDropExpression(Relation rel, const char *colName, bool missing_ok, LOCKMOD
 		}
 	}
 
+	if (attTup->attgenerated == ATTRIBUTE_GENERATED_VIRTUAL)
+	{
+		NewColumnValue *newval;
+		Expr	   *defval;
+
+		newval = (NewColumnValue *) palloc0(sizeof(NewColumnValue));
+		newval->attnum = attnum;
+		defval = (Expr *) build_generation_expression(rel, attnum);
+		newval->expr = defval;
+		newval->is_generated = true;
+
+		tab->newvals = lappend(tab->newvals, newval);
+		tab->rewrite |= AT_REWRITE_DEFAULT_VAL;
+	}
+
 	/*
 	 * Mark the column as no longer generated.  (The atthasdef flag needs to
 	 * get cleared too, but RemoveAttrDefault will handle that.)
diff --git a/src/test/regress/expected/fast_default.out b/src/test/regress/expected/fast_default.out
index ccbcdf8403f..7e99b822f11 100644
--- a/src/test/regress/expected/fast_default.out
+++ b/src/test/regress/expected/fast_default.out
@@ -70,6 +70,9 @@ NOTICE:  rewriting table has_volatile for reason 4
 -- stored generated columns need a rewrite
 ALTER TABLE has_volatile ADD col7 int GENERATED ALWAYS AS (55) stored;
 NOTICE:  rewriting table has_volatile for reason 2
+-- drop generation expression over virtual generated colum need rewrite
+ALTER TABLE has_volatile ALTER COLUMN col6 DROP EXPRESSION;
+NOTICE:  rewriting table has_volatile for reason 2
 -- Test a large sample of different datatypes
 CREATE TABLE T(pk INT NOT NULL PRIMARY KEY, c_int INT DEFAULT 1);
 SELECT set('t');
diff --git a/src/test/regress/expected/generated_virtual.out b/src/test/regress/expected/generated_virtual.out
index dc09c85938e..959cf4608ab 100644
--- a/src/test/regress/expected/generated_virtual.out
+++ b/src/test/regress/expected/generated_virtual.out
@@ -1133,27 +1133,24 @@ SELECT * FROM gtest29;
  a      | integer |           |          | 
  b      | integer |           |          | generated always as (a * 3)
 
-ALTER TABLE gtest29 ALTER COLUMN b DROP EXPRESSION;  -- not supported
-ERROR:  ALTER TABLE / DROP EXPRESSION is not supported for virtual generated columns
-DETAIL:  Column "b" of relation "gtest29" is a virtual generated column.
+ALTER TABLE gtest29 ALTER COLUMN b DROP EXPRESSION;
 INSERT INTO gtest29 (a) VALUES (5);
 INSERT INTO gtest29 (a, b) VALUES (6, 66);
-ERROR:  cannot insert a non-DEFAULT value into column "b"
-DETAIL:  Column "b" is a generated column.
-SELECT * FROM gtest29;
+SELECT * FROM gtest29 ORDER BY a, b;
  a | b  
 ---+----
  3 |  9
  4 | 12
- 5 | 15
-(3 rows)
+ 5 |   
+ 6 | 66
+(4 rows)
 
 \d gtest29
-                Table "generated_virtual_tests.gtest29"
- Column |  Type   | Collation | Nullable |           Default           
---------+---------+-----------+----------+-----------------------------
+      Table "generated_virtual_tests.gtest29"
+ Column |  Type   | Collation | Nullable | Default 
+--------+---------+-----------+----------+---------
  a      | integer |           |          | 
- b      | integer |           |          | generated always as (a * 3)
+ b      | integer |           |          | 
 
 -- check that dependencies between columns have also been removed
 --ALTER TABLE gtest29 DROP COLUMN a;  -- should not drop b
@@ -1165,22 +1162,20 @@ CREATE TABLE gtest30 (
 );
 CREATE TABLE gtest30_1 () INHERITS (gtest30);
 ALTER TABLE gtest30 ALTER COLUMN b DROP EXPRESSION;
-ERROR:  ALTER TABLE / DROP EXPRESSION is not supported for virtual generated columns
-DETAIL:  Column "b" of relation "gtest30" is a virtual generated column.
 \d gtest30
-                Table "generated_virtual_tests.gtest30"
- Column |  Type   | Collation | Nullable |           Default           
---------+---------+-----------+----------+-----------------------------
+      Table "generated_virtual_tests.gtest30"
+ Column |  Type   | Collation | Nullable | Default 
+--------+---------+-----------+----------+---------
  a      | integer |           |          | 
- b      | integer |           |          | generated always as (a * 2)
+ b      | integer |           |          | 
 Number of child tables: 1 (Use \d+ to list them.)
 
 \d gtest30_1
-               Table "generated_virtual_tests.gtest30_1"
- Column |  Type   | Collation | Nullable |           Default           
---------+---------+-----------+----------+-----------------------------
+     Table "generated_virtual_tests.gtest30_1"
+ Column |  Type   | Collation | Nullable | Default 
+--------+---------+-----------+----------+---------
  a      | integer |           |          | 
- b      | integer |           |          | generated always as (a * 2)
+ b      | integer |           |          | 
 Inherits: gtest30
 
 DROP TABLE gtest30 CASCADE;
@@ -1210,6 +1205,48 @@ Inherits: gtest30
 
 ALTER TABLE gtest30_1 ALTER COLUMN b DROP EXPRESSION;  -- error
 ERROR:  cannot drop generation expression from inherited column
+-- alter table drop expression with partitioning
+CREATE TABLE gtest_pp (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE (f1);
+CREATE TABLE gtest_ch1 PARTITION OF gtest_pp
+  FOR VALUES FROM ('2016-07-01') TO ('2016-08-01');  -- inherits gen expr
+CREATE TABLE gtest_ch2 PARTITION OF gtest_pp (
+    f3 WITH OPTIONS GENERATED ALWAYS AS (f2 * 22) VIRTUAL  -- overrides gen expr
+) FOR VALUES FROM ('2016-08-01') TO ('2016-09-01');
+INSERT INTO gtest_pp (f1, f2) VALUES ('2016-07-15', 1), ('2016-07-15', 2), ('2016-08-15', 3);
+ALTER TABLE ONLY gtest_pp ALTER COLUMN f3 SET EXPRESSION AS (f2 * 4);
+ALTER TABLE gtest_ch1 ALTER COLUMN f3 SET EXPRESSION AS (f2 * 10);
+SELECT tableoid::regclass, * FROM gtest_pp ORDER BY 1, 2, 3;
+ tableoid  |     f1     | f2 | f3 
+-----------+------------+----+----
+ gtest_ch1 | 07-15-2016 |  1 |  4
+ gtest_ch1 | 07-15-2016 |  2 |  8
+ gtest_ch2 | 08-15-2016 |  3 | 12
+(3 rows)
+
+ALTER TABLE gtest_ch1 ALTER COLUMN f3 DROP EXPRESSION; --error
+ERROR:  cannot drop generation expression from inherited column
+ALTER TABLE ONLY gtest_pp ALTER COLUMN f3 DROP EXPRESSION; --error
+ERROR:  ALTER TABLE / DROP EXPRESSION must be applied to child tables too
+ALTER TABLE gtest_pp ALTER COLUMN f3 DROP EXPRESSION; --ok
+SELECT tableoid::regclass, * FROM gtest_pp ORDER BY 1, 2, 3;
+ tableoid  |     f1     | f2 | f3 
+-----------+------------+----+----
+ gtest_ch1 | 07-15-2016 |  1 | 10
+ gtest_ch1 | 07-15-2016 |  2 | 20
+ gtest_ch2 | 08-15-2016 |  3 | 66
+(3 rows)
+
+\d gtest_pp
+Partitioned table "generated_virtual_tests.gtest_pp"
+ Column |  Type  | Collation | Nullable | Default 
+--------+--------+-----------+----------+---------
+ f1     | date   |           | not null | 
+ f2     | bigint |           |          | 
+ f3     | bigint |           |          | 
+Partition key: RANGE (f1)
+Number of partitions: 2 (Use \d+ to list them.)
+
+DROP TABLE gtest_pp;
 -- composite type dependencies
 CREATE TABLE gtest31_1 (a int, b text GENERATED ALWAYS AS ('hello') VIRTUAL, c text);
 CREATE TABLE gtest31_2 (x int, y gtest31_1);
diff --git a/src/test/regress/sql/fast_default.sql b/src/test/regress/sql/fast_default.sql
index 068dd0bc8aa..c0ef10b1ca3 100644
--- a/src/test/regress/sql/fast_default.sql
+++ b/src/test/regress/sql/fast_default.sql
@@ -77,6 +77,8 @@ ALTER TABLE has_volatile ALTER COLUMN col1 SET DATA TYPE float8,
 -- stored generated columns need a rewrite
 ALTER TABLE has_volatile ADD col7 int GENERATED ALWAYS AS (55) stored;
 
+-- drop generation expression over virtual generated colum need rewrite
+ALTER TABLE has_volatile ALTER COLUMN col6 DROP EXPRESSION;
 
 
 -- Test a large sample of different datatypes
diff --git a/src/test/regress/sql/generated_virtual.sql b/src/test/regress/sql/generated_virtual.sql
index dab8c92ef99..d278a868107 100644
--- a/src/test/regress/sql/generated_virtual.sql
+++ b/src/test/regress/sql/generated_virtual.sql
@@ -562,10 +562,10 @@ ALTER TABLE gtest29 ALTER COLUMN b SET EXPRESSION AS (a * 3);
 SELECT * FROM gtest29;
 \d gtest29
 
-ALTER TABLE gtest29 ALTER COLUMN b DROP EXPRESSION;  -- not supported
+ALTER TABLE gtest29 ALTER COLUMN b DROP EXPRESSION;
 INSERT INTO gtest29 (a) VALUES (5);
 INSERT INTO gtest29 (a, b) VALUES (6, 66);
-SELECT * FROM gtest29;
+SELECT * FROM gtest29 ORDER BY a, b;
 \d gtest29
 
 -- check that dependencies between columns have also been removed
@@ -592,6 +592,28 @@ ALTER TABLE ONLY gtest30 ALTER COLUMN b DROP EXPRESSION;  -- error
 \d gtest30_1
 ALTER TABLE gtest30_1 ALTER COLUMN b DROP EXPRESSION;  -- error
 
+-- alter table drop expression with partitioning
+CREATE TABLE gtest_pp (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE (f1);
+CREATE TABLE gtest_ch1 PARTITION OF gtest_pp
+  FOR VALUES FROM ('2016-07-01') TO ('2016-08-01');  -- inherits gen expr
+CREATE TABLE gtest_ch2 PARTITION OF gtest_pp (
+    f3 WITH OPTIONS GENERATED ALWAYS AS (f2 * 22) VIRTUAL  -- overrides gen expr
+) FOR VALUES FROM ('2016-08-01') TO ('2016-09-01');
+INSERT INTO gtest_pp (f1, f2) VALUES ('2016-07-15', 1), ('2016-07-15', 2), ('2016-08-15', 3);
+
+ALTER TABLE ONLY gtest_pp ALTER COLUMN f3 SET EXPRESSION AS (f2 * 4);
+ALTER TABLE gtest_ch1 ALTER COLUMN f3 SET EXPRESSION AS (f2 * 10);
+
+SELECT tableoid::regclass, * FROM gtest_pp ORDER BY 1, 2, 3;
+
+ALTER TABLE gtest_ch1 ALTER COLUMN f3 DROP EXPRESSION; --error
+ALTER TABLE ONLY gtest_pp ALTER COLUMN f3 DROP EXPRESSION; --error
+
+ALTER TABLE gtest_pp ALTER COLUMN f3 DROP EXPRESSION; --ok
+SELECT tableoid::regclass, * FROM gtest_pp ORDER BY 1, 2, 3;
+\d gtest_pp
+DROP TABLE gtest_pp;
+
 -- composite type dependencies
 CREATE TABLE gtest31_1 (a int, b text GENERATED ALWAYS AS ('hello') VIRTUAL, c text);
 CREATE TABLE gtest31_2 (x int, y gtest31_1);
-- 
2.34.1