bug: ALTER TABLE ADD VIRTUAL GENERATED COLUMN with table rewrite

From f247e3b8005e17b6446b243f7d03f7d02a5a82d7 Mon Sep 17 00:00:00 2001
From: jian he <jian.universality@gmail.com>
Date: Thu, 27 Feb 2025 22:54:03 +0800
Subject: [PATCH v1 1/1] fix ALTER TABLE ADD VIRTUAL GENERATED COLUMN when
 table rewrite

demo:
CREATE TABLE gtest20a (a int PRIMARY KEY, b int GENERATED ALWAYS AS (a * 2) VIRTUAL);
ALTER TABLE gtest20a ADD COLUMN c float8 DEFAULT RANDOM() CHECK (b < 60);
ERROR:  no generation expression found for column number 2 of table "pg_temp_17306"

in ATRewriteTable, the variable OIDNewHeap(if valid) corresponding pg_attrdef
default expression entry was not populated.  so OIDNewHeap cannot be used to
call expand_generated_columns_in_expr or build_generation_expression.  Therefore
in ATRewriteTable, we can only use the existing relation to expand the generated
expression.

discussion: https://postgr.es/m/
---
 src/backend/commands/tablecmds.c                | 2 +-
 src/test/regress/expected/generated_virtual.out | 4 ++++
 src/test/regress/sql/generated_virtual.sql      | 3 +++
 3 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/src/backend/commands/tablecmds.c b/src/backend/commands/tablecmds.c
index ce7d115667e..7d870c186ca 100644
--- a/src/backend/commands/tablecmds.c
+++ b/src/backend/commands/tablecmds.c
@@ -6149,7 +6149,7 @@ ATRewriteTable(AlteredTableInfo *tab, Oid OIDNewHeap)
 		{
 			case CONSTR_CHECK:
 				needscan = true;
-				con->qualstate = ExecPrepareExpr((Expr *) expand_generated_columns_in_expr(con->qual, newrel ? newrel : oldrel, 1), estate);
+				con->qualstate = ExecPrepareExpr((Expr *) expand_generated_columns_in_expr(con->qual, oldrel, 1), estate);
 				break;
 			case CONSTR_FOREIGN:
 				/* Nothing to do here */
diff --git a/src/test/regress/expected/generated_virtual.out b/src/test/regress/expected/generated_virtual.out
index b339fbcebfa..107562f10ff 100644
--- a/src/test/regress/expected/generated_virtual.out
+++ b/src/test/regress/expected/generated_virtual.out
@@ -647,6 +647,10 @@ INSERT INTO gtest20a (a) VALUES (10);
 INSERT INTO gtest20a (a) VALUES (30);
 ALTER TABLE gtest20a ADD CHECK (b < 50);  -- fails on existing row
 ERROR:  check constraint "gtest20a_b_check" of relation "gtest20a" is violated by some row
+--table rewrite cases.
+ALTER TABLE gtest20a ADD COLUMN c float8 DEFAULT RANDOM() CHECK (b < 50); -- fails on existing row
+ERROR:  check constraint "gtest20a_b_check" of relation "gtest20a" is violated by some row
+ALTER TABLE gtest20a ADD COLUMN c float8 DEFAULT RANDOM() CHECK (b < 61); --ok.
 CREATE TABLE gtest20b (a int PRIMARY KEY, b int GENERATED ALWAYS AS (a * 2) VIRTUAL);
 INSERT INTO gtest20b (a) VALUES (10);
 INSERT INTO gtest20b (a) VALUES (30);
diff --git a/src/test/regress/sql/generated_virtual.sql b/src/test/regress/sql/generated_virtual.sql
index c80630c11a5..09fb1c477e9 100644
--- a/src/test/regress/sql/generated_virtual.sql
+++ b/src/test/regress/sql/generated_virtual.sql
@@ -319,6 +319,9 @@ CREATE TABLE gtest20a (a int PRIMARY KEY, b int GENERATED ALWAYS AS (a * 2) VIRT
 INSERT INTO gtest20a (a) VALUES (10);
 INSERT INTO gtest20a (a) VALUES (30);
 ALTER TABLE gtest20a ADD CHECK (b < 50);  -- fails on existing row
+--table rewrite cases.
+ALTER TABLE gtest20a ADD COLUMN c float8 DEFAULT RANDOM() CHECK (b < 50); -- fails on existing row
+ALTER TABLE gtest20a ADD COLUMN c float8 DEFAULT RANDOM() CHECK (b < 61); --ok.
 
 CREATE TABLE gtest20b (a int PRIMARY KEY, b int GENERATED ALWAYS AS (a * 2) VIRTUAL);
 INSERT INTO gtest20b (a) VALUES (10);
-- 
2.34.1