Merge pull request #4613 from citusdata/fix/generated-cols-citus-local

When adding local table to metadata, we are dropping DEFAULT expressions
from shard relation. When finding columns having DEFAULT expressions,
we shouldn't rely on atthasdef since it might be true if column has
GENERATED ALWAYS AS (...) STORED expression.

On the other hand, we should not actually drop such GENERATED expressions from
shard relation since we don't evaluate such columns in coordinator and this would
result in inserting NULL values to such columns.

src/backend/distributed/commands/citus_add_local_table_to_metadata.c

@@ -68,8 +68,8 @@ static char * GetRenameShardTriggerCommand(Oid shardRelationId, char *triggerNam
 static void DropRelationTruncateTriggers(Oid relationId);
 static char * GetDropTriggerCommand(Oid relationId, char *triggerName);
 static List * GetRenameStatsCommandList(List *statsOidList, uint64 shardId);
-static void DropAndMoveDefaultSequenceOwnerships(Oid sourceRelationId,
-												 Oid targetRelationId);
+static void DropDefaultExpressionsAndMoveOwnedSequenceOwnerships(Oid sourceRelationId,
+																 Oid targetRelationId);
 static void DropDefaultColumnDefinition(Oid relationId, char *columnName);
 static void TransferSequenceOwnership(Oid ownedSequenceId, Oid targetRelationId,
 									  char *columnName);
@@ -309,7 +309,8 @@ CreateCitusLocalTable(Oid relationId, bool cascadeViaForeignKeys)
 	 * DEFAULT expressions from shard relation as we should evaluate such columns
 	 * in shell table when needed.
 	 */
-	DropAndMoveDefaultSequenceOwnerships(shardRelationId, shellRelationId);
+	DropDefaultExpressionsAndMoveOwnedSequenceOwnerships(shardRelationId,
+														 shellRelationId);
 
 	InsertMetadataForCitusLocalTable(shellRelationId, shardId);
 
@@ -883,18 +884,19 @@ GetRenameStatsCommandList(List *statsOidList, uint64 shardId)
 
 
 /*
- * DropAndMoveDefaultSequenceOwnerships drops default column definitions for
- * relation with sourceRelationId. Also, for each column that defaults to an
- * owned sequence, it grants ownership to the same named column of the relation
- * with targetRelationId.
+ * DropDefaultExpressionsAndMoveOwnedSequenceOwnerships drops default column
+ * definitions for relation with sourceRelationId. Also, for each column that
+ * defaults to an owned sequence, it grants ownership to the same named column
+ * of the relation with targetRelationId.
  */
 static void
-DropAndMoveDefaultSequenceOwnerships(Oid sourceRelationId, Oid targetRelationId)
+DropDefaultExpressionsAndMoveOwnedSequenceOwnerships(Oid sourceRelationId,
+													 Oid targetRelationId)
 {
 	List *columnNameList = NIL;
 	List *ownedSequenceIdList = NIL;
-	ExtractColumnsOwningSequences(sourceRelationId, &columnNameList,
-								  &ownedSequenceIdList);
+	ExtractDefaultColumnsAndOwnedSequences(sourceRelationId, &columnNameList,
+										   &ownedSequenceIdList);
 
 	ListCell *columnNameCell = NULL;
 	ListCell *ownedSequenceIdCell = NULL;

src/backend/distributed/commands/sequence.c

@@ -151,13 +151,14 @@ OptionsSpecifyOwnedBy(List *optionList, Oid *ownedByTableId)
 
 
 /*
- * ExtractColumnsOwningSequences finds each column of relation with relationId
- * defaulting to an owned sequence. Then, appends the column name and id of the
- * owned sequence -that the column defaults- to the lists passed as NIL initially.
+ * ExtractDefaultColumnsAndOwnedSequences finds each column of relation with
+ * relationId that has a DEFAULT expression and each sequence owned by such
+ * columns (if any). Then, appends the column name and id of the owned sequence
+ * -that the column defaults- to the lists passed as NIL initially.
  */
 void
-ExtractColumnsOwningSequences(Oid relationId, List **columnNameList,
-							  List **ownedSequenceIdList)
+ExtractDefaultColumnsAndOwnedSequences(Oid relationId, List **columnNameList,
+									   List **ownedSequenceIdList)
 {
 	Assert(*columnNameList == NIL && *ownedSequenceIdList == NIL);
 
@@ -177,6 +178,14 @@ ExtractColumnsOwningSequences(Oid relationId, List **columnNameList,
 			continue;
 		}
 
+#if PG_VERSION_NUM >= PG_VERSION_12
+		if (attributeForm->attgenerated == ATTRIBUTE_GENERATED_STORED)
+		{
+			/* skip columns with GENERATED AS ALWAYS expressions */
+			continue;
+		}
+#endif
+
 		char *columnName = NameStr(attributeForm->attname);
 		*columnNameList = lappend(*columnNameList, columnName);
 

src/backend/distributed/metadata/metadata_sync.c

@@ -1131,7 +1131,7 @@ SequenceDependencyCommandList(Oid relationId)
 	List *columnNameList = NIL;
 	List *sequenceIdList = NIL;
 
-	ExtractColumnsOwningSequences(relationId, &columnNameList, &sequenceIdList);
+	ExtractDefaultColumnsAndOwnedSequences(relationId, &columnNameList, &sequenceIdList);
 
 	ListCell *columnNameCell = NULL;
 	ListCell *sequenceIdCell = NULL;
@@ -1144,7 +1144,7 @@ SequenceDependencyCommandList(Oid relationId)
 		if (!OidIsValid(sequenceId))
 		{
 			/*
-			 * ExtractColumnsOwningSequences returns entries for all columns,
+			 * ExtractDefaultColumnsAndOwnedSequences returns entries for all columns,
 			 * but with 0 sequence ID unless there is default nextval(..).
 			 */
 			continue;

src/include/distributed/commands/sequence.h

@@ -13,8 +13,9 @@
 #include "nodes/pg_list.h"
 
 
-extern void ExtractColumnsOwningSequences(Oid relationId, List **columnNameList,
-										  List **ownedSequenceIdList);
+extern void ExtractDefaultColumnsAndOwnedSequences(Oid relationId,
+												   List **columnNameList,
+												   List **ownedSequenceIdList);
 
 
 #endif /* CITUS_SEQUENCE_H */

src/test/regress/expected/pg12.out

@@ -411,6 +411,52 @@ where val = 'asdf';
      3
 (1 row)
 
+-- not replicate reference tables from other test files
+SET citus.replicate_reference_tables_on_activate TO off;
+SELECT 1 FROM citus_add_node('localhost', :master_port, groupId => 0);
+ ?column?
+---------------------------------------------------------------------
+        1
+(1 row)
+
+BEGIN;
+  CREATE TABLE generated_stored_col_test (x int, y int generated always as (x+1) stored);
+  SELECT citus_add_local_table_to_metadata('generated_stored_col_test');
+ citus_add_local_table_to_metadata
+---------------------------------------------------------------------
+
+(1 row)
+
+  -- simply check if GENERATED ALWAYS AS (...) STORED expression works fine
+  INSERT INTO generated_stored_col_test VALUES(1), (2);
+  SELECT * FROM generated_stored_col_test ORDER BY 1,2;
+ x | y
+---------------------------------------------------------------------
+ 1 | 2
+ 2 | 3
+(2 rows)
+
+  -- show that we keep such expressions on shell relation and shard relation
+  SELECT s.relname, a.attname, a.attgenerated
+  FROM pg_class s
+  JOIN pg_attribute a ON a.attrelid=s.oid
+  WHERE s.relname LIKE 'generated_stored_col_test%' AND
+        attname = 'y'
+  ORDER BY 1,2;
+             relname             | attname | attgenerated
+---------------------------------------------------------------------
+ generated_stored_col_test       | y       | s
+ generated_stored_col_test_60040 | y       | s
+(2 rows)
+
+ROLLBACK;
+RESET citus.replicate_reference_tables_on_activate;
+SELECT citus_remove_node('localhost', :master_port);
+ citus_remove_node
+---------------------------------------------------------------------
+
+(1 row)
+
 \set VERBOSITY terse
 drop schema test_pg12 cascade;
 NOTICE:  drop cascades to 10 other objects

src/test/regress/sql/pg12.sql

@@ -275,6 +275,30 @@ select count(*)
 from col_test
 where val = 'asdf';
 
+-- not replicate reference tables from other test files
+SET citus.replicate_reference_tables_on_activate TO off;
+SELECT 1 FROM citus_add_node('localhost', :master_port, groupId => 0);
+
+BEGIN;
+  CREATE TABLE generated_stored_col_test (x int, y int generated always as (x+1) stored);
+  SELECT citus_add_local_table_to_metadata('generated_stored_col_test');
+
+  -- simply check if GENERATED ALWAYS AS (...) STORED expression works fine
+  INSERT INTO generated_stored_col_test VALUES(1), (2);
+  SELECT * FROM generated_stored_col_test ORDER BY 1,2;
+
+  -- show that we keep such expressions on shell relation and shard relation
+  SELECT s.relname, a.attname, a.attgenerated
+  FROM pg_class s
+  JOIN pg_attribute a ON a.attrelid=s.oid
+  WHERE s.relname LIKE 'generated_stored_col_test%' AND
+        attname = 'y'
+  ORDER BY 1,2;
+ROLLBACK;
+
+RESET citus.replicate_reference_tables_on_activate;
+SELECT citus_remove_node('localhost', :master_port);
+
 \set VERBOSITY terse
 drop schema test_pg12 cascade;
 \set VERBOSITY default