diff --git a/src/backend/commands/tablecmds.c b/src/backend/commands/tablecmds.c
index 2fa534413ea..af3a2c95c6b 100644
--- a/src/backend/commands/tablecmds.c
+++ b/src/backend/commands/tablecmds.c
@@ -15561,6 +15561,84 @@ ATExecAlterColumnType(AlteredTableInfo *tab, Relation rel,
 	return address;
 }
 
+static Oid
+get_parent_constraint(HeapTuple constraintTup)
+{
+	AttrMap    *attmap;
+	Relation	pg_constraints;
+	ScanKeyData key[3];
+	SysScanDesc scan;
+	HeapTuple	tuple;
+	Relation	parentRel;
+	Relation	conRel;
+	AttrNumber	child_attno;
+	Oid			result = InvalidOid;
+	Oid			parentRelid = InvalidOid;
+	Form_pg_constraint child_con = GETSTRUCT(constraintTup);
+
+	child_attno = child_con->contype == CONSTRAINT_NOTNULL ? extractNotNullColumn(constraintTup) : 0;
+
+	parentRelid = get_partition_parent(child_con->conrelid, false);
+
+	/* scan the parent constraints, matching by name for CHECK and by attno for NN */
+	pg_constraints = table_open(ConstraintRelationId, AccessShareLock);
+	ScanKeyInit(&key[0], Anum_pg_constraint_conrelid,
+				BTEqualStrategyNumber, F_OIDEQ, ObjectIdGetDatum(parentRelid));
+	ScanKeyInit(&key[1],
+				Anum_pg_constraint_contypid,
+				BTEqualStrategyNumber, F_OIDEQ,
+				ObjectIdGetDatum(InvalidOid));
+	if (child_con->contype == CONSTRAINT_CHECK)
+		ScanKeyInit(&key[2],
+					Anum_pg_constraint_conname,
+					BTEqualStrategyNumber, F_NAMEEQ,
+					CStringGetDatum(NameStr(child_con->conname)));
+
+	scan = systable_beginscan(pg_constraints,
+							  ConstraintRelidTypidNameIndexId, true, NULL,
+							  child_con->contype == CONSTRAINT_CHECK ? 3 : 2, key);
+
+	while (HeapTupleIsValid(tuple = systable_getnext(scan)))
+	{
+		Form_pg_constraint con = GETSTRUCT(tuple);
+
+		if (con->contype != child_con->contype)
+			continue;
+
+		if (con->connoinherit)
+			continue;
+
+		if (con->contype == CONSTRAINT_CHECK)
+		{
+			/* Got a match by name */
+			result = con->oid;
+			break;
+		}
+		if (con->contype == CONSTRAINT_NOTNULL)
+		{
+			AttrNumber	parent_attno = extractNotNullColumn(tuple);
+
+			parentRel = RelationIdGetRelation(parentRelid);
+			conRel = RelationIdGetRelation(child_con->conrelid);
+			attmap = build_attrmap_by_name(RelationGetDescr(parentRel),
+										   RelationGetDescr(conRel),
+										   true);
+			RelationClose(conRel);
+			RelationClose(parentRel);
+
+			/* Got a match by attno */
+			if (parent_attno == attmap->attnums[child_attno - 1])
+				result = con->oid;
+			break;
+		}
+	}
+
+	systable_endscan(scan);
+	table_close(pg_constraints, AccessShareLock);
+
+	return result;
+}
+
 /*
  * Subroutine for ATExecAlterColumnType and ATExecSetExpression: Find everything
  * that depends on the column (constraints, indexes, etc), and record enough
@@ -15634,9 +15712,30 @@ RememberAllDependentForRebuilding(AlteredTableInfo *tab, AlterTableType subtype,
 				}
 
 			case ConstraintRelationId:
-				Assert(foundObject.objectSubId == 0);
-				RememberConstraintForRebuilding(foundObject.objectId, tab);
-				break;
+				{
+					Oid			parentConstraintId = InvalidOid;
+					HeapTuple	tuple;
+					Form_pg_constraint con;
+
+					Assert(foundObject.objectSubId == 0);
+
+					tuple = SearchSysCache1(CONSTROID, ObjectIdGetDatum(foundObject.objectId));
+					if (!HeapTupleIsValid(tuple))
+						elog(ERROR, "cache lookup failed for constraint %u",
+							 foundObject.objectId);
+
+					con = (Form_pg_constraint) GETSTRUCT(tuple);
+
+					if (!con->conislocal && OidIsValid(parentConstraintId = get_parent_constraint(tuple)))
+					{
+						RememberConstraintForRebuilding(parentConstraintId, tab);
+					}
+
+					ReleaseSysCache(tuple);
+
+					RememberConstraintForRebuilding(foundObject.objectId, tab);
+					break;
+				}
 
 			case ProcedureRelationId:
 
diff --git a/src/test/regress/expected/alter_table.out b/src/test/regress/expected/alter_table.out
index e167a41ce79..60f5302d02e 100644
--- a/src/test/regress/expected/alter_table.out
+++ b/src/test/regress/expected/alter_table.out
@@ -4876,3 +4876,20 @@ drop publication pub1;
 drop schema alter1 cascade;
 drop schema alter2 cascade;
 NOTICE:  drop cascades to table alter2.t1
+-- Check that non-recursive SET EXPRESSION restores the constraints correctly
+create table tab (a int, b int not null generated always as (a+1) stored) partition by list (a);
+create table part partition of tab for values in (1);
+select conname, contype from pg_constraint where conrelid = 'part'::regclass;
+    conname     | contype 
+----------------+---------
+ tab_b_not_null | n
+(1 row)
+
+alter table part alter b set expression as (a);
+select conname, contype from pg_constraint where conrelid = 'part'::regclass;
+    conname     | contype 
+----------------+---------
+ tab_b_not_null | n
+(1 row)
+
+drop table tab;
diff --git a/src/test/regress/sql/alter_table.sql b/src/test/regress/sql/alter_table.sql
index f5f13bbd3e7..da37d591d73 100644
--- a/src/test/regress/sql/alter_table.sql
+++ b/src/test/regress/sql/alter_table.sql
@@ -3159,3 +3159,11 @@ alter table alter1.t1 set schema alter2;
 drop publication pub1;
 drop schema alter1 cascade;
 drop schema alter2 cascade;
+
+-- Check that non-recursive SET EXPRESSION restores the constraints correctly
+create table tab (a int, b int not null generated always as (a+1) stored) partition by list (a);
+create table part partition of tab for values in (1);
+select conname, contype from pg_constraint where conrelid = 'part'::regclass;
+alter table part alter b set expression as (a);
+select conname, contype from pg_constraint where conrelid = 'part'::regclass;
+drop table tab;