Cherry-pick Bug 3697586 fix

src/backend/distributed/transaction/backend_data.c

@@ -855,6 +855,16 @@
 void
 AssignDistributedTransactionId(void)
 {
+	/*
+	 * MyBackendData should always be available. However, we observed some
+	 * crashes where certain hooks were not executed.
+	 * Bug 3697586: Server crashes when assigning distributed transaction
+	 */
+	if (!MyBackendData)
+	{
+		ereport(ERROR, (errmsg("backend is not ready for distributed transactions")));
+	}
+
 	pg_atomic_uint64 *transactionNumberSequence =
 		&backendManagementShmemData->nextTransactionNumber;
 

src/test/regress/citus_tests/run_test.py

@@ -181,6 +181,9 @@ def extra_tests(self):
             "multi_test_helpers_superuser",
         ],
     ),
+    "foreign_key_to_reference_shard_rebalance": TestDeps(
+        "minimal_schedule", ["remove_coordinator_from_metadata"]
+    ),
 }
 
 

src/test/regress/expected/foreign_key_to_reference_shard_rebalance.out

@@ -210,7 +210,7 @@ select create_distributed_table('partitioned_tbl_with_fkey','x');
 
 create table partition_1_with_fkey partition of partitioned_tbl_with_fkey for values from ('2022-01-01') to ('2022-12-31');
 create table partition_2_with_fkey partition of partitioned_tbl_with_fkey for values from ('2023-01-01') to ('2023-12-31');
-create table partition_3_with_fkey partition of partitioned_tbl_with_fkey for values from ('2024-01-01') to ('2024-12-31');
+create table partition_3_with_fkey partition of partitioned_tbl_with_fkey DEFAULT;
 insert into partitioned_tbl_with_fkey (x,y) select s,s%10 from generate_series(1,100) s;
 ALTER TABLE partitioned_tbl_with_fkey ADD CONSTRAINT fkey_to_ref_tbl FOREIGN KEY (y) REFERENCES ref_table_with_fkey(id);
 WITH shardid AS (SELECT shardid FROM pg_dist_shard where logicalrelid = 'partitioned_tbl_with_fkey'::regclass ORDER BY shardid LIMIT 1)

src/test/regress/sql/foreign_key_to_reference_shard_rebalance.sql

@@ -84,7 +84,7 @@ create table partitioned_tbl_with_fkey (x int, y int, t timestamptz default now(
 select create_distributed_table('partitioned_tbl_with_fkey','x');
 create table partition_1_with_fkey partition of partitioned_tbl_with_fkey for values from ('2022-01-01') to ('2022-12-31');
 create table partition_2_with_fkey partition of partitioned_tbl_with_fkey for values from ('2023-01-01') to ('2023-12-31');
-create table partition_3_with_fkey partition of partitioned_tbl_with_fkey for values from ('2024-01-01') to ('2024-12-31');
+create table partition_3_with_fkey partition of partitioned_tbl_with_fkey DEFAULT;
 insert into partitioned_tbl_with_fkey (x,y) select s,s%10 from generate_series(1,100) s;
 
 ALTER TABLE partitioned_tbl_with_fkey ADD CONSTRAINT fkey_to_ref_tbl FOREIGN KEY (y) REFERENCES ref_table_with_fkey(id);