Mark auto-converted=false when creating fkeys

src/backend/distributed/commands/table.c

@@ -24,11 +24,12 @@
 #include "distributed/colocation_utils.h"
 #include "distributed/commands.h"
 #include "distributed/commands/utility_hook.h"
+#include "distributed/coordinator_protocol.h"
 #include "distributed/deparser.h"
 #include "distributed/deparse_shard_query.h"
 #include "distributed/distribution_column.h"
+#include "distributed/foreign_key_relationship.h"
 #include "distributed/listutils.h"
-#include "distributed/coordinator_protocol.h"
 #include "distributed/metadata_sync.h"
 #include "distributed/metadata/dependency.h"
 #include "distributed/metadata/distobject.h"
@@ -69,6 +70,8 @@ static void ErrorIfAttachCitusTableToPgLocalTable(Oid parentRelationId,
 												  Oid partitionRelationId);
 static bool AlterTableDefinesFKeyBetweenPostgresAndNonDistTable(
 	AlterTableStmt *alterTableStatement);
+static void MarkConnectedRelationsNotAutoConverted(Oid leftRelationId,
+												   Oid rightRelationId);
 static bool RelationIdListContainsCitusTableType(List *relationIdList,
 												 CitusTableType citusTableType);
 static bool RelationIdListContainsPostgresTable(List *relationIdList);
@@ -860,6 +863,8 @@ PreprocessAlterTableStmt(Node *node, const char *alterTableCommand,
 				rightRelationId = RangeVarGetRelid(constraint->pktable, lockmode,
 												   alterTableStatement->missing_ok);
 
+				MarkConnectedRelationsNotAutoConverted(leftRelationId, rightRelationId);
+
 				/*
 				 * Foreign constraint validations will be done in workers. If we do not
 				 * set this flag, PostgreSQL tries to do additional checking when we drop
@@ -1172,6 +1177,50 @@ AlterTableDefinesFKeyBetweenPostgresAndNonDistTable(AlterTableStmt *alterTableSt
 }
 
 
+/*
+ * MarkConnectedRelationsNotAutoConverted takes two relations. If both of them are Citus Local
+ * Tables, and one of them is auto-converted while the other one is not; then it
+ * marks both of them as not-auto-converted.
+ */
+static void
+MarkConnectedRelationsNotAutoConverted(Oid leftRelationId, Oid rightRelationId)
+{
+	if (!IsCitusTable(leftRelationId) || !IsCitusTable(rightRelationId))
+	{
+		return;
+	}
+
+	if (!IsCitusTableType(leftRelationId, CITUS_LOCAL_TABLE))
+	{
+		return;
+	}
+
+	if (!IsCitusTableType(rightRelationId, CITUS_LOCAL_TABLE))
+	{
+		return;
+	}
+
+	CitusTableCacheEntry *entryLeft = GetCitusTableCacheEntry(leftRelationId);
+	CitusTableCacheEntry *entryRight = GetCitusTableCacheEntry(rightRelationId);
+
+	if (entryLeft->autoConverted == entryRight->autoConverted)
+	{
+		return;
+	}
+
+	List *leftConnectedRelIds = GetForeignKeyConnectedRelationIdList(leftRelationId);
+	List *rightConnectedRelIds = GetForeignKeyConnectedRelationIdList(rightRelationId);
+	List *allConnectedRelations = list_concat_unique_oid(leftConnectedRelIds,
+														 rightConnectedRelIds);
+
+	Oid relationId = InvalidOid;
+	foreach_oid(relationId, allConnectedRelations)
+	{
+		UpdatePartitionAutoConverted(relationId, false);
+	}
+}
+
+
 /*
  * RelationIdListContainsCitusTableType returns true if given relationIdList
  * contains a citus table with given type.
@@ -1231,11 +1280,36 @@ ConvertPostgresLocalTablesToCitusLocalTables(AlterTableStmt *alterTableStatement
 	 */
 	relationRangeVarList = SortList(relationRangeVarList, CompareRangeVarsByOid);
 
+	bool autoConverted = true;
+	RangeVar *relationRangeVar;
+	foreach_ptr(relationRangeVar, relationRangeVarList)
+	{
+		/*
+		 * Here we iterate the relation list, and if at least one of the relations
+		 * is marked as not-auto-converted, we should mark all of them as
+		 * not-auto-converted. In that case, we set the local variable autoConverted
+		 * to false here, to later use it when converting relations.
+		 */
+		List *commandList = alterTableStatement->cmds;
+		LOCKMODE lockMode = AlterTableGetLockLevel(commandList);
+		bool missingOk = alterTableStatement->missing_ok;
+		Oid relationId = RangeVarGetRelid(relationRangeVar, lockMode, missingOk);
+		if (OidIsValid(relationId) && IsCitusTable(relationId) &&
+			IsCitusTableType(relationId, CITUS_LOCAL_TABLE))
+		{
+			CitusTableCacheEntry *entry = GetCitusTableCacheEntry(relationId);
+			if (!entry->autoConverted)
+			{
+				autoConverted = false;
+				break;
+			}
+		}
+	}
+
 	/*
 	 * Here we should operate on RangeVar objects since relations oid's would
 	 * change in below loop due to CreateCitusLocalTable.
 	 */
-	RangeVar *relationRangeVar;
 	foreach_ptr(relationRangeVar, relationRangeVarList)
 	{
 		List *commandList = alterTableStatement->cmds;
@@ -1252,13 +1326,24 @@ ConvertPostgresLocalTablesToCitusLocalTables(AlterTableStmt *alterTableStatement
 		}
 		else if (IsCitusTable(relationId))
 		{
-			/*
+			CitusTableCacheEntry *entry = GetCitusTableCacheEntry(relationId);
-			 * relationRangeVarList has also reference and citus local tables
+			if (!entry->autoConverted || autoConverted)
-			 * involved in this ADD FOREIGN KEY command. Moreover, even if
+			{
-			 * relationId was belonging to a postgres  local table initially,
+				/*
-			 * we might had already converted it to a citus local table by cascading.
+				 * relationRangeVarList has also reference and citus local tables
-			 */
+				 * involved in this ADD FOREIGN KEY command. Moreover, even if
-			continue;
+				 * relationId was belonging to a postgres  local table initially,
+				 * we might had already converted it to a citus local table by cascading.
+				 * Note that we cannot skip here, if the relation is marked as
+				 * auto-converted, but we are marking the relations in the command
+				 * auto-converted = false. Because in that case, it means that this
+				 * relation is marked as auto-converted because of a connection
+				 * with a reference table, but now it is connected to a Citus
+				 * Local Table. In that case, we need to mark this relation as
+				 * auto-converted = false, so cannot do a "continue" here.
+				 */
+				continue;
+			}
 		}
 
 		/*
@@ -1291,7 +1376,6 @@ ConvertPostgresLocalTablesToCitusLocalTables(AlterTableStmt *alterTableStatement
 			}
 			else
 			{
-				bool autoConverted = true;
 				CreateCitusLocalTable(relationId, cascade, autoConverted);
 			}
 		}

src/test/regress/expected/auto_undist_citus_local.out

@@ -594,6 +594,73 @@ SELECT logicalrelid, autoconverted FROM pg_dist_partition
 ---------------------------------------------------------------------
 (0 rows)
 
+-- verify creating fkeys update auto-converted to false
+CREATE TABLE table_ref(a int unique);
+CREATE TABLE table_auto_conv(a int unique references table_ref(a)) partition by range(a);
+CREATE TABLE table_auto_conv_child partition of table_auto_conv FOR VALUES FROM (1) TO (4);
+CREATE TABLE table_auto_conv_2(a int unique references table_auto_conv(a));
+CREATE TABLE table_not_auto_conv(a int unique);
+select create_reference_table('table_ref');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+-- table_not_auto_conv should not be here, as it's not converted yet
+-- other tables should be marked as auto-converted
+SELECT logicalrelid, autoconverted FROM pg_dist_partition
+  WHERE logicalrelid IN ('table_auto_conv'::regclass,
+                        'table_auto_conv_child'::regclass,
+                        'table_auto_conv_2'::regclass,
+                        'table_not_auto_conv'::regclass)
+  ORDER BY logicalrelid;
+     logicalrelid      | autoconverted
+---------------------------------------------------------------------
+ table_auto_conv       | t
+ table_auto_conv_2     | t
+ table_auto_conv_child | t
+(3 rows)
+
+select citus_add_local_table_to_metadata('table_not_auto_conv');
+ citus_add_local_table_to_metadata
+---------------------------------------------------------------------
+
+(1 row)
+
+alter table table_not_auto_conv add constraint fkey_to_mark_not_autoconverted foreign key (a) references table_auto_conv_2(a);
+-- all of them should be marked as auto-converted = false
+SELECT logicalrelid, autoconverted FROM pg_dist_partition
+  WHERE logicalrelid IN ('table_auto_conv'::regclass,
+                        'table_auto_conv_child'::regclass,
+                        'table_auto_conv_2'::regclass,
+                        'table_not_auto_conv'::regclass)
+  ORDER BY logicalrelid;
+     logicalrelid      | autoconverted
+---------------------------------------------------------------------
+ table_auto_conv       | f
+ table_auto_conv_2     | f
+ table_auto_conv_child | f
+ table_not_auto_conv   | f
+(4 rows)
+
+-- create&attach new partition, it should be marked as auto-converted = false, too
+CREATE TABLE table_auto_conv_child_2 partition of table_auto_conv FOR VALUES FROM (5) TO (8);
+SELECT logicalrelid, autoconverted FROM pg_dist_partition
+  WHERE logicalrelid IN ('table_auto_conv'::regclass,
+                        'table_auto_conv_child'::regclass,
+                        'table_auto_conv_2'::regclass,
+                        'table_not_auto_conv'::regclass,
+                        'table_auto_conv_child_2'::regclass)
+  ORDER BY logicalrelid;
+      logicalrelid       | autoconverted
+---------------------------------------------------------------------
+ table_auto_conv         | f
+ table_auto_conv_2       | f
+ table_auto_conv_child   | f
+ table_not_auto_conv     | f
+ table_auto_conv_child_2 | f
+(5 rows)
+
 -- a single drop table cascades into multiple undistributes
 DROP TABLE IF EXISTS citus_local_table_1, citus_local_table_2, citus_local_table_3, citus_local_table_2, reference_table_1;
 CREATE TABLE reference_table_1(r1 int UNIQUE, r2 int);
@@ -631,9 +698,9 @@ ALTER TABLE reference_table_1 OWNER TO another_user;
 SELECT run_command_on_placements('reference_table_1', 'ALTER TABLE %s OWNER TO another_user');
          run_command_on_placements
 ---------------------------------------------------------------------
- (localhost,57636,1810050,t,"ALTER TABLE")
+ (localhost,57636,1810056,t,"ALTER TABLE")
- (localhost,57637,1810050,t,"ALTER TABLE")
+ (localhost,57637,1810056,t,"ALTER TABLE")
- (localhost,57638,1810050,t,"ALTER TABLE")
+ (localhost,57638,1810056,t,"ALTER TABLE")
 (3 rows)
 
 SET citus.enable_ddl_propagation to ON;

src/test/regress/sql/auto_undist_citus_local.sql

@@ -287,6 +287,44 @@ SELECT logicalrelid, autoconverted FROM pg_dist_partition
                         'partitioned_table_2_1'::regclass)
   ORDER BY logicalrelid;
 
+-- verify creating fkeys update auto-converted to false
+CREATE TABLE table_ref(a int unique);
+CREATE TABLE table_auto_conv(a int unique references table_ref(a)) partition by range(a);
+CREATE TABLE table_auto_conv_child partition of table_auto_conv FOR VALUES FROM (1) TO (4);
+CREATE TABLE table_auto_conv_2(a int unique references table_auto_conv(a));
+CREATE TABLE table_not_auto_conv(a int unique);
+select create_reference_table('table_ref');
+
+-- table_not_auto_conv should not be here, as it's not converted yet
+-- other tables should be marked as auto-converted
+SELECT logicalrelid, autoconverted FROM pg_dist_partition
+  WHERE logicalrelid IN ('table_auto_conv'::regclass,
+                        'table_auto_conv_child'::regclass,
+                        'table_auto_conv_2'::regclass,
+                        'table_not_auto_conv'::regclass)
+  ORDER BY logicalrelid;
+
+select citus_add_local_table_to_metadata('table_not_auto_conv');
+alter table table_not_auto_conv add constraint fkey_to_mark_not_autoconverted foreign key (a) references table_auto_conv_2(a);
+
+-- all of them should be marked as auto-converted = false
+SELECT logicalrelid, autoconverted FROM pg_dist_partition
+  WHERE logicalrelid IN ('table_auto_conv'::regclass,
+                        'table_auto_conv_child'::regclass,
+                        'table_auto_conv_2'::regclass,
+                        'table_not_auto_conv'::regclass)
+  ORDER BY logicalrelid;
+
+-- create&attach new partition, it should be marked as auto-converted = false, too
+CREATE TABLE table_auto_conv_child_2 partition of table_auto_conv FOR VALUES FROM (5) TO (8);
+SELECT logicalrelid, autoconverted FROM pg_dist_partition
+  WHERE logicalrelid IN ('table_auto_conv'::regclass,
+                        'table_auto_conv_child'::regclass,
+                        'table_auto_conv_2'::regclass,
+                        'table_not_auto_conv'::regclass,
+                        'table_auto_conv_child_2'::regclass)
+  ORDER BY logicalrelid;
+
 -- a single drop table cascades into multiple undistributes
 DROP TABLE IF EXISTS citus_local_table_1, citus_local_table_2, citus_local_table_3, citus_local_table_2, reference_table_1;
 CREATE TABLE reference_table_1(r1 int UNIQUE, r2 int);