3 files changed, 134 insertions, 5 deletions

diff --git a/src/backend/catalog/pg_constraint.c b/src/backend/catalog/pg_constraint.c
index 1336c46d3f5..104e930a6fa 100644
--- a/src/backend/catalog/pg_constraint.c
+++ b/src/backend/catalog/pg_constraint.c
@@ -806,6 +806,104 @@ get_relation_constraint_oid(Oid relid, const char *conname, bool missing_ok)
 }
 
 /*
+ * get_relation_constraint_attnos
+ *		Find a constraint on the specified relation with the specified name
+ *		and return the constrained columns.
+ *
+ * Returns a Bitmapset of the column attnos of the constrained columns, with
+ * attnos being offset by FirstLowInvalidHeapAttributeNumber so that system
+ * columns can be represented.
+ *
+ * *constraintOid is set to the OID of the constraint, or InvalidOid on
+ * failure.
+ */
+Bitmapset *
+get_relation_constraint_attnos(Oid relid, const char *conname,
+							   bool missing_ok, Oid *constraintOid)
+{
+	Bitmapset  *conattnos = NULL;
+	Relation	pg_constraint;
+	HeapTuple	tuple;
+	SysScanDesc scan;
+	ScanKeyData skey[1];
+
+	/* Set *constraintOid, to avoid complaints about uninitialized vars */
+	*constraintOid = InvalidOid;
+
+	/*
+	 * Fetch the constraint tuple from pg_constraint.  There may be more than
+	 * one match, because constraints are not required to have unique names;
+	 * if so, error out.
+	 */
+	pg_constraint = heap_open(ConstraintRelationId, AccessShareLock);
+
+	ScanKeyInit(&skey[0],
+				Anum_pg_constraint_conrelid,
+				BTEqualStrategyNumber, F_OIDEQ,
+				ObjectIdGetDatum(relid));
+
+	scan = systable_beginscan(pg_constraint, ConstraintRelidIndexId, true,
+							  NULL, 1, skey);
+
+	while (HeapTupleIsValid(tuple = systable_getnext(scan)))
+	{
+		Form_pg_constraint con = (Form_pg_constraint) GETSTRUCT(tuple);
+		Datum		adatum;
+		bool		isNull;
+		ArrayType  *arr;
+		int16	   *attnums;
+		int			numcols;
+		int			i;
+
+		/* Check the constraint name */
+		if (strcmp(NameStr(con->conname), conname) != 0)
+			continue;
+		if (OidIsValid(*constraintOid))
+			ereport(ERROR,
+					(errcode(ERRCODE_DUPLICATE_OBJECT),
+					 errmsg("table \"%s\" has multiple constraints named \"%s\"",
+							get_rel_name(relid), conname)));
+
+		*constraintOid = HeapTupleGetOid(tuple);
+
+		/* Extract the conkey array, ie, attnums of constrained columns */
+		adatum = heap_getattr(tuple, Anum_pg_constraint_conkey,
+							  RelationGetDescr(pg_constraint), &isNull);
+		if (isNull)
+			continue;			/* no constrained columns */
+
+		arr = DatumGetArrayTypeP(adatum);	/* ensure not toasted */
+		numcols = ARR_DIMS(arr)[0];
+		if (ARR_NDIM(arr) != 1 ||
+			numcols < 0 ||
+			ARR_HASNULL(arr) ||
+			ARR_ELEMTYPE(arr) != INT2OID)
+			elog(ERROR, "conkey is not a 1-D smallint array");
+		attnums = (int16 *) ARR_DATA_PTR(arr);
+
+		/* Construct the result value */
+		for (i = 0; i < numcols; i++)
+		{
+			conattnos = bms_add_member(conattnos,
+									   attnums[i] - FirstLowInvalidHeapAttributeNumber);
+		}
+	}
+
+	systable_endscan(scan);
+
+	/* If no such constraint exists, complain */
+	if (!OidIsValid(*constraintOid) && !missing_ok)
+		ereport(ERROR,
+				(errcode(ERRCODE_UNDEFINED_OBJECT),
+				 errmsg("constraint \"%s\" for table \"%s\" does not exist",
+						conname, get_rel_name(relid))));
+
+	heap_close(pg_constraint, AccessShareLock);
+
+	return conattnos;
+}
+
+/*
  * get_domain_constraint_oid
  *		Find a constraint on the specified domain with the specified name.
  *		Returns constraint's OID.
diff --git a/src/backend/parser/parse_clause.c b/src/backend/parser/parse_clause.c
index af99e65aa7d..ed26517c266 100644
--- a/src/backend/parser/parse_clause.c
+++ b/src/backend/parser/parse_clause.c
@@ -3164,9 +3164,26 @@ transformOnConflictArbiter(ParseState *pstate,
 
 		pstate->p_namespace = save_namespace;
 
+		/*
+		 * If the arbiter is specified by constraint name, get the constraint
+		 * OID and mark the constrained columns as requiring SELECT privilege,
+		 * in the same way as would have happened if the arbiter had been
+		 * specified by explicit reference to the constraint's index columns.
+		 */
 		if (infer->conname)
-			*constraint = get_relation_constraint_oid(RelationGetRelid(pstate->p_target_relation),
-													  infer->conname, false);
+		{
+			Oid			relid = RelationGetRelid(pstate->p_target_relation);
+			RangeTblEntry *rte = pstate->p_target_rangetblentry;
+			Bitmapset  *conattnos;
+
+			conattnos = get_relation_constraint_attnos(relid, infer->conname,
+													   false, constraint);
+
+			/* Make sure the rel as a whole is marked for SELECT access */
+			rte->requiredPerms |= ACL_SELECT;
+			/* Mark the constrained columns as requiring SELECT access */
+			rte->selectedCols = bms_add_members(rte->selectedCols, conattnos);
+		}
 	}
 
 	/*
diff --git a/src/backend/rewrite/rowsecurity.c b/src/backend/rewrite/rowsecurity.c
index 9e7e54db67b..a0cd6b1075d 100644
--- a/src/backend/rewrite/rowsecurity.c
+++ b/src/backend/rewrite/rowsecurity.c
@@ -310,6 +310,8 @@ get_row_security_policies(Query *root, RangeTblEntry *rte, int rt_index,
 		{
 			List	   *conflict_permissive_policies;
 			List	   *conflict_restrictive_policies;
+			List	   *conflict_select_permissive_policies = NIL;
+			List	   *conflict_select_restrictive_policies = NIL;
 
 			/* Get the policies that apply to the auxiliary UPDATE */
 			get_policies_for_relation(rel, CMD_UPDATE, user_id,
@@ -339,9 +341,6 @@ get_row_security_policies(Query *root, RangeTblEntry *rte, int rt_index,
 			 */
 			if (rte->requiredPerms & ACL_SELECT)
 			{
-				List	   *conflict_select_permissive_policies = NIL;
-				List	   *conflict_select_restrictive_policies = NIL;
-
 				get_policies_for_relation(rel, CMD_SELECT, user_id,
 										  &conflict_select_permissive_policies,
 										  &conflict_select_restrictive_policies);
@@ -362,6 +361,21 @@ get_row_security_policies(Query *root, RangeTblEntry *rte, int rt_index,
 								   withCheckOptions,
 								   hasSubLinks,
 								   false);
+
+			/*
+			 * Add ALL/SELECT policies as WCO_RLS_UPDATE_CHECK WCOs, to ensure
+			 * that the final updated row is visible when taking the UPDATE
+			 * path of an INSERT .. ON CONFLICT DO UPDATE, if SELECT rights
+			 * are required for this relation.
+			 */
+			if (rte->requiredPerms & ACL_SELECT)
+				add_with_check_options(rel, rt_index,
+									   WCO_RLS_UPDATE_CHECK,
+									   conflict_select_permissive_policies,
+									   conflict_select_restrictive_policies,
+									   withCheckOptions,
+									   hasSubLinks,
+									   true);
 		}
 	}