From b565c793914275c6daf4b169a7d7cde3a2a028aa Mon Sep 17 00:00:00 2001
From: Amul Sul <amul.sul@enterprisedb.com>
Date: Wed, 15 Nov 2023 18:04:51 +0530
Subject: [PATCH v4 1/2] Allow to change generated column expression

---
 doc/src/sgml/ref/alter_table.sgml             |  13 ++
 src/backend/commands/tablecmds.c              | 111 +++++++++++-
 src/backend/parser/gram.y                     |  10 ++
 src/bin/psql/tab-complete.c                   |  10 +-
 src/include/nodes/parsenodes.h                |   1 +
 .../test_ddl_deparse/test_ddl_deparse.c       |   3 +
 src/test/regress/expected/generated.out       | 167 +++++++++++++++---
 src/test/regress/sql/generated.sql            |  36 +++-
 8 files changed, 321 insertions(+), 30 deletions(-)

diff --git a/doc/src/sgml/ref/alter_table.sgml b/doc/src/sgml/ref/alter_table.sgml
index 2c4138e4e9f..c5c0ac6e276 100644
--- a/doc/src/sgml/ref/alter_table.sgml
+++ b/doc/src/sgml/ref/alter_table.sgml
@@ -46,6 +46,7 @@ ALTER TABLE [ IF EXISTS ] <replaceable class="parameter">name</replaceable>
     ALTER [ COLUMN ] <replaceable class="parameter">column_name</replaceable> SET DEFAULT <replaceable class="parameter">expression</replaceable>
     ALTER [ COLUMN ] <replaceable class="parameter">column_name</replaceable> DROP DEFAULT
     ALTER [ COLUMN ] <replaceable class="parameter">column_name</replaceable> { SET | DROP } NOT NULL
+    ALTER [ COLUMN ] <replaceable class="parameter">column_name</replaceable> SET EXPRESSION AS ( <replaceable class="parameter">expression</replaceable> )
     ALTER [ COLUMN ] <replaceable class="parameter">column_name</replaceable> DROP EXPRESSION [ IF EXISTS ]
     ALTER [ COLUMN ] <replaceable class="parameter">column_name</replaceable> ADD GENERATED { ALWAYS | BY DEFAULT } AS IDENTITY [ ( <replaceable>sequence_options</replaceable> ) ]
     ALTER [ COLUMN ] <replaceable class="parameter">column_name</replaceable> { SET GENERATED { ALWAYS | BY DEFAULT } | SET <replaceable>sequence_option</replaceable> | RESTART [ [ WITH ] <replaceable class="parameter">restart</replaceable> ] } [...]
@@ -256,6 +257,18 @@ WITH ( MODULUS <replaceable class="parameter">numeric_literal</replaceable>, REM
     </listitem>
    </varlistentry>
 
+   <varlistentry id="sql-altertable-desc-set-expression">
+    <term><literal>SET EXPRESSION AS</literal></term>
+    <listitem>
+     <para>
+      This form replaces the expression of a generated column. Existing
+      data in the column is rewritten and all the future changes will
+      apply the new
+      generation expression.
+     </para>
+    </listitem>
+   </varlistentry>
+
    <varlistentry id="sql-altertable-desc-drop-expression">
     <term><literal>DROP EXPRESSION [ IF EXISTS ]</literal></term>
     <listitem>
diff --git a/src/backend/commands/tablecmds.c b/src/backend/commands/tablecmds.c
index 323d9bf8702..63debf1c376 100644
--- a/src/backend/commands/tablecmds.c
+++ b/src/backend/commands/tablecmds.c
@@ -455,6 +455,9 @@ static ObjectAddress ATExecAddIdentity(Relation rel, const char *colName,
 static ObjectAddress ATExecSetIdentity(Relation rel, const char *colName,
 									   Node *def, LOCKMODE lockmode);
 static ObjectAddress ATExecDropIdentity(Relation rel, const char *colName, bool missing_ok, LOCKMODE lockmode);
+static ObjectAddress ATExecSetExpression(AlteredTableInfo *tab, Relation rel,
+										 const char *colName, Node *newDefault,
+										 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 ATExecSetStatistics(Relation rel, const char *colName, int16 colNum,
@@ -4548,6 +4551,7 @@ AlterTableGetLockLevel(List *cmds)
 			case AT_AddIdentity:
 			case AT_DropIdentity:
 			case AT_SetIdentity:
+			case AT_SetExpression:
 			case AT_DropExpression:
 			case AT_SetCompression:
 				cmd_lockmode = AccessExclusiveLock;
@@ -4849,6 +4853,11 @@ ATPrepCmd(List **wqueue, Relation rel, AlterTableCmd *cmd,
 			ATSimpleRecursion(wqueue, rel, cmd, recurse, lockmode, context);
 			pass = AT_PASS_COL_ATTRS;
 			break;
+		case AT_SetExpression:	/* ALTER COLUMN SET EXPRESSION AS */
+			ATSimplePermissions(cmd->subtype, rel, ATT_TABLE | ATT_FOREIGN_TABLE);
+			ATSimpleRecursion(wqueue, rel, cmd, recurse, lockmode, context);
+			pass = AT_PASS_MISC;
+			break;
 		case AT_DropExpression: /* ALTER COLUMN DROP EXPRESSION */
 			ATSimplePermissions(cmd->subtype, rel, ATT_TABLE | ATT_FOREIGN_TABLE);
 			ATSimpleRecursion(wqueue, rel, cmd, recurse, lockmode, context);
@@ -5234,6 +5243,9 @@ ATExecCmd(List **wqueue, AlteredTableInfo *tab,
 		case AT_SetAttNotNull:	/* set pg_attribute.attnotnull */
 			address = ATExecSetAttNotNull(wqueue, rel, cmd->name, lockmode);
 			break;
+		case AT_SetExpression:
+			address = ATExecSetExpression(tab, rel, cmd->name, cmd->def, lockmode);
+			break;
 		case AT_DropExpression:
 			address = ATExecDropExpression(rel, cmd->name, cmd->missing_ok, lockmode);
 			break;
@@ -6361,6 +6373,8 @@ alter_table_type_to_string(AlterTableType cmdtype)
 			return "ALTER COLUMN ... SET NOT NULL";
 		case AT_SetAttNotNull:
 			return NULL;		/* not real grammar */
+		case AT_SetExpression:
+			return "ALTER COLUMN ... SET EXPRESSION AS";
 		case AT_DropExpression:
 			return "ALTER COLUMN ... DROP EXPRESSION";
 		case AT_SetStatistics:
@@ -8011,8 +8025,10 @@ ATExecColumnDefault(Relation rel, const char *colName,
 				(errcode(ERRCODE_SYNTAX_ERROR),
 				 errmsg("column \"%s\" of relation \"%s\" is a generated column",
 						colName, RelationGetRelationName(rel)),
-				 newDefault || TupleDescAttr(tupdesc, attnum - 1)->attgenerated != ATTRIBUTE_GENERATED_STORED ? 0 :
+				 TupleDescAttr(tupdesc, attnum - 1)->attgenerated != ATTRIBUTE_GENERATED_STORED ? 0 :
 		/* translator: %s is an SQL ALTER command */
+				 newDefault ? errhint("Use %s instead.",
+						 "ALTER TABLE ... ALTER COLUMN ... SET EXPRESSION") :
 				 errhint("Use %s instead.",
 						 "ALTER TABLE ... ALTER COLUMN ... DROP EXPRESSION")));
 
@@ -8311,6 +8327,99 @@ ATExecDropIdentity(Relation rel, const char *colName, bool missing_ok, LOCKMODE
 	return address;
 }
 
+/*
+ * ALTER TABLE ALTER COLUMN SET EXPRESSION AS
+ *
+ * Return the address of the affected column.
+ */
+static ObjectAddress
+ATExecSetExpression(AlteredTableInfo *tab, Relation rel, const char *colName,
+					Node *newDefault, LOCKMODE lockmode)
+{
+	HeapTuple	tuple;
+	Form_pg_attribute attTup;
+	AttrNumber	attnum;
+	Oid			attrdefoid;
+	ObjectAddress address;
+	Expr	   *defval;
+	NewColumnValue *newval;
+	RawColumnDefault *rawEnt;
+
+	tuple = SearchSysCacheAttName(RelationGetRelid(rel), colName);
+	if (!HeapTupleIsValid(tuple))
+		ereport(ERROR,
+				(errcode(ERRCODE_UNDEFINED_COLUMN),
+				 errmsg("column \"%s\" of relation \"%s\" does not exist",
+						colName, RelationGetRelationName(rel))));
+
+	attTup = (Form_pg_attribute) GETSTRUCT(tuple);
+	attnum = attTup->attnum;
+
+	if (attnum <= 0)
+		ereport(ERROR,
+				(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+				 errmsg("cannot alter system column \"%s\"",
+						colName)));
+
+	if (attTup->attgenerated != ATTRIBUTE_GENERATED_STORED)
+		ereport(ERROR,
+				(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
+				 errmsg("column \"%s\" of relation \"%s\" is not a stored generated column",
+						colName, RelationGetRelationName(rel))));
+	ReleaseSysCache(tuple);
+
+	/*
+	 * Drop the dependency records of the GENERATED expression, in particular
+	 * its INTERNAL dependency on the column, which would otherwise cause
+	 * dependency.c to refuse to perform the deletion.
+	 */
+	attrdefoid = GetAttrDefaultOid(RelationGetRelid(rel), attnum);
+	if (!OidIsValid(attrdefoid))
+		elog(ERROR, "could not find attrdef tuple for relation %u attnum %d",
+			 RelationGetRelid(rel), attnum);
+	(void) deleteDependencyRecordsFor(AttrDefaultRelationId, attrdefoid, false);
+
+	/* Make above changes visible */
+	CommandCounterIncrement();
+
+	/*
+	 * Get rid of the GENERATED expression itself.  We use RESTRICT here for
+	 * safety, but at present we do not expect anything to depend on the
+	 * default.
+	 */
+	RemoveAttrDefault(RelationGetRelid(rel), attnum, DROP_RESTRICT,
+					  false, false);
+
+	/* Prepare to store the new expression, in the catalogs */
+	rawEnt = (RawColumnDefault *) palloc(sizeof(RawColumnDefault));
+	rawEnt->attnum = attnum;
+	rawEnt->raw_default = newDefault;
+	rawEnt->missingMode = false;
+	rawEnt->generated = ATTRIBUTE_GENERATED_STORED;
+
+	/* Store the generated expression */
+	AddRelationNewConstraints(rel, list_make1(rawEnt), NIL,
+							  false, true, false, NULL);
+
+	/* Make above new expression visible */
+	CommandCounterIncrement();
+
+	/* Prepare for table rewrite */
+	defval = (Expr *) build_column_default(rel, attnum);
+
+	newval = (NewColumnValue *) palloc0(sizeof(NewColumnValue));
+	newval->attnum = attnum;
+	newval->expr = expression_planner(defval);
+	newval->is_generated = true;
+
+	tab->newvals = lappend(tab->newvals, newval);
+	tab->rewrite |= AT_REWRITE_DEFAULT_VAL;
+
+	ObjectAddressSubSet(address, RelationRelationId,
+						RelationGetRelid(rel), attnum);
+	return address;
+}
+
 /*
  * ALTER TABLE ALTER COLUMN DROP EXPRESSION
  */
diff --git a/src/backend/parser/gram.y b/src/backend/parser/gram.y
index c224df4eccc..b3c0e51230f 100644
--- a/src/backend/parser/gram.y
+++ b/src/backend/parser/gram.y
@@ -2404,6 +2404,16 @@ alter_table_cmd:
 					n->name = $3;
 					$$ = (Node *) n;
 				}
+			/* ALTER TABLE <name> ALTER [COLUMN] <colname> SET EXPRESSION AS <expr> */
+			| ALTER opt_column ColId SET EXPRESSION AS '(' a_expr ')'
+				{
+					AlterTableCmd *n = makeNode(AlterTableCmd);
+
+					n->subtype = AT_SetExpression;
+					n->name = $3;
+					n->def = $8;
+					$$ = (Node *) n;
+				}
 			/* ALTER TABLE <name> ALTER [COLUMN] <colname> DROP EXPRESSION */
 			| ALTER opt_column ColId DROP EXPRESSION
 				{
diff --git a/src/bin/psql/tab-complete.c b/src/bin/psql/tab-complete.c
index 006e10f5d2d..019d55a659d 100644
--- a/src/bin/psql/tab-complete.c
+++ b/src/bin/psql/tab-complete.c
@@ -2483,7 +2483,7 @@ psql_completion(const char *text, int start, int end)
 	/* ALTER TABLE ALTER [COLUMN] <foo> SET */
 	else if (Matches("ALTER", "TABLE", MatchAny, "ALTER", "COLUMN", MatchAny, "SET") ||
 			 Matches("ALTER", "TABLE", MatchAny, "ALTER", MatchAny, "SET"))
-		COMPLETE_WITH("(", "COMPRESSION", "DEFAULT", "GENERATED", "NOT NULL", "STATISTICS", "STORAGE",
+		COMPLETE_WITH("(", "COMPRESSION", "DEFAULT", "EXPRESSION", "GENERATED", "NOT NULL", "STATISTICS", "STORAGE",
 		/* a subset of ALTER SEQUENCE options */
 					  "INCREMENT", "MINVALUE", "MAXVALUE", "START", "NO", "CACHE", "CYCLE");
 	/* ALTER TABLE ALTER [COLUMN] <foo> SET ( */
@@ -2494,6 +2494,14 @@ psql_completion(const char *text, int start, int end)
 	else if (Matches("ALTER", "TABLE", MatchAny, "ALTER", "COLUMN", MatchAny, "SET", "COMPRESSION") ||
 			 Matches("ALTER", "TABLE", MatchAny, "ALTER", MatchAny, "SET", "COMPRESSION"))
 		COMPLETE_WITH("DEFAULT", "PGLZ", "LZ4");
+	/* ALTER TABLE ALTER [COLUMN] <foo> SET EXPRESSION */
+	else if (Matches("ALTER", "TABLE", MatchAny, "ALTER", "COLUMN", MatchAny, "SET", "EXPRESSION") ||
+			 Matches("ALTER", "TABLE", MatchAny, "ALTER", MatchAny, "SET", "EXPRESSION"))
+		COMPLETE_WITH("AS");
+	/* ALTER TABLE ALTER [COLUMN] <foo> SET EXPRESSION AS */
+	else if (Matches("ALTER", "TABLE", MatchAny, "ALTER", "COLUMN", MatchAny, "SET", "EXPRESSION", "AS") ||
+			 Matches("ALTER", "TABLE", MatchAny, "ALTER", MatchAny, "SET", "EXPRESSION", "AS"))
+		COMPLETE_WITH("(");
 	/* ALTER TABLE ALTER [COLUMN] <foo> SET GENERATED */
 	else if (Matches("ALTER", "TABLE", MatchAny, "ALTER", "COLUMN", MatchAny, "SET", "GENERATED") ||
 			 Matches("ALTER", "TABLE", MatchAny, "ALTER", MatchAny, "SET", "GENERATED"))
diff --git a/src/include/nodes/parsenodes.h b/src/include/nodes/parsenodes.h
index e494309da8d..cd08391dcb0 100644
--- a/src/include/nodes/parsenodes.h
+++ b/src/include/nodes/parsenodes.h
@@ -2192,6 +2192,7 @@ typedef enum AlterTableType
 	AT_DropNotNull,				/* alter column drop not null */
 	AT_SetNotNull,				/* alter column set not null */
 	AT_SetAttNotNull,			/* set attnotnull w/o a constraint */
+	AT_SetExpression,			/* alter column set expression */
 	AT_DropExpression,			/* alter column drop expression */
 	AT_SetStatistics,			/* alter column set statistics */
 	AT_SetOptions,				/* alter column set ( options ) */
diff --git a/src/test/modules/test_ddl_deparse/test_ddl_deparse.c b/src/test/modules/test_ddl_deparse/test_ddl_deparse.c
index 0302f79bb71..d34c7719dca 100644
--- a/src/test/modules/test_ddl_deparse/test_ddl_deparse.c
+++ b/src/test/modules/test_ddl_deparse/test_ddl_deparse.c
@@ -132,6 +132,9 @@ get_altertable_subcmdinfo(PG_FUNCTION_ARGS)
 			case AT_SetAttNotNull:
 				strtype = "SET ATTNOTNULL";
 				break;
+			case AT_SetExpression:
+				strtype = "SET EXPRESSION AS";
+				break;
 			case AT_DropExpression:
 				strtype = "DROP EXPRESSION";
 				break;
diff --git a/src/test/regress/expected/generated.out b/src/test/regress/expected/generated.out
index a2f38d0f50a..1cd7d6c6b35 100644
--- a/src/test/regress/expected/generated.out
+++ b/src/test/regress/expected/generated.out
@@ -785,30 +785,119 @@ Partition of: gtest_parent FOR VALUES FROM ('08-01-2016') TO ('09-01-2016')
 Partition of: gtest_parent FOR VALUES FROM ('09-01-2016') TO ('10-01-2016')
 
 INSERT INTO gtest_parent (f1, f2) VALUES ('2016-07-15', 1);
-SELECT * FROM gtest_parent;
-     f1     | f2 | f3 
-------------+----+----
- 07-15-2016 |  1 |  2
-(1 row)
+INSERT INTO gtest_parent (f1, f2) VALUES ('2016-07-15', 2);
+INSERT INTO gtest_parent (f1, f2) VALUES ('2016-08-15', 3);
+SELECT tableoid::regclass, * FROM gtest_parent ORDER BY 1, 2, 3;
+   tableoid   |     f1     | f2 | f3 
+--------------+------------+----+----
+ gtest_child  | 07-15-2016 |  1 |  2
+ gtest_child  | 07-15-2016 |  2 |  4
+ gtest_child2 | 08-15-2016 |  3 | 66
+(3 rows)
 
-SELECT * FROM gtest_child;
-     f1     | f2 | f3 
-------------+----+----
- 07-15-2016 |  1 |  2
-(1 row)
+UPDATE gtest_parent SET f1 = f1 + 60 WHERE f2 = 1;
+SELECT tableoid::regclass, * FROM gtest_parent ORDER BY 1, 2, 3;
+   tableoid   |     f1     | f2 | f3 
+--------------+------------+----+----
+ gtest_child  | 07-15-2016 |  2 |  4
+ gtest_child2 | 08-15-2016 |  3 | 66
+ gtest_child3 | 09-13-2016 |  1 | 33
+(3 rows)
 
-UPDATE gtest_parent SET f1 = f1 + 60;
-SELECT * FROM gtest_parent;
-     f1     | f2 | f3 
-------------+----+----
- 09-13-2016 |  1 | 33
-(1 row)
+-- alter only parent's and one child's generated expression
+ALTER TABLE ONLY gtest_parent ALTER COLUMN f3 SET EXPRESSION AS (f2 * 4);
+ALTER TABLE gtest_child ALTER COLUMN f3 SET EXPRESSION AS (f2 * 10);
+\d gtest_parent
+                   Partitioned table "public.gtest_parent"
+ Column |  Type  | Collation | Nullable |               Default               
+--------+--------+-----------+----------+-------------------------------------
+ f1     | date   |           | not null | 
+ f2     | bigint |           |          | 
+ f3     | bigint |           |          | generated always as (f2 * 4) stored
+Partition key: RANGE (f1)
+Number of partitions: 3 (Use \d+ to list them.)
 
-SELECT * FROM gtest_child3;
-     f1     | f2 | f3 
-------------+----+----
- 09-13-2016 |  1 | 33
-(1 row)
+\d gtest_child
+                          Table "public.gtest_child"
+ Column |  Type  | Collation | Nullable |               Default                
+--------+--------+-----------+----------+--------------------------------------
+ f1     | date   |           | not null | 
+ f2     | bigint |           |          | 
+ f3     | bigint |           |          | generated always as (f2 * 10) stored
+Partition of: gtest_parent FOR VALUES FROM ('07-01-2016') TO ('08-01-2016')
+
+\d gtest_child2
+                          Table "public.gtest_child2"
+ Column |  Type  | Collation | Nullable |               Default                
+--------+--------+-----------+----------+--------------------------------------
+ f1     | date   |           | not null | 
+ f2     | bigint |           |          | 
+ f3     | bigint |           |          | generated always as (f2 * 22) stored
+Partition of: gtest_parent FOR VALUES FROM ('08-01-2016') TO ('09-01-2016')
+
+\d gtest_child3
+                          Table "public.gtest_child3"
+ Column |  Type  | Collation | Nullable |               Default                
+--------+--------+-----------+----------+--------------------------------------
+ f1     | date   |           | not null | 
+ f2     | bigint |           |          | 
+ f3     | bigint |           |          | generated always as (f2 * 33) stored
+Partition of: gtest_parent FOR VALUES FROM ('09-01-2016') TO ('10-01-2016')
+
+SELECT tableoid::regclass, * FROM gtest_parent ORDER BY 1, 2, 3;
+   tableoid   |     f1     | f2 | f3 
+--------------+------------+----+----
+ gtest_child  | 07-15-2016 |  2 | 20
+ gtest_child2 | 08-15-2016 |  3 | 66
+ gtest_child3 | 09-13-2016 |  1 | 33
+(3 rows)
+
+-- alter generated expression of a parent and all it's child altogether
+ALTER TABLE gtest_parent ALTER COLUMN f3 SET EXPRESSION AS (f2 * 2);
+\d gtest_parent
+                   Partitioned table "public.gtest_parent"
+ Column |  Type  | Collation | Nullable |               Default               
+--------+--------+-----------+----------+-------------------------------------
+ f1     | date   |           | not null | 
+ f2     | bigint |           |          | 
+ f3     | bigint |           |          | generated always as (f2 * 2) stored
+Partition key: RANGE (f1)
+Number of partitions: 3 (Use \d+ to list them.)
+
+\d gtest_child
+                          Table "public.gtest_child"
+ Column |  Type  | Collation | Nullable |               Default               
+--------+--------+-----------+----------+-------------------------------------
+ f1     | date   |           | not null | 
+ f2     | bigint |           |          | 
+ f3     | bigint |           |          | generated always as (f2 * 2) stored
+Partition of: gtest_parent FOR VALUES FROM ('07-01-2016') TO ('08-01-2016')
+
+\d gtest_child2
+                         Table "public.gtest_child2"
+ Column |  Type  | Collation | Nullable |               Default               
+--------+--------+-----------+----------+-------------------------------------
+ f1     | date   |           | not null | 
+ f2     | bigint |           |          | 
+ f3     | bigint |           |          | generated always as (f2 * 2) stored
+Partition of: gtest_parent FOR VALUES FROM ('08-01-2016') TO ('09-01-2016')
+
+\d gtest_child3
+                         Table "public.gtest_child3"
+ Column |  Type  | Collation | Nullable |               Default               
+--------+--------+-----------+----------+-------------------------------------
+ f1     | date   |           | not null | 
+ f2     | bigint |           |          | 
+ f3     | bigint |           |          | generated always as (f2 * 2) stored
+Partition of: gtest_parent FOR VALUES FROM ('09-01-2016') TO ('10-01-2016')
+
+SELECT tableoid::regclass, * FROM gtest_parent ORDER BY 1, 2, 3;
+   tableoid   |     f1     | f2 | f3 
+--------------+------------+----+----
+ gtest_child  | 07-15-2016 |  2 |  4
+ gtest_child2 | 08-15-2016 |  3 |  6
+ gtest_child3 | 09-13-2016 |  1 |  2
+(3 rows)
 
 -- we leave these tables around for purposes of testing dump/reload/upgrade
 -- generated columns in partition key (not allowed)
@@ -935,18 +1024,50 @@ CREATE TABLE gtest29 (
     b int GENERATED ALWAYS AS (a * 2) STORED
 );
 INSERT INTO gtest29 (a) VALUES (3), (4);
+SELECT * FROM gtest29;
+ a | b 
+---+---
+ 3 | 6
+ 4 | 8
+(2 rows)
+
+\d gtest29
+                            Table "public.gtest29"
+ Column |  Type   | Collation | Nullable |              Default               
+--------+---------+-----------+----------+------------------------------------
+ a      | integer |           |          | 
+ b      | integer |           |          | generated always as (a * 2) stored
+
+ALTER TABLE gtest29 ALTER COLUMN a SET EXPRESSION AS (a * 3);  -- error
+ERROR:  column "a" of relation "gtest29" is not a stored generated column
 ALTER TABLE gtest29 ALTER COLUMN a DROP EXPRESSION;  -- error
 ERROR:  column "a" of relation "gtest29" is not a stored generated column
 ALTER TABLE gtest29 ALTER COLUMN a DROP EXPRESSION IF EXISTS;  -- notice
 NOTICE:  column "a" of relation "gtest29" is not a stored generated column, skipping
+-- Change the expression
+ALTER TABLE gtest29 ALTER COLUMN b SET EXPRESSION AS (a * 3);
+SELECT * FROM gtest29;
+ a | b  
+---+----
+ 3 |  9
+ 4 | 12
+(2 rows)
+
+\d gtest29
+                            Table "public.gtest29"
+ Column |  Type   | Collation | Nullable |              Default               
+--------+---------+-----------+----------+------------------------------------
+ a      | integer |           |          | 
+ b      | integer |           |          | generated always as (a * 3) stored
+
 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;
  a | b  
 ---+----
- 3 |  6
- 4 |  8
+ 3 |  9
+ 4 | 12
  5 |   
  6 | 66
 (4 rows)
diff --git a/src/test/regress/sql/generated.sql b/src/test/regress/sql/generated.sql
index 298f6b3aa8b..a0a1942056b 100644
--- a/src/test/regress/sql/generated.sql
+++ b/src/test/regress/sql/generated.sql
@@ -414,11 +414,28 @@ ALTER TABLE gtest_parent ATTACH PARTITION gtest_child3 FOR VALUES FROM ('2016-09
 \d gtest_child2
 \d gtest_child3
 INSERT INTO gtest_parent (f1, f2) VALUES ('2016-07-15', 1);
-SELECT * FROM gtest_parent;
-SELECT * FROM gtest_child;
-UPDATE gtest_parent SET f1 = f1 + 60;
-SELECT * FROM gtest_parent;
-SELECT * FROM gtest_child3;
+INSERT INTO gtest_parent (f1, f2) VALUES ('2016-07-15', 2);
+INSERT INTO gtest_parent (f1, f2) VALUES ('2016-08-15', 3);
+SELECT tableoid::regclass, * FROM gtest_parent ORDER BY 1, 2, 3;
+UPDATE gtest_parent SET f1 = f1 + 60 WHERE f2 = 1;
+SELECT tableoid::regclass, * FROM gtest_parent ORDER BY 1, 2, 3;
+
+-- alter only parent's and one child's generated expression
+ALTER TABLE ONLY gtest_parent ALTER COLUMN f3 SET EXPRESSION AS (f2 * 4);
+ALTER TABLE gtest_child ALTER COLUMN f3 SET EXPRESSION AS (f2 * 10);
+\d gtest_parent
+\d gtest_child
+\d gtest_child2
+\d gtest_child3
+SELECT tableoid::regclass, * FROM gtest_parent ORDER BY 1, 2, 3;
+
+-- alter generated expression of a parent and all it's child altogether
+ALTER TABLE gtest_parent ALTER COLUMN f3 SET EXPRESSION AS (f2 * 2);
+\d gtest_parent
+\d gtest_child
+\d gtest_child2
+\d gtest_child3
+SELECT tableoid::regclass, * FROM gtest_parent ORDER BY 1, 2, 3;
 -- we leave these tables around for purposes of testing dump/reload/upgrade
 
 -- generated columns in partition key (not allowed)
@@ -473,8 +490,17 @@ CREATE TABLE gtest29 (
     b int GENERATED ALWAYS AS (a * 2) STORED
 );
 INSERT INTO gtest29 (a) VALUES (3), (4);
+SELECT * FROM gtest29;
+\d gtest29
+ALTER TABLE gtest29 ALTER COLUMN a SET EXPRESSION AS (a * 3);  -- error
 ALTER TABLE gtest29 ALTER COLUMN a DROP EXPRESSION;  -- error
 ALTER TABLE gtest29 ALTER COLUMN a DROP EXPRESSION IF EXISTS;  -- notice
+
+-- Change the expression
+ALTER TABLE gtest29 ALTER COLUMN b SET EXPRESSION AS (a * 3);
+SELECT * FROM gtest29;
+\d gtest29
+
 ALTER TABLE gtest29 ALTER COLUMN b DROP EXPRESSION;
 INSERT INTO gtest29 (a) VALUES (5);
 INSERT INTO gtest29 (a, b) VALUES (6, 66);
-- 
2.18.0

