1 files changed, 65 insertions, 1 deletions

diff --git a/src/backend/commands/copy.c b/src/backend/commands/copy.c
index fbd7492a73f..6b8357634a7 100644
--- a/src/backend/commands/copy.c
+++ b/src/backend/commands/copy.c
@@ -37,7 +37,9 @@
 #include "optimizer/clauses.h"
 #include "optimizer/planner.h"
 #include "parser/parse_relation.h"
+#include "nodes/makefuncs.h"
 #include "rewrite/rewriteHandler.h"
+#include "rewrite/rowsecurity.h"
 #include "storage/fd.h"
 #include "tcop/tcopprot.h"
 #include "utils/acl.h"
@@ -784,6 +786,7 @@ DoCopy(const CopyStmt *stmt, const char *queryString, uint64 *processed)
 	bool		pipe = (stmt->filename == NULL);
 	Relation	rel;
 	Oid			relid;
+	Node	   *query = NULL;
 
 	/* Disallow COPY to/from file or program except to superusers. */
 	if (!pipe && !superuser())
@@ -837,11 +840,72 @@ DoCopy(const CopyStmt *stmt, const char *queryString, uint64 *processed)
 				rte->selectedCols = bms_add_member(rte->selectedCols, attno);
 		}
 		ExecCheckRTPerms(list_make1(rte), true);
+
+		/*
+		 * Permission check for row security.
+		 *
+		 * check_enable_rls will ereport(ERROR) if the user has requested
+		 * something invalid and will otherwise indicate if we should enable
+		 * RLS (returns RLS_ENABLED) or not for this COPY statement.
+		 *
+		 * If the relation has a row security policy and we are to apply it
+		 * then perform a "query" copy and allow the normal query processing to
+		 * handle the policies.
+		 *
+		 * If RLS is not enabled for this, then just fall through to the
+		 * normal non-filtering relation handling.
+		 */
+		if (check_enable_rls(rte->relid, InvalidOid) == RLS_ENABLED)
+		{
+			SelectStmt *select;
+			ColumnRef  *cr;
+			ResTarget  *target;
+			RangeVar   *from;
+
+			if (is_from)
+				ereport(ERROR,
+						(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+						 errmsg("COPY FROM not supported with row security."),
+						 errhint("Use direct INSERT statements instead.")));
+
+			/* Build target list */
+			cr = makeNode(ColumnRef);
+
+			if (!stmt->attlist)
+				cr->fields = list_make1(makeNode(A_Star));
+			else
+				cr->fields = stmt->attlist;
+
+			cr->location = 1;
+
+			target = makeNode(ResTarget);
+			target->name = NULL;
+			target->indirection = NIL;
+			target->val = (Node *) cr;
+			target->location = 1;
+
+			/* Build FROM clause */
+			from = makeRangeVar(NULL, RelationGetRelationName(rel), 1);
+
+			/* Build query */
+			select = makeNode(SelectStmt);
+			select->targetList = list_make1(target);
+			select->fromClause = list_make1(from);
+
+			query = (Node*) select;
+
+			relid = InvalidOid;
+
+			/* Close the handle to the relation as it is no longer needed. */
+			heap_close(rel, (is_from ? RowExclusiveLock : AccessShareLock));
+			rel = NULL;
+		}
 	}
 	else
 	{
 		Assert(stmt->query);
 
+		query = stmt->query;
 		relid = InvalidOid;
 		rel = NULL;
 	}
@@ -861,7 +925,7 @@ DoCopy(const CopyStmt *stmt, const char *queryString, uint64 *processed)
 	}
 	else
 	{
-		cstate = BeginCopyTo(rel, stmt->query, queryString,
+		cstate = BeginCopyTo(rel, query, queryString,
 							 stmt->filename, stmt->is_program,
 							 stmt->attlist, stmt->options);
 		*processed = DoCopyTo(cstate);	/* copy from database to file */