Make RLS work with UPDATE ... WHERE CURRENT OF

UPDATE ... WHERE CURRENT OF would not work in conjunction with
RLS. Arrange to allow the CURRENT OF expression to be pushed down.
Issue noted by Peter Geoghegan. Patch by Dean Rasheed. Back patch
to 9.5 where RLS was introduced.

src/backend/optimizer/path/allpaths.c

@@ -2177,6 +2177,46 @@ subquery_push_qual(Query *subquery, RangeTblEntry *rte, Index rti, Node *qual)
 		recurse_push_qual(subquery->setOperations, subquery,
 						  rte, rti, qual);
 	}
+	else if (IsA(qual, CurrentOfExpr))
+	{
+		/*
+		 * This is possible when a WHERE CURRENT OF expression is applied to a
+		 * table with row-level security.  In that case, the subquery should
+		 * contain precisely one rtable entry for the table, and we can safely
+		 * push the expression down into the subquery.  This will cause a TID
+		 * scan subquery plan to be generated allowing the target relation to
+		 * be updated.
+		 *
+		 * Someday we might also be able to use a WHERE CURRENT OF expression
+		 * on a view, but currently the rewriter prevents that, so we should
+		 * never see any other case here, but generate sane error messages in
+		 * case it does somehow happen.
+		 */
+		if (subquery->rtable == NIL)
+			ereport(ERROR,
+					(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+					 errmsg("WHERE CURRENT OF is not supported on a view with no underlying relation")));
+
+		if (list_length(subquery->rtable) > 1)
+			ereport(ERROR,
+					(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+					 errmsg("WHERE CURRENT OF is not supported on a view with more than one underlying relation")));
+
+		if (subquery->hasAggs || subquery->groupClause || subquery->groupingSets || subquery->havingQual)
+			ereport(ERROR,
+					(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+					 errmsg("WHERE CURRENT OF is not supported on a view with grouping or aggregation")));
+
+		/*
+		 * Adjust the CURRENT OF expression to refer to the underlying table
+		 * in the subquery, and attach it to the subquery's WHERE clause.
+		 */
+		qual = copyObject(qual);
+		((CurrentOfExpr *) qual)->cvarno = 1;
+
+		subquery->jointree->quals =
+			make_and_qual(subquery->jointree->quals, qual);
+	}
 	else
 	{
 		/*

src/backend/optimizer/util/clauses.c

@@ -1492,6 +1492,16 @@ contain_leaked_vars_walker(Node *node, void *context)
 			}
 			break;
 
+		case T_CurrentOfExpr:
+
+			/*
+			 * WHERE CURRENT OF doesn't contain function calls.  Moreover, it
+			 * is important that this can be pushed down into a
+			 * security_barrier view, since the planner must always generate
+			 * a TID scan when CURRENT OF is present -- c.f. cost_tidscan.
+			 */
+			return false;
+
 		default:
 
 			/*

src/test/regress/expected/rowsecurity.out

@@ -2729,6 +2729,113 @@ COPY copy_t FROM STDIN; --fail - permission denied.
 ERROR:  permission denied for relation copy_t
 RESET SESSION AUTHORIZATION;
 DROP TABLE copy_t;
+-- Check WHERE CURRENT OF
+SET SESSION AUTHORIZATION rls_regress_user0;
+CREATE TABLE current_check (currentid int, payload text, rlsuser text);
+GRANT ALL ON current_check TO PUBLIC;
+INSERT INTO current_check VALUES
+    (1, 'abc', 'rls_regress_user1'),
+    (2, 'bcd', 'rls_regress_user1'),
+    (3, 'cde', 'rls_regress_user1'),
+    (4, 'def', 'rls_regress_user1');
+CREATE POLICY p1 ON current_check FOR SELECT USING (currentid % 2 = 0);
+CREATE POLICY p2 ON current_check FOR DELETE USING (currentid = 4 AND rlsuser = current_user);
+CREATE POLICY p3 ON current_check FOR UPDATE USING (currentid = 4) WITH CHECK (rlsuser = current_user);
+ALTER TABLE current_check ENABLE ROW LEVEL SECURITY;
+SET SESSION AUTHORIZATION rls_regress_user1;
+-- Can SELECT even rows
+SELECT * FROM current_check;
+ currentid | payload |      rlsuser      
+-----------+---------+-------------------
+         2 | bcd     | rls_regress_user1
+         4 | def     | rls_regress_user1
+(2 rows)
+
+-- Cannot UPDATE row 2
+UPDATE current_check SET payload = payload || '_new' WHERE currentid = 2 RETURNING *;
+ currentid | payload | rlsuser 
+-----------+---------+---------
+(0 rows)
+
+BEGIN;
+DECLARE current_check_cursor SCROLL CURSOR FOR SELECT * FROM current_check;
+-- Returns rows that can be seen according to SELECT policy, like plain SELECT
+-- above (even rows)
+FETCH ABSOLUTE 1 FROM current_check_cursor;
+ currentid | payload |      rlsuser      
+-----------+---------+-------------------
+         2 | bcd     | rls_regress_user1
+(1 row)
+
+-- Still cannot UPDATE row 2 through cursor
+UPDATE current_check SET payload = payload || '_new' WHERE CURRENT OF current_check_cursor RETURNING *;
+ currentid | payload | rlsuser 
+-----------+---------+---------
+(0 rows)
+
+-- Can update row 4 through cursor, which is the next visible row
+FETCH RELATIVE 1 FROM current_check_cursor;
+ currentid | payload |      rlsuser      
+-----------+---------+-------------------
+         4 | def     | rls_regress_user1
+(1 row)
+
+UPDATE current_check SET payload = payload || '_new' WHERE CURRENT OF current_check_cursor RETURNING *;
+ currentid | payload |      rlsuser      
+-----------+---------+-------------------
+         4 | def_new | rls_regress_user1
+(1 row)
+
+SELECT * FROM current_check;
+ currentid | payload |      rlsuser      
+-----------+---------+-------------------
+         2 | bcd     | rls_regress_user1
+         4 | def_new | rls_regress_user1
+(2 rows)
+
+-- Plan should be a subquery TID scan
+EXPLAIN (COSTS OFF) UPDATE current_check SET payload = payload WHERE CURRENT OF current_check_cursor;
+                          QUERY PLAN                           
+---------------------------------------------------------------
+ Update on current_check current_check_1
+   ->  Subquery Scan on current_check
+         ->  LockRows
+               ->  Tid Scan on current_check current_check_2
+                     TID Cond: CURRENT OF current_check_cursor
+                     Filter: (currentid = 4)
+(6 rows)
+
+-- Similarly can only delete row 4
+FETCH ABSOLUTE 1 FROM current_check_cursor;
+ currentid | payload |      rlsuser      
+-----------+---------+-------------------
+         2 | bcd     | rls_regress_user1
+(1 row)
+
+DELETE FROM current_check WHERE CURRENT OF current_check_cursor RETURNING *;
+ currentid | payload | rlsuser 
+-----------+---------+---------
+(0 rows)
+
+FETCH RELATIVE 1 FROM current_check_cursor;
+ currentid | payload |      rlsuser      
+-----------+---------+-------------------
+         4 | def     | rls_regress_user1
+(1 row)
+
+DELETE FROM current_check WHERE CURRENT OF current_check_cursor RETURNING *;
+ currentid | payload |      rlsuser      
+-----------+---------+-------------------
+         4 | def_new | rls_regress_user1
+(1 row)
+
+SELECT * FROM current_check;
+ currentid | payload |      rlsuser      
+-----------+---------+-------------------
+         2 | bcd     | rls_regress_user1
+(1 row)
+
+COMMIT;
 --
 -- Collation support
 --

src/test/regress/sql/rowsecurity.sql

@@ -1087,6 +1087,55 @@ COPY copy_t FROM STDIN; --fail - permission denied.
 RESET SESSION AUTHORIZATION;
 DROP TABLE copy_t;
 
+-- Check WHERE CURRENT OF
+SET SESSION AUTHORIZATION rls_regress_user0;
+
+CREATE TABLE current_check (currentid int, payload text, rlsuser text);
+GRANT ALL ON current_check TO PUBLIC;
+
+INSERT INTO current_check VALUES
+    (1, 'abc', 'rls_regress_user1'),
+    (2, 'bcd', 'rls_regress_user1'),
+    (3, 'cde', 'rls_regress_user1'),
+    (4, 'def', 'rls_regress_user1');
+
+CREATE POLICY p1 ON current_check FOR SELECT USING (currentid % 2 = 0);
+CREATE POLICY p2 ON current_check FOR DELETE USING (currentid = 4 AND rlsuser = current_user);
+CREATE POLICY p3 ON current_check FOR UPDATE USING (currentid = 4) WITH CHECK (rlsuser = current_user);
+
+ALTER TABLE current_check ENABLE ROW LEVEL SECURITY;
+
+SET SESSION AUTHORIZATION rls_regress_user1;
+
+-- Can SELECT even rows
+SELECT * FROM current_check;
+
+-- Cannot UPDATE row 2
+UPDATE current_check SET payload = payload || '_new' WHERE currentid = 2 RETURNING *;
+
+BEGIN;
+
+DECLARE current_check_cursor SCROLL CURSOR FOR SELECT * FROM current_check;
+-- Returns rows that can be seen according to SELECT policy, like plain SELECT
+-- above (even rows)
+FETCH ABSOLUTE 1 FROM current_check_cursor;
+-- Still cannot UPDATE row 2 through cursor
+UPDATE current_check SET payload = payload || '_new' WHERE CURRENT OF current_check_cursor RETURNING *;
+-- Can update row 4 through cursor, which is the next visible row
+FETCH RELATIVE 1 FROM current_check_cursor;
+UPDATE current_check SET payload = payload || '_new' WHERE CURRENT OF current_check_cursor RETURNING *;
+SELECT * FROM current_check;
+-- Plan should be a subquery TID scan
+EXPLAIN (COSTS OFF) UPDATE current_check SET payload = payload WHERE CURRENT OF current_check_cursor;
+-- Similarly can only delete row 4
+FETCH ABSOLUTE 1 FROM current_check_cursor;
+DELETE FROM current_check WHERE CURRENT OF current_check_cursor RETURNING *;
+FETCH RELATIVE 1 FROM current_check_cursor;
+DELETE FROM current_check WHERE CURRENT OF current_check_cursor RETURNING *;
+SELECT * FROM current_check;
+
+COMMIT;
+
 --
 -- Collation support
 --