From 4e525fc170bf5c449166343424096b0ed071a94e Mon Sep 17 00:00:00 2001
From: EPoikans <eriks.poikans@gmail.com>
Date: Fri, 14 Jun 2024 11:26:27 +0000
Subject: [PATCH 1/2] CommentMigration

---
 .../distributed/commands/alter_table.c        |  39 +++++-
 src/test/regress/Makefile                     |   3 +
 src/test/regress/comment_migration_schedule   |   7 ++
 .../expected/comment_migration_column.out     |  65 ++++++++++
 ...comment_migration_column_joined_tables.out | 101 +++++++++++++++
 ...ment_migration_column_joined_tables_FK.out |  98 +++++++++++++++
 .../expected/comment_migration_table.out      |  61 +++++++++
 .../comment_migration_table_joined_tables.out | 119 ++++++++++++++++++
 ...mment_migration_table_joined_tables_FK.out | 116 +++++++++++++++++
 src/test/regress/multi_schedule               |  10 ++
 .../regress/sql/comment_migration_column.sql  |  33 +++++
 ...comment_migration_column_joined_tables.sql |  54 ++++++++
 ...ment_migration_column_joined_tables_FK.sql |  56 +++++++++
 .../regress/sql/comment_migration_table.sql   |  29 +++++
 .../comment_migration_table_joined_tables.sql |  49 ++++++++
 ...mment_migration_table_joined_tables_FK.sql |  51 ++++++++
 16 files changed, 887 insertions(+), 4 deletions(-)
 create mode 100644 src/test/regress/comment_migration_schedule
 create mode 100644 src/test/regress/expected/comment_migration_column.out
 create mode 100644 src/test/regress/expected/comment_migration_column_joined_tables.out
 create mode 100644 src/test/regress/expected/comment_migration_column_joined_tables_FK.out
 create mode 100644 src/test/regress/expected/comment_migration_table.out
 create mode 100644 src/test/regress/expected/comment_migration_table_joined_tables.out
 create mode 100644 src/test/regress/expected/comment_migration_table_joined_tables_FK.out
 create mode 100644 src/test/regress/sql/comment_migration_column.sql
 create mode 100644 src/test/regress/sql/comment_migration_column_joined_tables.sql
 create mode 100644 src/test/regress/sql/comment_migration_column_joined_tables_FK.sql
 create mode 100644 src/test/regress/sql/comment_migration_table.sql
 create mode 100644 src/test/regress/sql/comment_migration_table_joined_tables.sql
 create mode 100644 src/test/regress/sql/comment_migration_table_joined_tables_FK.sql

diff --git a/src/backend/distributed/commands/alter_table.c b/src/backend/distributed/commands/alter_table.c
index 030dbbe7869..c62a02b7197 100644
--- a/src/backend/distributed/commands/alter_table.c
+++ b/src/backend/distributed/commands/alter_table.c
@@ -35,6 +35,7 @@
 #include "catalog/pg_depend.h"
 #include "catalog/pg_rewrite_d.h"
 #include "commands/defrem.h"
+#include "commands/comment.h"
 #include "executor/spi.h"
 #include "nodes/pg_list.h"
 #include "utils/builtins.h"
@@ -844,7 +845,6 @@ ConvertTableInternal(TableConversionState *con)
 				 */
 				.cascadeViaForeignKeys = false
 			};
-
 			TableConversionReturn *partitionReturn = con->function(&partitionParam);
 			if (cascadeOption == CASCADE_TO_COLOCATED_NO_ALREADY_CASCADED)
 			{
@@ -1032,7 +1032,6 @@ ConvertTableInternal(TableConversionState *con)
 
 	/* increment command counter so that next command can see the new table */
 	CommandCounterIncrement();
-
 	SetLocalEnableLocalReferenceForeignKeys(oldEnableLocalReferenceForeignKeys);
 
 	InTableTypeConversionFunctionCall = false;
@@ -1780,6 +1779,35 @@ CreateMaterializedViewDDLCommand(Oid matViewOid)
 	return query->data;
 }
 
+/*
+* MigrateColumnComments migrates distributed table column comments to the target undistributed table columns.
+*/
+static void
+MigrateColumnComments(Oid sourceId, Oid targetId){
+	Relation relation = relation_open(sourceId, AccessShareLock);
+	TupleDesc tupleDesc = RelationGetDescr(relation);
+	for (int attrNum = 0; attrNum < tupleDesc->natts; attrNum++)
+	{
+		Form_pg_attribute attr = TupleDescAttr(tupleDesc, attrNum);
+		if (!attr->attisdropped)
+		{
+			char *columnComment = GetComment(sourceId, RelationRelationId ,attrNum + 1);
+			CreateComments(targetId, RelationRelationId, attrNum + 1, columnComment);
+		}
+	}
+	relation_close(relation, AccessShareLock);
+}
+
+/*
+ * MigrateTableComment migrates the comment of the source distributed table to the target undistributed table.
+*/
+static void
+MigrateTableComment(Oid sourceId, Oid targetId){
+	char *comment = GetComment(sourceId, RelationRelationId, 0);
+	if(comment != NULL) {
+		CreateComments(targetId, RelationRelationId, 0, comment);
+	}
+}
 
 /*
  * ReplaceTable replaces the source table with the target table.
@@ -1796,11 +1824,11 @@ ReplaceTable(Oid sourceId, Oid targetId, List *justBeforeDropCommands,
 	char *sourceName = get_rel_name(sourceId);
 	char *qualifiedSourceName = generate_qualified_relation_name(sourceId);
 	char *qualifiedTargetName = generate_qualified_relation_name(targetId);
-
 	StringInfo query = makeStringInfo();
 
 	if (!PartitionedTable(sourceId) && !IsForeignTable(sourceId))
 	{
+		
 		if (!suppressNoticeMessages)
 		{
 			ereport(NOTICE, (errmsg("moving the data of %s", qualifiedSourceName)));
@@ -1814,6 +1842,7 @@ ReplaceTable(Oid sourceId, Oid targetId, List *justBeforeDropCommands,
 			 */
 			appendStringInfo(query, "INSERT INTO %s SELECT * FROM %s",
 							 qualifiedTargetName, qualifiedSourceName);
+			MigrateColumnComments(sourceId, targetId);
 		}
 		else
 		{
@@ -1833,6 +1862,7 @@ ReplaceTable(Oid sourceId, Oid targetId, List *justBeforeDropCommands,
 		}
 
 		ExecuteQueryViaSPI(query->data, SPI_OK_INSERT);
+		MigrateColumnComments(sourceId, targetId);
 	}
 
 	/*
@@ -1883,6 +1913,8 @@ ReplaceTable(Oid sourceId, Oid targetId, List *justBeforeDropCommands,
 		ereport(NOTICE, (errmsg("dropping the old %s", qualifiedSourceName)));
 	}
 
+	MigrateTableComment(sourceId, targetId);
+
 	resetStringInfo(query);
 	appendStringInfo(query, "DROP %sTABLE %s CASCADE",
 					 IsForeignTable(sourceId) ? "FOREIGN " : "",
@@ -1901,7 +1933,6 @@ ReplaceTable(Oid sourceId, Oid targetId, List *justBeforeDropCommands,
 	ExecuteQueryViaSPI(query->data, SPI_OK_UTILITY);
 }
 
-
 /*
  * HasAnyGeneratedStoredColumns decides if relation has any columns that we
  * might need to copy the data of when replacing table.
diff --git a/src/test/regress/Makefile b/src/test/regress/Makefile
index 4bdc7a1b8a5..ead7dc97ae2 100644
--- a/src/test/regress/Makefile
+++ b/src/test/regress/Makefile
@@ -146,6 +146,9 @@ check-isolation-custom-schedule-vg: all  $(isolation_test_files)
 	--valgrind --pg_ctl-timeout=360 --connection-timeout=500000 --valgrind-path=valgrind --valgrind-log-file=$(CITUS_VALGRIND_LOG_FILE) \
 	-- $(MULTI_REGRESS_OPTS) --inputdir=$(citus_abs_srcdir)/build --schedule=$(citus_abs_srcdir)/$(SCHEDULE) $(EXTRA_TESTS)
 
+check-comment-migration: all
+	$(pg_regress_multi_check) --load-extension=citus \
+	-- $(MULTI_REGRESS_OPTS) --schedule=$(citus_abs_srcdir)/comment_migration_schedule $(EXTRA_TESTS)
 
 check-empty: all
 	$(pg_regress_multi_check) --load-extension=citus \
diff --git a/src/test/regress/comment_migration_schedule b/src/test/regress/comment_migration_schedule
new file mode 100644
index 00000000000..2c11f1458b3
--- /dev/null
+++ b/src/test/regress/comment_migration_schedule
@@ -0,0 +1,7 @@
+# Tests comment migration during undistribute_table
+test: comment_migration_table
+test: comment_migration_table_joined_tables
+test: comment_migration_table_joined_tables_FK
+test: comment_migration_column
+test: comment_migration_column_joined_tables
+test: comment_migration_column_joined_tables_FK
\ No newline at end of file
diff --git a/src/test/regress/expected/comment_migration_column.out b/src/test/regress/expected/comment_migration_column.out
new file mode 100644
index 00000000000..1880801a013
--- /dev/null
+++ b/src/test/regress/expected/comment_migration_column.out
@@ -0,0 +1,65 @@
+CREATE SCHEMA comment_migration_column;
+CREATE TABLE comment_migration_column.table_1
+(
+    id            bigserial,
+    code          varchar(200)                           not null,
+    name          varchar(200),
+    date_created  timestamp with time zone default NOW() not null,
+    active        boolean                  default true,
+    CONSTRAINT table_1_pkey PRIMARY KEY (id)
+)
+    WITH (autovacuum_enabled = TRUE);
+SET citus.shard_replication_factor = 1;
+comment on column comment_migration_column.table_1.id is 'table id';
+comment on column comment_migration_column.table_1.code is 'table code';
+comment on column comment_migration_column.table_1.name is 'table name';
+comment on column comment_migration_column.table_1.date_created is 'table date_created';
+comment on column comment_migration_column.table_1.active is 'table active';
+select col_description('comment_migration_column.table_1'::regclass,1), col_description('comment_migration_column.table_1'::regclass,2), col_description('comment_migration_column.table_1'::regclass,3), col_description('comment_migration_column.table_1'::regclass,4), col_description('comment_migration_column.table_1'::regclass,5);
+ col_description | col_description | col_description |  col_description   | col_description
+---------------------------------------------------------------------
+ table id        | table code      | table name      | table date_created | table active
+(1 row)
+
+SELECT create_distributed_table('comment_migration_column.table_1', 'id');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+select col_description('comment_migration_column.table_1'::regclass,1), col_description('comment_migration_column.table_1'::regclass,2), col_description('comment_migration_column.table_1'::regclass,3), col_description('comment_migration_column.table_1'::regclass,4), col_description('comment_migration_column.table_1'::regclass,5);
+ col_description | col_description | col_description |  col_description   | col_description
+---------------------------------------------------------------------
+ table id        | table code      | table name      | table date_created | table active
+(1 row)
+
+select undistribute_table('comment_migration_column.table_1');
+NOTICE:  creating a new table for comment_migration_column.table_1
+NOTICE:  moving the data of comment_migration_column.table_1
+NOTICE:  dropping the old comment_migration_column.table_1
+NOTICE:  renaming the new table to comment_migration_column.table_1
+ undistribute_table
+---------------------------------------------------------------------
+
+(1 row)
+
+select col_description('comment_migration_column.table_1'::regclass,1), col_description('comment_migration_column.table_1'::regclass,2), col_description('comment_migration_column.table_1'::regclass,3), col_description('comment_migration_column.table_1'::regclass,4), col_description('comment_migration_column.table_1'::regclass,5);
+ col_description | col_description | col_description |  col_description   | col_description
+---------------------------------------------------------------------
+ table id        | table code      | table name      | table date_created | table active
+(1 row)
+
+SELECT create_distributed_table('comment_migration_column.table_1', 'id');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+select col_description('comment_migration_column.table_1'::regclass,1), col_description('comment_migration_column.table_1'::regclass,2), col_description('comment_migration_column.table_1'::regclass,3), col_description('comment_migration_column.table_1'::regclass,4), col_description('comment_migration_column.table_1'::regclass,5);
+ col_description | col_description | col_description |  col_description   | col_description
+---------------------------------------------------------------------
+ table id        | table code      | table name      | table date_created | table active
+(1 row)
+
+DROP TABLE comment_migration_column.table_1;
+DROP SCHEMA comment_migration_column;
diff --git a/src/test/regress/expected/comment_migration_column_joined_tables.out b/src/test/regress/expected/comment_migration_column_joined_tables.out
new file mode 100644
index 00000000000..8a09237d110
--- /dev/null
+++ b/src/test/regress/expected/comment_migration_column_joined_tables.out
@@ -0,0 +1,101 @@
+CREATE SCHEMA comment_migration_column;
+CREATE TABLE comment_migration_column.table_1
+(
+    id            bigserial,
+    code          varchar(200)                           not null,
+    name          varchar(200),
+    date_created  timestamp with time zone default NOW() not null,
+    active        boolean                  default true,
+    tenant_id     bigint,
+    CONSTRAINT table_1_pkey PRIMARY KEY (id, tenant_id)
+)
+    WITH (autovacuum_enabled = TRUE);
+CREATE TABLE comment_migration_column.table_2
+(
+    id            bigserial,
+    table_1_id    bigint ,
+    description   varchar(200),
+    tenant_id     bigint,
+    CONSTRAINT table_2_pkey PRIMARY KEY (id, tenant_id)
+)
+    WITH (autovacuum_enabled = TRUE);
+SET citus.shard_replication_factor = 1;
+comment on column comment_migration_column.table_1.id is 'table 1 id';
+comment on column comment_migration_column.table_1.code is 'table 1 code';
+comment on column comment_migration_column.table_1.name is 'table 1 name';
+comment on column comment_migration_column.table_1.date_created is 'table 1 date_created';
+comment on column comment_migration_column.table_1.active is 'table 1 active';
+comment on column comment_migration_column.table_2.id is 'table 2 id';
+comment on column comment_migration_column.table_2.table_1_id is 'table 2 foreign key to table 1';
+comment on column comment_migration_column.table_2.description is 'table 2 description';
+select col_description('comment_migration_column.table_1'::regclass,1), col_description('comment_migration_column.table_1'::regclass,2), col_description('comment_migration_column.table_1'::regclass,3), col_description('comment_migration_column.table_1'::regclass,4), col_description('comment_migration_column.table_1'::regclass,5);
+ col_description | col_description | col_description |   col_description    | col_description
+---------------------------------------------------------------------
+ table 1 id      | table 1 code    | table 1 name    | table 1 date_created | table 1 active
+(1 row)
+
+select col_description('comment_migration_column.table_2'::regclass,1), col_description('comment_migration_column.table_2'::regclass,2), col_description('comment_migration_column.table_2'::regclass,3);
+ col_description |        col_description         |   col_description
+---------------------------------------------------------------------
+ table 2 id      | table 2 foreign key to table 1 | table 2 description
+(1 row)
+
+SELECT create_distributed_table('comment_migration_column.table_1', 'tenant_id');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT create_distributed_table('comment_migration_column.table_2', 'tenant_id', colocate_with=>'comment_migration_column.table_1');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT undistribute_table('comment_migration_column.table_1');
+NOTICE:  creating a new table for comment_migration_column.table_1
+NOTICE:  moving the data of comment_migration_column.table_1
+NOTICE:  dropping the old comment_migration_column.table_1
+NOTICE:  renaming the new table to comment_migration_column.table_1
+ undistribute_table
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT undistribute_table('comment_migration_column.table_2');
+NOTICE:  creating a new table for comment_migration_column.table_2
+NOTICE:  moving the data of comment_migration_column.table_2
+NOTICE:  dropping the old comment_migration_column.table_2
+NOTICE:  renaming the new table to comment_migration_column.table_2
+ undistribute_table
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT create_distributed_table('comment_migration_column.table_1', 'tenant_id');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT create_distributed_table('comment_migration_column.table_2', 'tenant_id', colocate_with=>'comment_migration_column.table_1');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+select col_description('comment_migration_column.table_1'::regclass,1), col_description('comment_migration_column.table_1'::regclass,2), col_description('comment_migration_column.table_1'::regclass,3), col_description('comment_migration_column.table_1'::regclass,4), col_description('comment_migration_column.table_1'::regclass,5);
+ col_description | col_description | col_description |   col_description    | col_description
+---------------------------------------------------------------------
+ table 1 id      | table 1 code    | table 1 name    | table 1 date_created | table 1 active
+(1 row)
+
+select col_description('comment_migration_column.table_2'::regclass,1), col_description('comment_migration_column.table_2'::regclass,2), col_description('comment_migration_column.table_2'::regclass,3);
+ col_description |        col_description         |   col_description
+---------------------------------------------------------------------
+ table 2 id      | table 2 foreign key to table 1 | table 2 description
+(1 row)
+
+DROP TABLE comment_migration_column.table_2;
+DROP TABLE comment_migration_column.table_1;
+DROP SCHEMA comment_migration_column;
diff --git a/src/test/regress/expected/comment_migration_column_joined_tables_FK.out b/src/test/regress/expected/comment_migration_column_joined_tables_FK.out
new file mode 100644
index 00000000000..a343b974639
--- /dev/null
+++ b/src/test/regress/expected/comment_migration_column_joined_tables_FK.out
@@ -0,0 +1,98 @@
+CREATE SCHEMA comment_migration_column;
+CREATE TABLE comment_migration_column.table_1
+(
+    id            bigserial,
+    code          varchar(200)                           not null,
+    name          varchar(200),
+    date_created  timestamp with time zone default NOW() not null,
+    active        boolean                  default true,
+    owner_id     bigint,
+    CONSTRAINT table_1_pkey PRIMARY KEY (id)
+)
+    WITH (autovacuum_enabled = TRUE);
+CREATE TABLE comment_migration_column.table_2
+(
+    id            bigserial,
+    table_1_id    bigint ,
+    description   varchar(200),
+    owner_id     bigint,
+    CONSTRAINT table_2_pkey PRIMARY KEY (id)
+)
+    WITH (autovacuum_enabled = TRUE);
+SET citus.shard_replication_factor = 1;
+comment on column comment_migration_column.table_1.id is 'table 1 id';
+comment on column comment_migration_column.table_1.code is 'table 1 code';
+comment on column comment_migration_column.table_1.name is 'table 1 name';
+comment on column comment_migration_column.table_1.date_created is 'table 1 date_created';
+comment on column comment_migration_column.table_1.active is 'table 1 active';
+comment on column comment_migration_column.table_2.id is 'table 2 id';
+comment on column comment_migration_column.table_2.table_1_id is 'table 2 foreign key to table 1';
+comment on column comment_migration_column.table_2.description is 'table 2 description';
+select col_description('comment_migration_column.table_1'::regclass,1), col_description('comment_migration_column.table_1'::regclass,2), col_description('comment_migration_column.table_1'::regclass,3), col_description('comment_migration_column.table_1'::regclass,4), col_description('comment_migration_column.table_1'::regclass,5);
+ col_description | col_description | col_description |   col_description    | col_description
+---------------------------------------------------------------------
+ table 1 id      | table 1 code    | table 1 name    | table 1 date_created | table 1 active
+(1 row)
+
+select col_description('comment_migration_column.table_2'::regclass,1), col_description('comment_migration_column.table_2'::regclass,2), col_description('comment_migration_column.table_2'::regclass,3);
+ col_description |        col_description         |   col_description
+---------------------------------------------------------------------
+ table 2 id      | table 2 foreign key to table 1 | table 2 description
+(1 row)
+
+SELECT create_reference_table('comment_migration_column.table_1');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT create_distributed_table('comment_migration_column.table_2', 'id');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+ALTER TABLE comment_migration_column.table_2 ADD CONSTRAINT table2_table1_fk FOREIGN KEY (table_1_id) REFERENCES comment_migration_column.table_1(id);
+SELECT undistribute_table('comment_migration_column.table_1', cascade_via_foreign_keys=>true);
+NOTICE:  creating a new table for comment_migration_column.table_1
+NOTICE:  moving the data of comment_migration_column.table_1
+NOTICE:  dropping the old comment_migration_column.table_1
+NOTICE:  renaming the new table to comment_migration_column.table_1
+NOTICE:  creating a new table for comment_migration_column.table_2
+NOTICE:  moving the data of comment_migration_column.table_2
+NOTICE:  dropping the old comment_migration_column.table_2
+NOTICE:  renaming the new table to comment_migration_column.table_2
+ undistribute_table
+---------------------------------------------------------------------
+
+(1 row)
+
+ALTER TABLE comment_migration_column.table_2 DROP CONSTRAINT table2_table1_fk;
+SELECT create_reference_table('comment_migration_column.table_1');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT create_distributed_table('comment_migration_column.table_2', 'id');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+ALTER TABLE comment_migration_column.table_2 ADD CONSTRAINT table2_table1_fk FOREIGN KEY (table_1_id) REFERENCES comment_migration_column.table_1(id);
+select col_description('comment_migration_column.table_1'::regclass,1), col_description('comment_migration_column.table_1'::regclass,2), col_description('comment_migration_column.table_1'::regclass,3), col_description('comment_migration_column.table_1'::regclass,4), col_description('comment_migration_column.table_1'::regclass,5);
+ col_description | col_description | col_description |   col_description    | col_description
+---------------------------------------------------------------------
+ table 1 id      | table 1 code    | table 1 name    | table 1 date_created | table 1 active
+(1 row)
+
+select col_description('comment_migration_column.table_2'::regclass,1), col_description('comment_migration_column.table_2'::regclass,2), col_description('comment_migration_column.table_2'::regclass,3);
+ col_description |        col_description         |   col_description
+---------------------------------------------------------------------
+ table 2 id      | table 2 foreign key to table 1 | table 2 description
+(1 row)
+
+DROP TABLE comment_migration_column.table_2;
+DROP TABLE comment_migration_column.table_1;
+DROP SCHEMA comment_migration_column;
diff --git a/src/test/regress/expected/comment_migration_table.out b/src/test/regress/expected/comment_migration_table.out
new file mode 100644
index 00000000000..4eeddc0a493
--- /dev/null
+++ b/src/test/regress/expected/comment_migration_table.out
@@ -0,0 +1,61 @@
+CREATE SCHEMA comment_migration_table;
+CREATE TABLE comment_migration_table.table_1
+(
+    id            bigserial,
+    code          varchar(200)                           not null,
+    name          varchar(200),
+    date_created  timestamp with time zone default NOW() not null,
+    active        boolean                  default true,
+    CONSTRAINT table_1_pkey PRIMARY KEY (id)
+)
+    WITH (autovacuum_enabled = TRUE);
+SET citus.shard_replication_factor = 1;
+comment on table comment_migration_table.table_1 is 'Table 1';
+select obj_description('comment_migration_table.table_1'::regclass);
+ obj_description
+---------------------------------------------------------------------
+ Table 1
+(1 row)
+
+SELECT create_distributed_table('comment_migration_table.table_1', 'id');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+select obj_description('comment_migration_table.table_1'::regclass);
+ obj_description
+---------------------------------------------------------------------
+ Table 1
+(1 row)
+
+select undistribute_table('comment_migration_table.table_1');
+NOTICE:  creating a new table for comment_migration_table.table_1
+NOTICE:  moving the data of comment_migration_table.table_1
+NOTICE:  dropping the old comment_migration_table.table_1
+NOTICE:  renaming the new table to comment_migration_table.table_1
+ undistribute_table
+---------------------------------------------------------------------
+
+(1 row)
+
+select obj_description('comment_migration_table.table_1'::regclass);
+ obj_description
+---------------------------------------------------------------------
+ Table 1
+(1 row)
+
+SELECT create_distributed_table('comment_migration_table.table_1', 'id');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+select obj_description('comment_migration_table.table_1'::regclass);
+ obj_description
+---------------------------------------------------------------------
+ Table 1
+(1 row)
+
+DROP TABLE comment_migration_table.table_1;
+DROP SCHEMA comment_migration_table;
diff --git a/src/test/regress/expected/comment_migration_table_joined_tables.out b/src/test/regress/expected/comment_migration_table_joined_tables.out
new file mode 100644
index 00000000000..91126289177
--- /dev/null
+++ b/src/test/regress/expected/comment_migration_table_joined_tables.out
@@ -0,0 +1,119 @@
+CREATE SCHEMA comment_migration_table;
+CREATE TABLE comment_migration_table.table_1
+(
+    id            bigserial,
+    code          varchar(200)                           not null,
+    name          varchar(200),
+    date_created  timestamp with time zone default NOW() not null,
+    active        boolean                  default true,
+    tenant_id     bigint,
+    CONSTRAINT table_1_pkey PRIMARY KEY (id, tenant_id)
+)
+    WITH (autovacuum_enabled = TRUE);
+CREATE TABLE comment_migration_table.table_2
+(
+    id            bigserial,
+    table_1_id    bigint ,
+    description   varchar(200),
+    tenant_id     bigint,
+    CONSTRAINT table_2_pkey PRIMARY KEY (id, tenant_id)
+)
+    WITH (autovacuum_enabled = TRUE);
+SET citus.shard_replication_factor = 1;
+comment on table comment_migration_table.table_1 is 'Table 1';
+comment on table comment_migration_table.table_2 is 'Table 2';
+select obj_description('comment_migration_table.table_1'::regclass);
+ obj_description
+---------------------------------------------------------------------
+ Table 1
+(1 row)
+
+select obj_description('comment_migration_table.table_2'::regclass);
+ obj_description
+---------------------------------------------------------------------
+ Table 2
+(1 row)
+
+SELECT create_distributed_table('comment_migration_table.table_1', 'tenant_id');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT create_distributed_table('comment_migration_table.table_2', 'tenant_id', colocate_with=>'comment_migration_table.table_1');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+select obj_description('comment_migration_table.table_1'::regclass);
+ obj_description
+---------------------------------------------------------------------
+ Table 1
+(1 row)
+
+select obj_description('comment_migration_table.table_2'::regclass);
+ obj_description
+---------------------------------------------------------------------
+ Table 2
+(1 row)
+
+SELECT undistribute_table('comment_migration_table.table_1');
+NOTICE:  creating a new table for comment_migration_table.table_1
+NOTICE:  moving the data of comment_migration_table.table_1
+NOTICE:  dropping the old comment_migration_table.table_1
+NOTICE:  renaming the new table to comment_migration_table.table_1
+ undistribute_table
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT undistribute_table('comment_migration_table.table_2');
+NOTICE:  creating a new table for comment_migration_table.table_2
+NOTICE:  moving the data of comment_migration_table.table_2
+NOTICE:  dropping the old comment_migration_table.table_2
+NOTICE:  renaming the new table to comment_migration_table.table_2
+ undistribute_table
+---------------------------------------------------------------------
+
+(1 row)
+
+select obj_description('comment_migration_table.table_1'::regclass);
+ obj_description
+---------------------------------------------------------------------
+ Table 1
+(1 row)
+
+select obj_description('comment_migration_table.table_2'::regclass);
+ obj_description
+---------------------------------------------------------------------
+ Table 2
+(1 row)
+
+SELECT create_distributed_table('comment_migration_table.table_1', 'tenant_id');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT create_distributed_table('comment_migration_table.table_2', 'tenant_id', colocate_with=>'comment_migration_table.table_1');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+select obj_description('comment_migration_table.table_1'::regclass);
+ obj_description
+---------------------------------------------------------------------
+ Table 1
+(1 row)
+
+select obj_description('comment_migration_table.table_2'::regclass);
+ obj_description
+---------------------------------------------------------------------
+ Table 2
+(1 row)
+
+DROP TABLE comment_migration_table.table_2;
+DROP TABLE comment_migration_table.table_1;
+DROP SCHEMA comment_migration_table;
diff --git a/src/test/regress/expected/comment_migration_table_joined_tables_FK.out b/src/test/regress/expected/comment_migration_table_joined_tables_FK.out
new file mode 100644
index 00000000000..ad3d4fc201d
--- /dev/null
+++ b/src/test/regress/expected/comment_migration_table_joined_tables_FK.out
@@ -0,0 +1,116 @@
+CREATE SCHEMA comment_migration_table;
+CREATE TABLE comment_migration_table.table_1
+(
+    id            bigserial,
+    code          varchar(200)                           not null,
+    name          varchar(200),
+    date_created  timestamp with time zone default NOW() not null,
+    active        boolean                  default true,
+    owner_id     bigint,
+    CONSTRAINT table_1_pkey PRIMARY KEY (id)
+)
+    WITH (autovacuum_enabled = TRUE);
+CREATE TABLE comment_migration_table.table_2
+(
+    id            bigserial,
+    table_1_id    bigint ,
+    description   varchar(200),
+    owner_id     bigint,
+    CONSTRAINT table_2_pkey PRIMARY KEY (id)
+)
+    WITH (autovacuum_enabled = TRUE);
+SET citus.shard_replication_factor = 1;
+comment on table comment_migration_table.table_1 is 'Table 1';
+comment on table comment_migration_table.table_2 is 'Table 2';
+select obj_description('comment_migration_table.table_1'::regclass);
+ obj_description
+---------------------------------------------------------------------
+ Table 1
+(1 row)
+
+select obj_description('comment_migration_table.table_2'::regclass);
+ obj_description
+---------------------------------------------------------------------
+ Table 2
+(1 row)
+
+SELECT create_reference_table('comment_migration_table.table_1');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT create_distributed_table('comment_migration_table.table_2', 'id');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+ALTER TABLE comment_migration_table.table_2 ADD CONSTRAINT table2_table1_fk FOREIGN KEY (table_1_id) REFERENCES comment_migration_table.table_1(id);
+select obj_description('comment_migration_table.table_1'::regclass);
+ obj_description
+---------------------------------------------------------------------
+ Table 1
+(1 row)
+
+select obj_description('comment_migration_table.table_2'::regclass);
+ obj_description
+---------------------------------------------------------------------
+ Table 2
+(1 row)
+
+SELECT undistribute_table('comment_migration_table.table_1', cascade_via_foreign_keys=>true);
+NOTICE:  creating a new table for comment_migration_table.table_1
+NOTICE:  moving the data of comment_migration_table.table_1
+NOTICE:  dropping the old comment_migration_table.table_1
+NOTICE:  renaming the new table to comment_migration_table.table_1
+NOTICE:  creating a new table for comment_migration_table.table_2
+NOTICE:  moving the data of comment_migration_table.table_2
+NOTICE:  dropping the old comment_migration_table.table_2
+NOTICE:  renaming the new table to comment_migration_table.table_2
+ undistribute_table
+---------------------------------------------------------------------
+
+(1 row)
+
+select obj_description('comment_migration_table.table_1'::regclass);
+ obj_description
+---------------------------------------------------------------------
+ Table 1
+(1 row)
+
+select obj_description('comment_migration_table.table_2'::regclass);
+ obj_description
+---------------------------------------------------------------------
+ Table 2
+(1 row)
+
+ALTER TABLE comment_migration_table.table_2 DROP CONSTRAINT table2_table1_fk;
+SELECT create_reference_table('comment_migration_table.table_1');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT create_distributed_table('comment_migration_table.table_2', 'id');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+ALTER TABLE comment_migration_table.table_2 ADD CONSTRAINT table2_table1_fk FOREIGN KEY (table_1_id) REFERENCES comment_migration_table.table_1(id);
+select obj_description('comment_migration_table.table_1'::regclass);
+ obj_description
+---------------------------------------------------------------------
+ Table 1
+(1 row)
+
+select obj_description('comment_migration_table.table_2'::regclass);
+ obj_description
+---------------------------------------------------------------------
+ Table 2
+(1 row)
+
+DROP TABLE comment_migration_table.table_2;
+DROP TABLE comment_migration_table.table_1;
+DROP SCHEMA comment_migration_table;
diff --git a/src/test/regress/multi_schedule b/src/test/regress/multi_schedule
index fca36f5ab06..4ed8e135ada 100644
--- a/src/test/regress/multi_schedule
+++ b/src/test/regress/multi_schedule
@@ -135,3 +135,13 @@ test: check_mx
 test: generated_identity
 test: drop_database
 test: check_cluster_state
+
+# ----------
+# Tests for migrating comments while undistributing tables
+# ----------
+test: comment_migration_table
+test: comment_migration_table_joined_tables
+test: comment_migration_table_joined_tables_FK
+test: comment_migration_column
+test: comment_migration_column_joined_tables
+test: comment_migration_column_joined_tables_FK
\ No newline at end of file
diff --git a/src/test/regress/sql/comment_migration_column.sql b/src/test/regress/sql/comment_migration_column.sql
new file mode 100644
index 00000000000..3f9753894ab
--- /dev/null
+++ b/src/test/regress/sql/comment_migration_column.sql
@@ -0,0 +1,33 @@
+CREATE SCHEMA comment_migration_column;
+
+CREATE TABLE comment_migration_column.table_1
+(
+    id            bigserial,
+    code          varchar(200)                           not null,
+    name          varchar(200),
+    date_created  timestamp with time zone default NOW() not null,
+    active        boolean                  default true,
+    CONSTRAINT table_1_pkey PRIMARY KEY (id)
+)
+    WITH (autovacuum_enabled = TRUE);
+
+SET citus.shard_replication_factor = 1;
+
+comment on column comment_migration_column.table_1.id is 'table id';
+comment on column comment_migration_column.table_1.code is 'table code';
+comment on column comment_migration_column.table_1.name is 'table name';
+comment on column comment_migration_column.table_1.date_created is 'table date_created';
+comment on column comment_migration_column.table_1.active is 'table active';
+select col_description('comment_migration_column.table_1'::regclass,1), col_description('comment_migration_column.table_1'::regclass,2), col_description('comment_migration_column.table_1'::regclass,3), col_description('comment_migration_column.table_1'::regclass,4), col_description('comment_migration_column.table_1'::regclass,5);
+
+SELECT create_distributed_table('comment_migration_column.table_1', 'id');
+select col_description('comment_migration_column.table_1'::regclass,1), col_description('comment_migration_column.table_1'::regclass,2), col_description('comment_migration_column.table_1'::regclass,3), col_description('comment_migration_column.table_1'::regclass,4), col_description('comment_migration_column.table_1'::regclass,5);
+
+select undistribute_table('comment_migration_column.table_1');
+select col_description('comment_migration_column.table_1'::regclass,1), col_description('comment_migration_column.table_1'::regclass,2), col_description('comment_migration_column.table_1'::regclass,3), col_description('comment_migration_column.table_1'::regclass,4), col_description('comment_migration_column.table_1'::regclass,5);
+
+SELECT create_distributed_table('comment_migration_column.table_1', 'id');
+select col_description('comment_migration_column.table_1'::regclass,1), col_description('comment_migration_column.table_1'::regclass,2), col_description('comment_migration_column.table_1'::regclass,3), col_description('comment_migration_column.table_1'::regclass,4), col_description('comment_migration_column.table_1'::regclass,5);
+
+DROP TABLE comment_migration_column.table_1;
+DROP SCHEMA comment_migration_column;
\ No newline at end of file
diff --git a/src/test/regress/sql/comment_migration_column_joined_tables.sql b/src/test/regress/sql/comment_migration_column_joined_tables.sql
new file mode 100644
index 00000000000..f3220e9b23b
--- /dev/null
+++ b/src/test/regress/sql/comment_migration_column_joined_tables.sql
@@ -0,0 +1,54 @@
+CREATE SCHEMA comment_migration_column;
+
+CREATE TABLE comment_migration_column.table_1
+(
+    id            bigserial,
+    code          varchar(200)                           not null,
+    name          varchar(200),
+    date_created  timestamp with time zone default NOW() not null,
+    active        boolean                  default true,
+    tenant_id     bigint,
+    CONSTRAINT table_1_pkey PRIMARY KEY (id, tenant_id)
+)
+    WITH (autovacuum_enabled = TRUE);
+
+CREATE TABLE comment_migration_column.table_2
+(
+    id            bigserial,
+    table_1_id    bigint ,
+    description   varchar(200),
+    tenant_id     bigint,
+    CONSTRAINT table_2_pkey PRIMARY KEY (id, tenant_id)
+)
+    WITH (autovacuum_enabled = TRUE);
+
+SET citus.shard_replication_factor = 1;
+
+comment on column comment_migration_column.table_1.id is 'table 1 id';
+comment on column comment_migration_column.table_1.code is 'table 1 code';
+comment on column comment_migration_column.table_1.name is 'table 1 name';
+comment on column comment_migration_column.table_1.date_created is 'table 1 date_created';
+comment on column comment_migration_column.table_1.active is 'table 1 active';
+
+comment on column comment_migration_column.table_2.id is 'table 2 id';
+comment on column comment_migration_column.table_2.table_1_id is 'table 2 foreign key to table 1';
+comment on column comment_migration_column.table_2.description is 'table 2 description';
+
+select col_description('comment_migration_column.table_1'::regclass,1), col_description('comment_migration_column.table_1'::regclass,2), col_description('comment_migration_column.table_1'::regclass,3), col_description('comment_migration_column.table_1'::regclass,4), col_description('comment_migration_column.table_1'::regclass,5);
+select col_description('comment_migration_column.table_2'::regclass,1), col_description('comment_migration_column.table_2'::regclass,2), col_description('comment_migration_column.table_2'::regclass,3);
+
+SELECT create_distributed_table('comment_migration_column.table_1', 'tenant_id');
+SELECT create_distributed_table('comment_migration_column.table_2', 'tenant_id', colocate_with=>'comment_migration_column.table_1');
+
+SELECT undistribute_table('comment_migration_column.table_1');
+SELECT undistribute_table('comment_migration_column.table_2');
+
+SELECT create_distributed_table('comment_migration_column.table_1', 'tenant_id');
+SELECT create_distributed_table('comment_migration_column.table_2', 'tenant_id', colocate_with=>'comment_migration_column.table_1');
+
+select col_description('comment_migration_column.table_1'::regclass,1), col_description('comment_migration_column.table_1'::regclass,2), col_description('comment_migration_column.table_1'::regclass,3), col_description('comment_migration_column.table_1'::regclass,4), col_description('comment_migration_column.table_1'::regclass,5);
+select col_description('comment_migration_column.table_2'::regclass,1), col_description('comment_migration_column.table_2'::regclass,2), col_description('comment_migration_column.table_2'::regclass,3);
+
+DROP TABLE comment_migration_column.table_2;
+DROP TABLE comment_migration_column.table_1;
+DROP SCHEMA comment_migration_column;
\ No newline at end of file
diff --git a/src/test/regress/sql/comment_migration_column_joined_tables_FK.sql b/src/test/regress/sql/comment_migration_column_joined_tables_FK.sql
new file mode 100644
index 00000000000..67e3c5cc415
--- /dev/null
+++ b/src/test/regress/sql/comment_migration_column_joined_tables_FK.sql
@@ -0,0 +1,56 @@
+CREATE SCHEMA comment_migration_column;
+
+CREATE TABLE comment_migration_column.table_1
+(
+    id            bigserial,
+    code          varchar(200)                           not null,
+    name          varchar(200),
+    date_created  timestamp with time zone default NOW() not null,
+    active        boolean                  default true,
+    owner_id     bigint,
+    CONSTRAINT table_1_pkey PRIMARY KEY (id)
+)
+    WITH (autovacuum_enabled = TRUE);
+
+CREATE TABLE comment_migration_column.table_2
+(
+    id            bigserial,
+    table_1_id    bigint ,
+    description   varchar(200),
+    owner_id     bigint,
+    CONSTRAINT table_2_pkey PRIMARY KEY (id)
+)
+    WITH (autovacuum_enabled = TRUE);
+
+SET citus.shard_replication_factor = 1;
+
+comment on column comment_migration_column.table_1.id is 'table 1 id';
+comment on column comment_migration_column.table_1.code is 'table 1 code';
+comment on column comment_migration_column.table_1.name is 'table 1 name';
+comment on column comment_migration_column.table_1.date_created is 'table 1 date_created';
+comment on column comment_migration_column.table_1.active is 'table 1 active';
+
+comment on column comment_migration_column.table_2.id is 'table 2 id';
+comment on column comment_migration_column.table_2.table_1_id is 'table 2 foreign key to table 1';
+comment on column comment_migration_column.table_2.description is 'table 2 description';
+
+select col_description('comment_migration_column.table_1'::regclass,1), col_description('comment_migration_column.table_1'::regclass,2), col_description('comment_migration_column.table_1'::regclass,3), col_description('comment_migration_column.table_1'::regclass,4), col_description('comment_migration_column.table_1'::regclass,5);
+select col_description('comment_migration_column.table_2'::regclass,1), col_description('comment_migration_column.table_2'::regclass,2), col_description('comment_migration_column.table_2'::regclass,3);
+
+SELECT create_reference_table('comment_migration_column.table_1');
+SELECT create_distributed_table('comment_migration_column.table_2', 'id');
+ALTER TABLE comment_migration_column.table_2 ADD CONSTRAINT table2_table1_fk FOREIGN KEY (table_1_id) REFERENCES comment_migration_column.table_1(id);
+
+SELECT undistribute_table('comment_migration_column.table_1', cascade_via_foreign_keys=>true);
+
+ALTER TABLE comment_migration_column.table_2 DROP CONSTRAINT table2_table1_fk;
+SELECT create_reference_table('comment_migration_column.table_1');
+SELECT create_distributed_table('comment_migration_column.table_2', 'id');
+ALTER TABLE comment_migration_column.table_2 ADD CONSTRAINT table2_table1_fk FOREIGN KEY (table_1_id) REFERENCES comment_migration_column.table_1(id);
+
+select col_description('comment_migration_column.table_1'::regclass,1), col_description('comment_migration_column.table_1'::regclass,2), col_description('comment_migration_column.table_1'::regclass,3), col_description('comment_migration_column.table_1'::regclass,4), col_description('comment_migration_column.table_1'::regclass,5);
+select col_description('comment_migration_column.table_2'::regclass,1), col_description('comment_migration_column.table_2'::regclass,2), col_description('comment_migration_column.table_2'::regclass,3);
+
+DROP TABLE comment_migration_column.table_2;
+DROP TABLE comment_migration_column.table_1;
+DROP SCHEMA comment_migration_column;
\ No newline at end of file
diff --git a/src/test/regress/sql/comment_migration_table.sql b/src/test/regress/sql/comment_migration_table.sql
new file mode 100644
index 00000000000..2049e070683
--- /dev/null
+++ b/src/test/regress/sql/comment_migration_table.sql
@@ -0,0 +1,29 @@
+CREATE SCHEMA comment_migration_table;
+
+CREATE TABLE comment_migration_table.table_1
+(
+    id            bigserial,
+    code          varchar(200)                           not null,
+    name          varchar(200),
+    date_created  timestamp with time zone default NOW() not null,
+    active        boolean                  default true,
+    CONSTRAINT table_1_pkey PRIMARY KEY (id)
+)
+    WITH (autovacuum_enabled = TRUE);
+
+SET citus.shard_replication_factor = 1;
+
+comment on table comment_migration_table.table_1 is 'Table 1';
+select obj_description('comment_migration_table.table_1'::regclass); 
+
+SELECT create_distributed_table('comment_migration_table.table_1', 'id');
+select obj_description('comment_migration_table.table_1'::regclass); 
+
+select undistribute_table('comment_migration_table.table_1');
+select obj_description('comment_migration_table.table_1'::regclass); 
+
+SELECT create_distributed_table('comment_migration_table.table_1', 'id');
+select obj_description('comment_migration_table.table_1'::regclass); 
+
+DROP TABLE comment_migration_table.table_1;
+DROP SCHEMA comment_migration_table;
\ No newline at end of file
diff --git a/src/test/regress/sql/comment_migration_table_joined_tables.sql b/src/test/regress/sql/comment_migration_table_joined_tables.sql
new file mode 100644
index 00000000000..63f8988473c
--- /dev/null
+++ b/src/test/regress/sql/comment_migration_table_joined_tables.sql
@@ -0,0 +1,49 @@
+CREATE SCHEMA comment_migration_table;
+
+CREATE TABLE comment_migration_table.table_1
+(
+    id            bigserial,
+    code          varchar(200)                           not null,
+    name          varchar(200),
+    date_created  timestamp with time zone default NOW() not null,
+    active        boolean                  default true,
+    tenant_id     bigint,
+    CONSTRAINT table_1_pkey PRIMARY KEY (id, tenant_id)
+)
+    WITH (autovacuum_enabled = TRUE);
+
+CREATE TABLE comment_migration_table.table_2
+(
+    id            bigserial,
+    table_1_id    bigint ,
+    description   varchar(200),
+    tenant_id     bigint,
+    CONSTRAINT table_2_pkey PRIMARY KEY (id, tenant_id)
+)
+    WITH (autovacuum_enabled = TRUE);
+
+SET citus.shard_replication_factor = 1;
+
+comment on table comment_migration_table.table_1 is 'Table 1';
+comment on table comment_migration_table.table_2 is 'Table 2';
+select obj_description('comment_migration_table.table_1'::regclass);
+select obj_description('comment_migration_table.table_2'::regclass);
+
+SELECT create_distributed_table('comment_migration_table.table_1', 'tenant_id');
+SELECT create_distributed_table('comment_migration_table.table_2', 'tenant_id', colocate_with=>'comment_migration_table.table_1');
+select obj_description('comment_migration_table.table_1'::regclass);
+select obj_description('comment_migration_table.table_2'::regclass);
+
+SELECT undistribute_table('comment_migration_table.table_1');
+SELECT undistribute_table('comment_migration_table.table_2');
+select obj_description('comment_migration_table.table_1'::regclass);
+select obj_description('comment_migration_table.table_2'::regclass);
+
+SELECT create_distributed_table('comment_migration_table.table_1', 'tenant_id');
+SELECT create_distributed_table('comment_migration_table.table_2', 'tenant_id', colocate_with=>'comment_migration_table.table_1');
+select obj_description('comment_migration_table.table_1'::regclass);
+select obj_description('comment_migration_table.table_2'::regclass);
+
+DROP TABLE comment_migration_table.table_2;
+DROP TABLE comment_migration_table.table_1;
+DROP SCHEMA comment_migration_table;
\ No newline at end of file
diff --git a/src/test/regress/sql/comment_migration_table_joined_tables_FK.sql b/src/test/regress/sql/comment_migration_table_joined_tables_FK.sql
new file mode 100644
index 00000000000..05427e94fce
--- /dev/null
+++ b/src/test/regress/sql/comment_migration_table_joined_tables_FK.sql
@@ -0,0 +1,51 @@
+CREATE SCHEMA comment_migration_table;
+
+CREATE TABLE comment_migration_table.table_1
+(
+    id            bigserial,
+    code          varchar(200)                           not null,
+    name          varchar(200),
+    date_created  timestamp with time zone default NOW() not null,
+    active        boolean                  default true,
+    owner_id     bigint,
+    CONSTRAINT table_1_pkey PRIMARY KEY (id)
+)
+    WITH (autovacuum_enabled = TRUE);
+
+CREATE TABLE comment_migration_table.table_2
+(
+    id            bigserial,
+    table_1_id    bigint ,
+    description   varchar(200),
+    owner_id     bigint,
+    CONSTRAINT table_2_pkey PRIMARY KEY (id)
+)
+    WITH (autovacuum_enabled = TRUE);
+
+SET citus.shard_replication_factor = 1;
+
+comment on table comment_migration_table.table_1 is 'Table 1';
+comment on table comment_migration_table.table_2 is 'Table 2';
+select obj_description('comment_migration_table.table_1'::regclass);
+select obj_description('comment_migration_table.table_2'::regclass);
+
+SELECT create_reference_table('comment_migration_table.table_1');
+SELECT create_distributed_table('comment_migration_table.table_2', 'id');
+ALTER TABLE comment_migration_table.table_2 ADD CONSTRAINT table2_table1_fk FOREIGN KEY (table_1_id) REFERENCES comment_migration_table.table_1(id);
+select obj_description('comment_migration_table.table_1'::regclass);
+select obj_description('comment_migration_table.table_2'::regclass);
+
+SELECT undistribute_table('comment_migration_table.table_1', cascade_via_foreign_keys=>true);
+select obj_description('comment_migration_table.table_1'::regclass);
+select obj_description('comment_migration_table.table_2'::regclass);
+
+ALTER TABLE comment_migration_table.table_2 DROP CONSTRAINT table2_table1_fk;
+SELECT create_reference_table('comment_migration_table.table_1');
+SELECT create_distributed_table('comment_migration_table.table_2', 'id');
+ALTER TABLE comment_migration_table.table_2 ADD CONSTRAINT table2_table1_fk FOREIGN KEY (table_1_id) REFERENCES comment_migration_table.table_1(id);
+select obj_description('comment_migration_table.table_1'::regclass);
+select obj_description('comment_migration_table.table_2'::regclass);
+
+DROP TABLE comment_migration_table.table_2;
+DROP TABLE comment_migration_table.table_1;
+DROP SCHEMA comment_migration_table;
\ No newline at end of file

From 45fca5059833dd8307922ae07eb609e1092f34e9 Mon Sep 17 00:00:00 2001
From: EriksP <Eriks.Poikans@zzdats.lv>
Date: Fri, 21 Jun 2024 13:56:43 +0300
Subject: [PATCH 2/2] Code formatting cleanup

---
 .../distributed/commands/alter_table.c        | 24 ++++++++++++-------
 src/test/regress/comment_migration_schedule   |  2 +-
 src/test/regress/multi_schedule               |  2 +-
 .../regress/sql/comment_migration_column.sql  |  2 +-
 ...comment_migration_column_joined_tables.sql |  2 +-
 ...ment_migration_column_joined_tables_FK.sql |  2 +-
 .../regress/sql/comment_migration_table.sql   | 10 ++++----
 .../comment_migration_table_joined_tables.sql |  2 +-
 ...mment_migration_table_joined_tables_FK.sql |  2 +-
 9 files changed, 27 insertions(+), 21 deletions(-)

diff --git a/src/backend/distributed/commands/alter_table.c b/src/backend/distributed/commands/alter_table.c
index c62a02b7197..62764945cb3 100644
--- a/src/backend/distributed/commands/alter_table.c
+++ b/src/backend/distributed/commands/alter_table.c
@@ -34,8 +34,8 @@
 #include "catalog/pg_am.h"
 #include "catalog/pg_depend.h"
 #include "catalog/pg_rewrite_d.h"
-#include "commands/defrem.h"
 #include "commands/comment.h"
+#include "commands/defrem.h"
 #include "executor/spi.h"
 #include "nodes/pg_list.h"
 #include "utils/builtins.h"
@@ -1779,11 +1779,13 @@ CreateMaterializedViewDDLCommand(Oid matViewOid)
 	return query->data;
 }
 
+
 /*
-* MigrateColumnComments migrates distributed table column comments to the target undistributed table columns.
-*/
+ * MigrateColumnComments migrates distributed table column comments to the target undistributed table columns.
+ */
 static void
-MigrateColumnComments(Oid sourceId, Oid targetId){
+MigrateColumnComments(Oid sourceId, Oid targetId)
+{
 	Relation relation = relation_open(sourceId, AccessShareLock);
 	TupleDesc tupleDesc = RelationGetDescr(relation);
 	for (int attrNum = 0; attrNum < tupleDesc->natts; attrNum++)
@@ -1791,24 +1793,28 @@ MigrateColumnComments(Oid sourceId, Oid targetId){
 		Form_pg_attribute attr = TupleDescAttr(tupleDesc, attrNum);
 		if (!attr->attisdropped)
 		{
-			char *columnComment = GetComment(sourceId, RelationRelationId ,attrNum + 1);
+			char *columnComment = GetComment(sourceId, RelationRelationId, attrNum + 1);
 			CreateComments(targetId, RelationRelationId, attrNum + 1, columnComment);
 		}
 	}
 	relation_close(relation, AccessShareLock);
 }
 
+
 /*
  * MigrateTableComment migrates the comment of the source distributed table to the target undistributed table.
-*/
+ */
 static void
-MigrateTableComment(Oid sourceId, Oid targetId){
+MigrateTableComment(Oid sourceId, Oid targetId)
+{
 	char *comment = GetComment(sourceId, RelationRelationId, 0);
-	if(comment != NULL) {
+	if (comment != NULL)
+	{
 		CreateComments(targetId, RelationRelationId, 0, comment);
 	}
 }
 
+
 /*
  * ReplaceTable replaces the source table with the target table.
  * It moves all the rows of the source table to target table with INSERT SELECT.
@@ -1828,7 +1834,6 @@ ReplaceTable(Oid sourceId, Oid targetId, List *justBeforeDropCommands,
 
 	if (!PartitionedTable(sourceId) && !IsForeignTable(sourceId))
 	{
-		
 		if (!suppressNoticeMessages)
 		{
 			ereport(NOTICE, (errmsg("moving the data of %s", qualifiedSourceName)));
@@ -1933,6 +1938,7 @@ ReplaceTable(Oid sourceId, Oid targetId, List *justBeforeDropCommands,
 	ExecuteQueryViaSPI(query->data, SPI_OK_UTILITY);
 }
 
+
 /*
  * HasAnyGeneratedStoredColumns decides if relation has any columns that we
  * might need to copy the data of when replacing table.
diff --git a/src/test/regress/comment_migration_schedule b/src/test/regress/comment_migration_schedule
index 2c11f1458b3..50146af0e5b 100644
--- a/src/test/regress/comment_migration_schedule
+++ b/src/test/regress/comment_migration_schedule
@@ -4,4 +4,4 @@ test: comment_migration_table_joined_tables
 test: comment_migration_table_joined_tables_FK
 test: comment_migration_column
 test: comment_migration_column_joined_tables
-test: comment_migration_column_joined_tables_FK
\ No newline at end of file
+test: comment_migration_column_joined_tables_FK
diff --git a/src/test/regress/multi_schedule b/src/test/regress/multi_schedule
index 4ed8e135ada..bb0c53e098e 100644
--- a/src/test/regress/multi_schedule
+++ b/src/test/regress/multi_schedule
@@ -144,4 +144,4 @@ test: comment_migration_table_joined_tables
 test: comment_migration_table_joined_tables_FK
 test: comment_migration_column
 test: comment_migration_column_joined_tables
-test: comment_migration_column_joined_tables_FK
\ No newline at end of file
+test: comment_migration_column_joined_tables_FK
diff --git a/src/test/regress/sql/comment_migration_column.sql b/src/test/regress/sql/comment_migration_column.sql
index 3f9753894ab..f5f5b453cbd 100644
--- a/src/test/regress/sql/comment_migration_column.sql
+++ b/src/test/regress/sql/comment_migration_column.sql
@@ -30,4 +30,4 @@ SELECT create_distributed_table('comment_migration_column.table_1', 'id');
 select col_description('comment_migration_column.table_1'::regclass,1), col_description('comment_migration_column.table_1'::regclass,2), col_description('comment_migration_column.table_1'::regclass,3), col_description('comment_migration_column.table_1'::regclass,4), col_description('comment_migration_column.table_1'::regclass,5);
 
 DROP TABLE comment_migration_column.table_1;
-DROP SCHEMA comment_migration_column;
\ No newline at end of file
+DROP SCHEMA comment_migration_column;
diff --git a/src/test/regress/sql/comment_migration_column_joined_tables.sql b/src/test/regress/sql/comment_migration_column_joined_tables.sql
index f3220e9b23b..10e0b1cca25 100644
--- a/src/test/regress/sql/comment_migration_column_joined_tables.sql
+++ b/src/test/regress/sql/comment_migration_column_joined_tables.sql
@@ -51,4 +51,4 @@ select col_description('comment_migration_column.table_2'::regclass,1), col_desc
 
 DROP TABLE comment_migration_column.table_2;
 DROP TABLE comment_migration_column.table_1;
-DROP SCHEMA comment_migration_column;
\ No newline at end of file
+DROP SCHEMA comment_migration_column;
diff --git a/src/test/regress/sql/comment_migration_column_joined_tables_FK.sql b/src/test/regress/sql/comment_migration_column_joined_tables_FK.sql
index 67e3c5cc415..3fcc8c8f205 100644
--- a/src/test/regress/sql/comment_migration_column_joined_tables_FK.sql
+++ b/src/test/regress/sql/comment_migration_column_joined_tables_FK.sql
@@ -53,4 +53,4 @@ select col_description('comment_migration_column.table_2'::regclass,1), col_desc
 
 DROP TABLE comment_migration_column.table_2;
 DROP TABLE comment_migration_column.table_1;
-DROP SCHEMA comment_migration_column;
\ No newline at end of file
+DROP SCHEMA comment_migration_column;
diff --git a/src/test/regress/sql/comment_migration_table.sql b/src/test/regress/sql/comment_migration_table.sql
index 2049e070683..638bddba98d 100644
--- a/src/test/regress/sql/comment_migration_table.sql
+++ b/src/test/regress/sql/comment_migration_table.sql
@@ -14,16 +14,16 @@ CREATE TABLE comment_migration_table.table_1
 SET citus.shard_replication_factor = 1;
 
 comment on table comment_migration_table.table_1 is 'Table 1';
-select obj_description('comment_migration_table.table_1'::regclass); 
+select obj_description('comment_migration_table.table_1'::regclass);
 
 SELECT create_distributed_table('comment_migration_table.table_1', 'id');
-select obj_description('comment_migration_table.table_1'::regclass); 
+select obj_description('comment_migration_table.table_1'::regclass);
 
 select undistribute_table('comment_migration_table.table_1');
-select obj_description('comment_migration_table.table_1'::regclass); 
+select obj_description('comment_migration_table.table_1'::regclass);
 
 SELECT create_distributed_table('comment_migration_table.table_1', 'id');
-select obj_description('comment_migration_table.table_1'::regclass); 
+select obj_description('comment_migration_table.table_1'::regclass);
 
 DROP TABLE comment_migration_table.table_1;
-DROP SCHEMA comment_migration_table;
\ No newline at end of file
+DROP SCHEMA comment_migration_table;
diff --git a/src/test/regress/sql/comment_migration_table_joined_tables.sql b/src/test/regress/sql/comment_migration_table_joined_tables.sql
index 63f8988473c..4c7f962e5d9 100644
--- a/src/test/regress/sql/comment_migration_table_joined_tables.sql
+++ b/src/test/regress/sql/comment_migration_table_joined_tables.sql
@@ -46,4 +46,4 @@ select obj_description('comment_migration_table.table_2'::regclass);
 
 DROP TABLE comment_migration_table.table_2;
 DROP TABLE comment_migration_table.table_1;
-DROP SCHEMA comment_migration_table;
\ No newline at end of file
+DROP SCHEMA comment_migration_table;
diff --git a/src/test/regress/sql/comment_migration_table_joined_tables_FK.sql b/src/test/regress/sql/comment_migration_table_joined_tables_FK.sql
index 05427e94fce..093177f93c8 100644
--- a/src/test/regress/sql/comment_migration_table_joined_tables_FK.sql
+++ b/src/test/regress/sql/comment_migration_table_joined_tables_FK.sql
@@ -48,4 +48,4 @@ select obj_description('comment_migration_table.table_2'::regclass);
 
 DROP TABLE comment_migration_table.table_2;
 DROP TABLE comment_migration_table.table_1;
-DROP SCHEMA comment_migration_table;
\ No newline at end of file
+DROP SCHEMA comment_migration_table;