Skip copying GENERATED ALWAYS AS STORED cols in ReplaceTable (#4616)

Postgres doesn't allow inserting into columns having GENERATED ALWAYS
AS (...) STORED expressions.
For this reason, when executing undistribute_table or an alter_* udf,
we should skip copying such columns.
This is not bad since Postgres would already generate such columns.

src/backend/distributed/commands/alter_table.c

@@ -185,6 +185,8 @@ static void CreateCitusTableLike(TableConversionState *con);
 static List * GetViewCreationCommandsOfTable(Oid relationId);
 static void ReplaceTable(Oid sourceId, Oid targetId, List *justBeforeDropCommands,
 						 bool suppressNoticeMessages);
+static bool HasAnyGeneratedStoredColumns(Oid relationId);
+static List * GetNonGeneratedStoredColumnNameList(Oid relationId);
 static void CheckAlterDistributedTableConversionParameters(TableConversionState *con);
 static char * CreateWorkerChangeSequenceDependencyCommand(char *sequenceSchemaName,
 														  char *sequenceName,
@@ -1122,9 +1124,33 @@ ReplaceTable(Oid sourceId, Oid targetId, List *justBeforeDropCommands,
 									quote_qualified_identifier(schemaName, sourceName))));
 		}
 
+		if (!HasAnyGeneratedStoredColumns(sourceId))
+		{
+			/*
+			 * Relation has no GENERATED STORED columns, copy the table via plain
+			 * "INSERT INTO .. SELECT *"".
+			 */
 			appendStringInfo(query, "INSERT INTO %s SELECT * FROM %s",
 							 quote_qualified_identifier(schemaName, targetName),
 							 quote_qualified_identifier(schemaName, sourceName));
+		}
+		else
+		{
+			/*
+			 * Skip columns having GENERATED ALWAYS AS (...) STORED expressions
+			 * since Postgres doesn't allow inserting into such columns.
+			 * This is not bad since Postgres would already generate such columns.
+			 * Note that here we intentionally don't skip columns having DEFAULT
+			 * expressions since user might have inserted non-default values.
+			 */
+			List *nonStoredColumnNameList = GetNonGeneratedStoredColumnNameList(sourceId);
+			char *insertColumnString = StringJoin(nonStoredColumnNameList, ',');
+			appendStringInfo(query, "INSERT INTO %s (%s) SELECT %s FROM %s",
+							 quote_qualified_identifier(schemaName, targetName),
+							 insertColumnString, insertColumnString,
+							 quote_qualified_identifier(schemaName, sourceName));
+		}
+
 		ExecuteQueryViaSPI(query->data, SPI_OK_INSERT);
 	}
 
@@ -1183,6 +1209,55 @@ ReplaceTable(Oid sourceId, Oid targetId, List *justBeforeDropCommands,
 }
 
 
+/*
+ * HasAnyGeneratedStoredColumns decides if relation has any columns that we
+ * might need to copy the data of when replacing table.
+ */
+static bool
+HasAnyGeneratedStoredColumns(Oid relationId)
+{
+	return list_length(GetNonGeneratedStoredColumnNameList(relationId)) > 0;
+}
+
+
+/*
+ * GetNonGeneratedStoredColumnNameList returns a list of column names for
+ * columns not having GENERATED ALWAYS AS (...) STORED expressions.
+ */
+static List *
+GetNonGeneratedStoredColumnNameList(Oid relationId)
+{
+	List *nonStoredColumnNameList = NIL;
+
+	Relation relation = relation_open(relationId, AccessShareLock);
+	TupleDesc tupleDescriptor = RelationGetDescr(relation);
+	for (int columnIndex = 0; columnIndex < tupleDescriptor->natts; columnIndex++)
+	{
+		Form_pg_attribute currentColumn = TupleDescAttr(tupleDescriptor, columnIndex);
+		if (currentColumn->attisdropped)
+		{
+			/* skip dropped columns */
+			continue;
+		}
+
+#if PG_VERSION_NUM >= 120000
+		if (currentColumn->attgenerated == ATTRIBUTE_GENERATED_STORED)
+		{
+			continue;
+		}
+#endif
+
+		const char *quotedColumnName = quote_identifier(NameStr(currentColumn->attname));
+		nonStoredColumnNameList = lappend(nonStoredColumnNameList,
+										  pstrdup(quotedColumnName));
+	}
+
+	relation_close(relation, NoLock);
+
+	return nonStoredColumnNameList;
+}
+
+
 /*
  * CheckAlterDistributedTableConversionParameters errors for the cases where
  * alter_distributed_table UDF wouldn't work.

src/test/regress/expected/pg12.out

@@ -449,6 +449,199 @@ BEGIN;
  generated_stored_col_test_60040 | y       | s
 (2 rows)
 
+ROLLBACK;
+CREATE TABLE generated_stored_dist (
+  col_1 int,
+  "col\'_2" text,
+  col_3 text generated always as (UPPER("col\'_2")) stored
+);
+SELECT create_distributed_table ('generated_stored_dist', 'col_1');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+INSERT INTO generated_stored_dist VALUES (1, 'text_1'), (2, 'text_2');
+SELECT * FROM generated_stored_dist ORDER BY 1,2,3;
+ col_1 | col\'_2 | col_3
+---------------------------------------------------------------------
+     1 | text_1  | TEXT_1
+     2 | text_2  | TEXT_2
+(2 rows)
+
+INSERT INTO generated_stored_dist VALUES (1, 'text_1'), (2, 'text_2');
+SELECT alter_distributed_table('generated_stored_dist', shard_count := 5, cascade_to_colocated := false);
+NOTICE:  creating a new table for test_pg12.generated_stored_dist
+NOTICE:  Moving the data of test_pg12.generated_stored_dist
+NOTICE:  Dropping the old test_pg12.generated_stored_dist
+NOTICE:  Renaming the new table to test_pg12.generated_stored_dist
+ alter_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT * FROM generated_stored_dist ORDER BY 1,2,3;
+ col_1 | col\'_2 | col_3
+---------------------------------------------------------------------
+     1 | text_1  | TEXT_1
+     1 | text_1  | TEXT_1
+     2 | text_2  | TEXT_2
+     2 | text_2  | TEXT_2
+(4 rows)
+
+CREATE TABLE generated_stored_local (
+  col_1 int,
+  "col\'_2" text,
+  col_3 text generated always as (UPPER("col\'_2")) stored
+);
+SELECT citus_add_local_table_to_metadata('generated_stored_local');
+ citus_add_local_table_to_metadata
+---------------------------------------------------------------------
+
+(1 row)
+
+INSERT INTO generated_stored_local VALUES (1, 'text_1'), (2, 'text_2');
+SELECT * FROM generated_stored_local ORDER BY 1,2,3;
+ col_1 | col\'_2 | col_3
+---------------------------------------------------------------------
+     1 | text_1  | TEXT_1
+     2 | text_2  | TEXT_2
+(2 rows)
+
+SELECT create_distributed_table ('generated_stored_local', 'col_1');
+NOTICE:  Copying data from local table...
+NOTICE:  copying the data has completed
+DETAIL:  The local data in the table is no longer visible, but is still on disk.
+HINT:  To remove the local data, run: SELECT truncate_local_data_after_distributing_table($$test_pg12.generated_stored_local$$)
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+INSERT INTO generated_stored_local VALUES (1, 'text_1'), (2, 'text_2');
+SELECT * FROM generated_stored_local ORDER BY 1,2,3;
+ col_1 | col\'_2 | col_3
+---------------------------------------------------------------------
+     1 | text_1  | TEXT_1
+     1 | text_1  | TEXT_1
+     2 | text_2  | TEXT_2
+     2 | text_2  | TEXT_2
+(4 rows)
+
+create table generated_stored_columnar(i int) partition by range(i);
+create table generated_stored_columnar_p0 partition of generated_stored_columnar for values from (0) to (10);
+create table generated_stored_columnar_p1 partition of generated_stored_columnar for values from (10) to (20);
+SELECT alter_table_set_access_method('generated_stored_columnar_p0', 'columnar');
+NOTICE:  creating a new table for test_pg12.generated_stored_columnar_p0
+NOTICE:  Moving the data of test_pg12.generated_stored_columnar_p0
+NOTICE:  Dropping the old test_pg12.generated_stored_columnar_p0
+NOTICE:  Renaming the new table to test_pg12.generated_stored_columnar_p0
+ alter_table_set_access_method
+---------------------------------------------------------------------
+
+(1 row)
+
+CREATE TABLE generated_stored_ref (
+  col_1 int,
+  col_2 int,
+  col_3 int generated always as (col_1+col_2) stored,
+  col_4 int,
+  col_5 int generated always as (col_4*2-col_1) stored
+);
+SELECT create_reference_table ('generated_stored_ref');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+INSERT INTO generated_stored_ref (col_1, col_4) VALUES (1,2), (11,12);
+INSERT INTO generated_stored_ref (col_1, col_2, col_4) VALUES (100,101,102), (200,201,202);
+SELECT * FROM generated_stored_ref ORDER BY 1,2,3,4,5;
+ col_1 | col_2 | col_3 | col_4 | col_5
+---------------------------------------------------------------------
+     1 |       |       |     2 |     3
+    11 |       |       |    12 |    13
+   100 |   101 |   201 |   102 |   104
+   200 |   201 |   401 |   202 |   204
+(4 rows)
+
+BEGIN;
+  SELECT undistribute_table('generated_stored_ref');
+NOTICE:  creating a new table for test_pg12.generated_stored_ref
+NOTICE:  Moving the data of test_pg12.generated_stored_ref
+NOTICE:  Dropping the old test_pg12.generated_stored_ref
+NOTICE:  Renaming the new table to test_pg12.generated_stored_ref
+ undistribute_table
+---------------------------------------------------------------------
+
+(1 row)
+
+  INSERT INTO generated_stored_ref (col_1, col_4) VALUES (11,12), (21,22);
+  INSERT INTO generated_stored_ref (col_1, col_2, col_4) VALUES (200,201,202), (300,301,302);
+  SELECT * FROM generated_stored_ref ORDER BY 1,2,3,4,5;
+ col_1 | col_2 | col_3 | col_4 | col_5
+---------------------------------------------------------------------
+     1 |       |       |     2 |     3
+    11 |       |       |    12 |    13
+    11 |       |       |    12 |    13
+    21 |       |       |    22 |    23
+   100 |   101 |   201 |   102 |   104
+   200 |   201 |   401 |   202 |   204
+   200 |   201 |   401 |   202 |   204
+   300 |   301 |   601 |   302 |   304
+(8 rows)
+
+ROLLBACK;
+BEGIN;
+  -- drop some of the columns not having "generated always as stored" expressions
+  -- this would drop generated columns too
+  ALTER TABLE generated_stored_ref DROP COLUMN col_1;
+  ALTER TABLE generated_stored_ref DROP COLUMN col_4;
+  -- show that undistribute_table works fine
+  SELECT undistribute_table('generated_stored_ref');
+NOTICE:  creating a new table for test_pg12.generated_stored_ref
+NOTICE:  Moving the data of test_pg12.generated_stored_ref
+NOTICE:  Dropping the old test_pg12.generated_stored_ref
+NOTICE:  Renaming the new table to test_pg12.generated_stored_ref
+ undistribute_table
+---------------------------------------------------------------------
+
+(1 row)
+
+  INSERT INTO generated_stored_ref VALUES (5);
+  SELECT * FROM generated_stored_REF ORDER BY 1;
+ col_2
+---------------------------------------------------------------------
+     5
+   101
+   201
+
+
+(5 rows)
+
+ROLLBACK;
+BEGIN;
+  -- now drop all columns
+  ALTER TABLE generated_stored_ref DROP COLUMN col_3;
+  ALTER TABLE generated_stored_ref DROP COLUMN col_5;
+  ALTER TABLE generated_stored_ref DROP COLUMN col_1;
+  ALTER TABLE generated_stored_ref DROP COLUMN col_2;
+  ALTER TABLE generated_stored_ref DROP COLUMN col_4;
+  -- show that undistribute_table works fine
+  SELECT undistribute_table('generated_stored_ref');
+NOTICE:  creating a new table for test_pg12.generated_stored_ref
+NOTICE:  Moving the data of test_pg12.generated_stored_ref
+NOTICE:  Dropping the old test_pg12.generated_stored_ref
+NOTICE:  Renaming the new table to test_pg12.generated_stored_ref
+ undistribute_table
+---------------------------------------------------------------------
+
+(1 row)
+
+  SELECT * FROM generated_stored_ref;
+--
+(4 rows)
+
 ROLLBACK;
 RESET citus.replicate_reference_tables_on_activate;
 SELECT citus_remove_node('localhost', :master_port);
@@ -459,6 +652,6 @@ SELECT citus_remove_node('localhost', :master_port);
 
 \set VERBOSITY terse
 drop schema test_pg12 cascade;
-NOTICE:  drop cascades to 10 other objects
+NOTICE:  drop cascades to 15 other objects
 \set VERBOSITY default
 SET citus.shard_replication_factor to 2;

src/test/regress/sql/pg12.sql

@@ -296,6 +296,90 @@ BEGIN;
   ORDER BY 1,2;
 ROLLBACK;
 
+CREATE TABLE generated_stored_dist (
+  col_1 int,
+  "col\'_2" text,
+  col_3 text generated always as (UPPER("col\'_2")) stored
+);
+
+SELECT create_distributed_table ('generated_stored_dist', 'col_1');
+
+INSERT INTO generated_stored_dist VALUES (1, 'text_1'), (2, 'text_2');
+SELECT * FROM generated_stored_dist ORDER BY 1,2,3;
+
+INSERT INTO generated_stored_dist VALUES (1, 'text_1'), (2, 'text_2');
+SELECT alter_distributed_table('generated_stored_dist', shard_count := 5, cascade_to_colocated := false);
+SELECT * FROM generated_stored_dist ORDER BY 1,2,3;
+
+CREATE TABLE generated_stored_local (
+  col_1 int,
+  "col\'_2" text,
+  col_3 text generated always as (UPPER("col\'_2")) stored
+);
+
+SELECT citus_add_local_table_to_metadata('generated_stored_local');
+
+INSERT INTO generated_stored_local VALUES (1, 'text_1'), (2, 'text_2');
+SELECT * FROM generated_stored_local ORDER BY 1,2,3;
+
+SELECT create_distributed_table ('generated_stored_local', 'col_1');
+
+INSERT INTO generated_stored_local VALUES (1, 'text_1'), (2, 'text_2');
+SELECT * FROM generated_stored_local ORDER BY 1,2,3;
+
+create table generated_stored_columnar(i int) partition by range(i);
+create table generated_stored_columnar_p0 partition of generated_stored_columnar for values from (0) to (10);
+create table generated_stored_columnar_p1 partition of generated_stored_columnar for values from (10) to (20);
+SELECT alter_table_set_access_method('generated_stored_columnar_p0', 'columnar');
+
+CREATE TABLE generated_stored_ref (
+  col_1 int,
+  col_2 int,
+  col_3 int generated always as (col_1+col_2) stored,
+  col_4 int,
+  col_5 int generated always as (col_4*2-col_1) stored
+);
+
+SELECT create_reference_table ('generated_stored_ref');
+
+INSERT INTO generated_stored_ref (col_1, col_4) VALUES (1,2), (11,12);
+INSERT INTO generated_stored_ref (col_1, col_2, col_4) VALUES (100,101,102), (200,201,202);
+
+SELECT * FROM generated_stored_ref ORDER BY 1,2,3,4,5;
+
+BEGIN;
+  SELECT undistribute_table('generated_stored_ref');
+  INSERT INTO generated_stored_ref (col_1, col_4) VALUES (11,12), (21,22);
+  INSERT INTO generated_stored_ref (col_1, col_2, col_4) VALUES (200,201,202), (300,301,302);
+  SELECT * FROM generated_stored_ref ORDER BY 1,2,3,4,5;
+ROLLBACK;
+
+BEGIN;
+  -- drop some of the columns not having "generated always as stored" expressions
+  -- this would drop generated columns too
+  ALTER TABLE generated_stored_ref DROP COLUMN col_1;
+  ALTER TABLE generated_stored_ref DROP COLUMN col_4;
+
+  -- show that undistribute_table works fine
+  SELECT undistribute_table('generated_stored_ref');
+  INSERT INTO generated_stored_ref VALUES (5);
+  SELECT * FROM generated_stored_REF ORDER BY 1;
+ROLLBACK;
+
+BEGIN;
+  -- now drop all columns
+  ALTER TABLE generated_stored_ref DROP COLUMN col_3;
+  ALTER TABLE generated_stored_ref DROP COLUMN col_5;
+  ALTER TABLE generated_stored_ref DROP COLUMN col_1;
+  ALTER TABLE generated_stored_ref DROP COLUMN col_2;
+  ALTER TABLE generated_stored_ref DROP COLUMN col_4;
+
+  -- show that undistribute_table works fine
+  SELECT undistribute_table('generated_stored_ref');
+
+  SELECT * FROM generated_stored_ref;
+ROLLBACK;
+
 RESET citus.replicate_reference_tables_on_activate;
 SELECT citus_remove_node('localhost', :master_port);