[CF 6178] v1 - COPY FROM with RLS

This branch was automatically generated by a robot using patches from an
email thread registered at:

https://commitfest.postgresql.org/patch/6178

The branch will be overwritten each time a new patch version is posted to
the thread, and also periodically to check for bitrot caused by changes
on the master branch.

Patch(es): https://www.postgresql.org/message-id/CACJufxFbmnoa5O-vL43DPTCGt6oagY4dXgKxy=rcD9-e9g0zEg@mail.gmail.com
Author(s): Jian He

Author: Commitfest Bot

doc/src/sgml/ref/copy.sgml

@@ -569,9 +569,11 @@ COPY <replaceable class="parameter">count</replaceable>
     If row-level security is enabled for the table, the relevant
     <command>SELECT</command> policies will apply to <literal>COPY
     <replaceable class="parameter">table</replaceable> TO</literal> statements.
-    Currently, <command>COPY FROM</command> is not supported for tables
-    with row-level security. Use equivalent <command>INSERT</command>
-    statements instead.
+    <command>COPY FROM</command> is supported for tables with row-level security.
+    However if any row-level security policy’s <literal>USING</literal> or
+    <literal>WITH CHECK</literal> expression contains a subquery, then
+    <command>COPY FROM</command> is not supported. In that case, Use equivalent
+    <command>INSERT</command> statements instead.
    </para>
 
    <para>

src/backend/commands/copy.c

@@ -218,20 +218,14 @@ DoCopy(ParseState *pstate, const CopyStmt *stmt,
 		 * If RLS is not enabled for this, then just fall through to the
 		 * normal non-filtering relation handling.
 		 */
-		if (check_enable_rls(relid, InvalidOid, false) == RLS_ENABLED)
+		if (!is_from && check_enable_rls(relid, InvalidOid, false) == RLS_ENABLED)
 		{
 			SelectStmt *select;
 			ColumnRef  *cr;
 			ResTarget  *target;
 			RangeVar   *from;
 			List	   *targetList = NIL;
 
-			if (is_from)
-				ereport(ERROR,
-						(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
-						 errmsg("COPY FROM not supported with row-level security"),
-						 errhint("Use INSERT statements instead.")));
-
 			/*
 			 * Build target list
 			 *

src/backend/commands/copyfrom.c

@@ -39,16 +39,20 @@
 #include "foreign/fdwapi.h"
 #include "mb/pg_wchar.h"
 #include "miscadmin.h"
+#include "nodes/makefuncs.h"
 #include "nodes/miscnodes.h"
 #include "optimizer/optimizer.h"
+#include "parser/parse_relation.h"
 #include "pgstat.h"
 #include "rewrite/rewriteHandler.h"
+#include "rewrite/rowsecurity.h"
 #include "storage/fd.h"
 #include "tcop/tcopprot.h"
 #include "utils/lsyscache.h"
 #include "utils/memutils.h"
 #include "utils/portal.h"
 #include "utils/rel.h"
+#include "utils/rls.h"
 #include "utils/snapmgr.h"
 
 /*
@@ -957,6 +961,109 @@ CopyFrom(CopyFromState cstate)
 
 	Assert(resultRelInfo->ri_BatchSize >= 1);
 
+	if (check_enable_rls(RelationGetRelid(cstate->rel), InvalidOid, false) == RLS_ENABLED)
+	{
+		List	*securityQuals = NIL;
+		List	*withCheckOptions = NIL;
+		List	*newWithCheckOptions = NIL;
+		List	*wcoExprs = NIL;
+		char	*refname;
+
+		Query	*root	= makeNode(Query);
+		ModifyTable *node = makeNode(ModifyTable);
+		RangeTblEntry *rte = makeNode(RangeTblEntry);
+		bool		hasRowSecurity = false;
+		bool		hasSubLinks = false;
+		RTEPermissionInfo *perminfo;
+
+		/*
+		 * We use the Query and RTE nodes to retrieve the COPY FROM relation's
+		 * security policies (get_row_security_policies) and transform them into
+		 * WithCheckOption nodes.  Later, we initialize these WCO node exprstate
+		 * and pass these initialized WCOs to the resultRelInfo.
+		 */
+		rte->alias = NULL;
+		refname = RelationGetRelationName(cstate->rel);
+		rte->eref = makeAlias(refname, NIL);
+		rte->rtekind = RTE_RELATION;
+		rte->relid = RelationGetRelid(cstate->rel);
+		rte->inh = false;
+		rte->relkind = cstate->rel->rd_rel->relkind;
+		rte->rellockmode = RowExclusiveLock;
+		rte->lateral = false;
+		rte->inFromCl = false;
+
+		perminfo = addRTEPermissionInfo(&root->rteperminfos, rte);
+		perminfo->requiredPerms = ACL_INSERT;
+		perminfo->checkAsUser = InvalidOid;
+
+		/* we already did the permission check on DoCopy */
+		foreach_int(cur, cstate->attnumlist)
+		{
+			int			attno;
+			Bitmapset **bms;
+
+			attno = cur - FirstLowInvalidHeapAttributeNumber;
+			bms =  &perminfo->insertedCols;
+
+			*bms = bms_add_member(*bms, attno);
+		}
+
+		root->resultRelation = 1;
+		root->rtable = list_make1(rte);
+		root->commandType = CMD_INSERT;
+		root->stmt_location = -1;
+
+		get_row_security_policies(root, rte, 1,
+								  &securityQuals,
+								  &withCheckOptions,
+								  &hasRowSecurity,
+								  &hasSubLinks);
+
+		/* policy contain subquery, maybe doable? */
+		if (hasSubLinks)
+			ereport(ERROR,
+					errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+					errmsg("COPY FROM not supported with row-level security policy contain subquery"),
+					errhint("Use INSERT statements instead."));
+
+		foreach_node(WithCheckOption, wco, withCheckOptions)
+		{
+			wco->qual = eval_const_expressions(NULL, wco->qual);
+			wco->qual = (Node *) canonicalize_qual((Expr *) wco->qual, false);
+			wco->qual = (Node *) make_ands_implicit((Expr *) wco->qual);
+
+			if (wco->qual != NULL)
+				newWithCheckOptions = lappend(newWithCheckOptions, wco);
+		}
+
+		foreach_node(WithCheckOption, wco, newWithCheckOptions)
+		{
+			ExprState  *wcoExpr = ExecInitQual(castNode(List, wco->qual),
+											   &mtstate->ps);
+
+			wcoExprs = lappend(wcoExprs, wcoExpr);
+		}
+
+		resultRelInfo->ri_WithCheckOptions = newWithCheckOptions;
+		resultRelInfo->ri_WithCheckOptionExprs = wcoExprs;
+
+		/* see make_modifytable */
+		node->operation = CMD_INSERT;
+
+		/*
+		 * INSERT applies to a single relation only, so rootRelation is always 0
+		 */
+		node->rootRelation = 0;
+		node->returningOldAlias = NULL;
+		node->returningNewAlias = NULL;
+		node->resultRelations = list_make1_int(1);
+		node->onConflictAction = ONCONFLICT_NONE;
+		node->withCheckOptionLists = list_make1(list_copy(newWithCheckOptions));
+
+		mtstate->ps.plan = (Plan *) node;
+	}
+
 	/* Prepare to catch AFTER triggers. */
 	AfterTriggerBeginQuery();
 
@@ -1349,6 +1456,13 @@ CopyFrom(CopyFromState cstate)
 					ExecComputeStoredGenerated(resultRelInfo, estate, myslot,
 											   CMD_INSERT);
 
+				/* do row level security policy check */
+				if (resultRelInfo->ri_WithCheckOptions != NIL)
+					ExecWithCheckOptions(WCO_RLS_INSERT_CHECK,
+										 resultRelInfo,
+										 myslot,
+										 estate);
+
 				/*
 				 * If the target is a plain table, check the constraints of
 				 * the tuple.

src/test/regress/expected/rowsecurity.out

@@ -615,6 +615,10 @@ EXPLAIN (COSTS OFF) SELECT * FROM document NATURAL JOIN category WHERE f_leak(dt
 -- back from p1r for this because it sorts first
 INSERT INTO document VALUES (100, 44, 1, 'regress_rls_dave', 'testing sorting of policies'); -- fail
 ERROR:  new row violates row-level security policy "p1r" for table "document"
+-- fail, COPY FROM, security policy with subquery not supported
+COPY document FROM STDIN;
+ERROR:  COPY FROM not supported with row-level security policy contain subquery
+HINT:  Use INSERT statements instead.
 -- Just to see a p2r error
 INSERT INTO document VALUES (100, 55, 1, 'regress_rls_dave', 'testing sorting of policies'); -- fail
 ERROR:  new row violates row-level security policy "p2r" for table "document"
@@ -1351,6 +1355,10 @@ EXPLAIN (COSTS OFF) SELECT * FROM part_document WHERE f_leak(dtitle);
 -- pp1 ERROR
 INSERT INTO part_document VALUES (100, 11, 5, 'regress_rls_dave', 'testing pp1'); -- fail
 ERROR:  new row violates row-level security policy for table "part_document"
+-- fail, COPY FROM, security policy with subquery not supported
+COPY part_document FROM STDIN;
+ERROR:  COPY FROM not supported with row-level security policy contain subquery
+HINT:  Use INSERT statements instead.
 -- pp1r ERROR
 INSERT INTO part_document VALUES (100, 99, 1, 'regress_rls_dave', 'testing pp1r'); -- fail
 ERROR:  new row violates row-level security policy "pp1r" for table "part_document"
@@ -1635,6 +1643,49 @@ EXPLAIN (COSTS OFF) SELECT * FROM dependent; -- After drop, should be unqualifie
  Seq Scan on dependent
 (1 row)
 
+--COPY FROM with RLS
+RESET SESSION AUTHORIZATION;
+CREATE TABLE pp (id int,val int) PARTITION BY RANGE (id);
+CREATE TABLE pp_1 (val int, id int);
+ALTER TABLE pp ATTACH PARTITION pp_1 FOR VALUES FROM (1) TO (10);
+CREATE TABLE pp_2 PARTITION OF pp FOR VALUES FROM (10) TO (20);
+ALTER TABLE pp ENABLE ROW LEVEL SECURITY;
+ALTER TABLE pp_1 ENABLE ROW LEVEL SECURITY;
+ALTER TABLE pp_2 ENABLE ROW LEVEL SECURITY;
+CREATE POLICY pp_1_p1 ON pp_1 FOR INSERT WITH CHECK (id = 6);
+CREATE POLICY p2_pp ON pp FOR ALL USING(id = 1 or id = 2);
+GRANT SELECT, INSERT ON pp TO regress_rls_alice;
+GRANT SELECT, INSERT ON pp_1 TO regress_rls_alice;
+GRANT SELECT, INSERT ON pp_2 TO regress_rls_alice;
+SET SESSION AUTHORIZATION regress_rls_alice;
+INSERT INTO pp_1 VALUES (13, 2); --error
+ERROR:  new row violates row-level security policy for table "pp_1"
+INSERT INTO pp_1 VALUES (16, 6); --ok
+COPY pp_1 FROM STDIN WITH DELIMITER ','; --second record not ok
+ERROR:  new row violates row-level security policy for table "pp_1"
+CONTEXT:  COPY pp_1, line 2: "13,2"
+INSERT INTO pp VALUES (1,11), (2,12);
+INSERT INTO pp values (5,11); --error
+ERROR:  new row violates row-level security policy for table "pp"
+INSERT INTO pp values (6,11); --error
+ERROR:  new row violates row-level security policy for table "pp"
+COPY pp FROM STDIN WITH DELIMITER ','; --error, second record not ok
+ERROR:  new row violates row-level security policy for table "pp"
+CONTEXT:  COPY pp, line 2: "5,11"
+COPY pp FROM STDIN WITH DELIMITER ','; --error, second record not ok
+ERROR:  new row violates row-level security policy for table "pp"
+CONTEXT:  COPY pp, line 2: "6,11"
+RESET SESSION AUTHORIZATION;
+CREATE POLICY p1_pp ON pp FOR INSERT WITH CHECK(id > 4);
+SET SESSION AUTHORIZATION regress_rls_alice;
+INSERT INTO pp VALUES (5, 15), (6, 16); --ok
+INSERT INTO pp VALUES (4, 14); --error
+ERROR:  new row violates row-level security policy for table "pp"
+COPY pp FROM STDIN WITH DELIMITER ','; --third record will result error
+ERROR:  new row violates row-level security policy for table "pp"
+CONTEXT:  COPY pp, line 3: "4,14"
+RESET SESSION AUTHORIZATION;
+DROP TABLE PP;
 -----   RECURSION    ----
 --
 -- Simple recursion
@@ -4079,9 +4130,7 @@ SET row_security TO OFF;
 COPY copy_t FROM STDIN; --fail - would be affected by RLS.
 ERROR:  query would be affected by row-level security policy for table "copy_t"
 SET row_security TO ON;
-COPY copy_t FROM STDIN; --fail - COPY FROM not supported by RLS.
-ERROR:  COPY FROM not supported with row-level security
-HINT:  Use INSERT statements instead.
+COPY copy_t FROM STDIN; --no error
 -- Check COPY FROM as user with permissions and BYPASSRLS
 SET SESSION AUTHORIZATION regress_rls_exempt_user;
 SET row_security TO ON;
@@ -4597,6 +4646,7 @@ ALTER TABLE r1 ENABLE ROW LEVEL SECURITY;
 ALTER TABLE r1 FORCE ROW LEVEL SECURITY;
 -- Works fine
 INSERT INTO r1 VALUES (10), (20);
+COPY r1 FROM STDIN;
 -- No error, but no rows
 TABLE r1;
  a 

src/test/regress/sql/rowsecurity.sql

@@ -283,6 +283,11 @@ EXPLAIN (COSTS OFF) SELECT * FROM document NATURAL JOIN category WHERE f_leak(dt
 -- 44 would technically fail for both p2r and p1r, but we should get an error
 -- back from p1r for this because it sorts first
 INSERT INTO document VALUES (100, 44, 1, 'regress_rls_dave', 'testing sorting of policies'); -- fail
+
+-- fail, COPY FROM, security policy with subquery not supported
+COPY document FROM STDIN;
+\.
+
 -- Just to see a p2r error
 INSERT INTO document VALUES (100, 55, 1, 'regress_rls_dave', 'testing sorting of policies'); -- fail
 
@@ -518,6 +523,10 @@ EXPLAIN (COSTS OFF) SELECT * FROM part_document WHERE f_leak(dtitle);
 
 -- pp1 ERROR
 INSERT INTO part_document VALUES (100, 11, 5, 'regress_rls_dave', 'testing pp1'); -- fail
+-- fail, COPY FROM, security policy with subquery not supported
+COPY part_document FROM STDIN;
+\.
+
 -- pp1r ERROR
 INSERT INTO part_document VALUES (100, 99, 1, 'regress_rls_dave', 'testing pp1r'); -- fail
 
@@ -618,6 +627,56 @@ DROP TABLE dependee CASCADE;
 
 EXPLAIN (COSTS OFF) SELECT * FROM dependent; -- After drop, should be unqualified
 
+--COPY FROM with RLS
+RESET SESSION AUTHORIZATION;
+CREATE TABLE pp (id int,val int) PARTITION BY RANGE (id);
+CREATE TABLE pp_1 (val int, id int);
+ALTER TABLE pp ATTACH PARTITION pp_1 FOR VALUES FROM (1) TO (10);
+CREATE TABLE pp_2 PARTITION OF pp FOR VALUES FROM (10) TO (20);
+ALTER TABLE pp ENABLE ROW LEVEL SECURITY;
+ALTER TABLE pp_1 ENABLE ROW LEVEL SECURITY;
+ALTER TABLE pp_2 ENABLE ROW LEVEL SECURITY;
+CREATE POLICY pp_1_p1 ON pp_1 FOR INSERT WITH CHECK (id = 6);
+CREATE POLICY p2_pp ON pp FOR ALL USING(id = 1 or id = 2);
+GRANT SELECT, INSERT ON pp TO regress_rls_alice;
+GRANT SELECT, INSERT ON pp_1 TO regress_rls_alice;
+GRANT SELECT, INSERT ON pp_2 TO regress_rls_alice;
+
+SET SESSION AUTHORIZATION regress_rls_alice;
+INSERT INTO pp_1 VALUES (13, 2); --error
+INSERT INTO pp_1 VALUES (16, 6); --ok
+COPY pp_1 FROM STDIN WITH DELIMITER ','; --second record not ok
+16,6
+13,2
+\.
+
+INSERT INTO pp VALUES (1,11), (2,12);
+INSERT INTO pp values (5,11); --error
+INSERT INTO pp values (6,11); --error
+COPY pp FROM STDIN WITH DELIMITER ','; --error, second record not ok
+1,11
+5,11
+\.
+
+COPY pp FROM STDIN WITH DELIMITER ','; --error, second record not ok
+2,12
+6,11
+\.
+
+RESET SESSION AUTHORIZATION;
+CREATE POLICY p1_pp ON pp FOR INSERT WITH CHECK(id > 4);
+SET SESSION AUTHORIZATION regress_rls_alice;
+INSERT INTO pp VALUES (5, 15), (6, 16); --ok
+INSERT INTO pp VALUES (4, 14); --error
+COPY pp FROM STDIN WITH DELIMITER ','; --third record will result error
+5,15
+6,16
+4,14
+\.
+
+RESET SESSION AUTHORIZATION;
+DROP TABLE PP;
+
 -----   RECURSION    ----
 
 --
@@ -1795,8 +1854,10 @@ COPY copy_t FROM STDIN; --ok
 SET SESSION AUTHORIZATION regress_rls_bob;
 SET row_security TO OFF;
 COPY copy_t FROM STDIN; --fail - would be affected by RLS.
+\.
 SET row_security TO ON;
-COPY copy_t FROM STDIN; --fail - COPY FROM not supported by RLS.
+COPY copy_t FROM STDIN; --no error
+\.
 
 -- Check COPY FROM as user with permissions and BYPASSRLS
 SET SESSION AUTHORIZATION regress_rls_exempt_user;
@@ -2154,6 +2215,9 @@ ALTER TABLE r1 FORCE ROW LEVEL SECURITY;
 
 -- Works fine
 INSERT INTO r1 VALUES (10), (20);
+COPY r1 FROM STDIN;
+10
+\.
 
 -- No error, but no rows
 TABLE r1;