Don't allow distributing by a generated column

src/backend/distributed/commands/create_distributed_table.c

@@ -103,7 +103,8 @@ static bool LocalTableEmpty(Oid tableId);
 static void CopyLocalDataIntoShards(Oid relationId);
 static List * TupleDescColumnNameList(TupleDesc tupleDescriptor);
 static bool RelationUsesIdentityColumns(TupleDesc relationDesc);
-static bool RelationUsesGeneratedStoredColumns(TupleDesc relationDesc);
+static bool DistributionColumnUsesGeneratedStoredColumn(TupleDesc relationDesc,
+														Var *distributionColumn);
 static bool RelationUsesHeapAccessMethodOrNone(Relation relation);
 static bool CanUseExclusiveConnections(Oid relationId, bool localTableEmpty);
 
@@ -684,12 +685,13 @@ EnsureRelationCanBeDistributed(Oid relationId, Var *distributionColumn,
 								  "... AS IDENTITY.")));
 	}
 
-	/* verify target relation does not use generated columns */
+	/* verify target relation is not distributed by a generated columns */
-	if (RelationUsesGeneratedStoredColumns(relationDesc))
+	if (distributionMethod != DISTRIBUTE_BY_NONE &&
+		DistributionColumnUsesGeneratedStoredColumn(relationDesc, distributionColumn))
 	{
 		ereport(ERROR, (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
 						errmsg("cannot distribute relation: %s", relationName),
-						errdetail("Distributed relations must not use GENERATED ALWAYS "
+						errdetail("Distribution column must not use GENERATED ALWAYS "
 								  "AS (...) STORED.")));
 	}
 
@@ -1385,8 +1387,8 @@ TupleDescColumnNameList(TupleDesc tupleDescriptor)
 
 
 /*
- * RelationUsesIdentityColumns returns whether a given relation uses the SQL
+ * RelationUsesIdentityColumns returns whether a given relation uses
- * GENERATED ... AS IDENTITY features introduced as of PostgreSQL 10.
+ * GENERATED ... AS IDENTITY
  */
 static bool
 RelationUsesIdentityColumns(TupleDesc relationDesc)
@@ -1408,23 +1410,20 @@ RelationUsesIdentityColumns(TupleDesc relationDesc)
 
 
 /*
- * RelationUsesIdentityColumns returns whether a given relation uses the SQL
+ * DistributionColumnUsesGeneratedStoredColumn returns whether a given relation uses
- * GENERATED ... AS IDENTITY features introduced as of PostgreSQL 10.
+ * GENERATED ALWAYS AS (...) STORED on distribution column
  */
 static bool
-RelationUsesGeneratedStoredColumns(TupleDesc relationDesc)
+DistributionColumnUsesGeneratedStoredColumn(TupleDesc relationDesc,
+											Var *distributionColumn)
 {
 #if PG_VERSION_NUM >= 120000
-	int attributeIndex = 0;
+	Form_pg_attribute attributeForm = TupleDescAttr(relationDesc,
+													distributionColumn->varattno - 1);
 
-	for (attributeIndex = 0; attributeIndex < relationDesc->natts; attributeIndex++)
+	if (attributeForm->attgenerated == ATTRIBUTE_GENERATED_STORED)
 	{
-		Form_pg_attribute attributeForm = TupleDescAttr(relationDesc, attributeIndex);
+		return true;
-
-		if (attributeForm->attgenerated == ATTRIBUTE_GENERATED_STORED)
-		{
-			return true;
-		}
 	}
 #endif
 

src/test/regress/expected/pg12.out

@@ -33,37 +33,41 @@ create table gen2 (
 );
 insert into gen1 (id, val1) values (1,4),(3,6),(5,2),(7,2);
 insert into gen2 (id, val1) values (1,4),(3,6),(5,2),(7,2);
-select * from create_distributed_table('gen1', 'id');
+select create_distributed_table('gen1', 'id');
-ERROR:  cannot distribute relation: gen1
+NOTICE:  Copying data from local table...
-DETAIL:  Distributed relations must not use GENERATED ALWAYS AS (...) STORED.
+ create_distributed_table 
-select * from create_distributed_table('gen2', 'val2');
+--------------------------
+ 
+(1 row)
+
+select create_distributed_table('gen2', 'val2');
 ERROR:  cannot distribute relation: gen2
-DETAIL:  Distributed relations must not use GENERATED ALWAYS AS (...) STORED.
+DETAIL:  Distribution column must not use GENERATED ALWAYS AS (...) STORED.
 insert into gen1 (id, val1) values (2,4),(4,6),(6,2),(8,2);
 insert into gen2 (id, val1) values (2,4),(4,6),(6,2),(8,2);
-select * from gen1;
+select * from gen1 order by 1,2,3;
  id | val1 | val2 
 ----+------+------
   1 |    4 |    6
-  3 |    6 |    8
-  5 |    2 |    4
-  7 |    2 |    4
   2 |    4 |    6
+  3 |    6 |    8
   4 |    6 |    8
+  5 |    2 |    4
   6 |    2 |    4
+  7 |    2 |    4
   8 |    2 |    4
 (8 rows)
 
-select * from gen2;
+select * from gen2 order by 1,2,3;
  id | val1 | val2 
 ----+------+------
   1 |    4 |    6
-  3 |    6 |    8
-  5 |    2 |    4
-  7 |    2 |    4
   2 |    4 |    6
+  3 |    6 |    8
   4 |    6 |    8
+  5 |    2 |    4
   6 |    2 |    4
+  7 |    2 |    4
   8 |    2 |    4
 (8 rows)
 
@@ -169,7 +173,7 @@ $Q$);
              coordinator_plan             
 ------------------------------------------
  Custom Scan (Citus Adaptive)
-   ->  Distributed Subplan 5_1
+   ->  Distributed Subplan 7_1
          ->  Custom Scan (Citus Adaptive)
                Task Count: 4
 (4 rows)
@@ -214,7 +218,7 @@ $Q$);
 ------------------------------------------------
  Aggregate
    ->  Custom Scan (Citus Adaptive)
-         ->  Distributed Subplan 8_1
+         ->  Distributed Subplan 10_1
                ->  Custom Scan (Citus Adaptive)
                      Task Count: 4
 (5 rows)
@@ -235,7 +239,7 @@ $Q$);
 ------------------------------------------------
  Aggregate
    ->  Custom Scan (Citus Adaptive)
-         ->  Distributed Subplan 10_1
+         ->  Distributed Subplan 12_1
                ->  Custom Scan (Citus Adaptive)
                      Task Count: 4
 (5 rows)
@@ -280,8 +284,8 @@ NOTICE:  Copying data from local table...
 
 -- should still fail because of fkey
 INSERT INTO collection_users VALUES (1, 1000, 1);
-ERROR:  insert or update on table "collection_users_60024" violates foreign key constraint "collection_users_fkey_60024"
+ERROR:  insert or update on table "collection_users_60028" violates foreign key constraint "collection_users_fkey_60028"
-DETAIL:  Key (key, collection_id)=(1, 1000) is not present in table "collections_list_60012".
+DETAIL:  Key (key, collection_id)=(1, 1000) is not present in table "collections_list_60016".
 CONTEXT:  while executing command on localhost:57637
 -- whereas new record with partition should go through
 INSERT INTO collections_list VALUES (2, 1, '1.2');

src/test/regress/sql/pg12.sql

@@ -39,14 +39,14 @@ create table gen2 (
 insert into gen1 (id, val1) values (1,4),(3,6),(5,2),(7,2);
 insert into gen2 (id, val1) values (1,4),(3,6),(5,2),(7,2);
 
-select * from create_distributed_table('gen1', 'id');
+select create_distributed_table('gen1', 'id');
-select * from create_distributed_table('gen2', 'val2');
+select create_distributed_table('gen2', 'val2');
 
 insert into gen1 (id, val1) values (2,4),(4,6),(6,2),(8,2);
 insert into gen2 (id, val1) values (2,4),(4,6),(6,2),(8,2);
 
-select * from gen1;
+select * from gen1 order by 1,2,3;
-select * from gen2;
+select * from gen2 order by 1,2,3;
 
 -- Test new VACUUM/ANALYZE options
 analyze (skip_locked) gen1;