From c2c49cd10f001a5ee7a2d52083b2fcd3232fc53e Mon Sep 17 00:00:00 2001
From: Dean Rasheed <dean.a.rasheed@gmail.com>
Date: Thu, 27 Mar 2025 14:08:09 +0000
Subject: [PATCH v1 1/2] New RLS tests to test policies applied by command
 type.

The existing RLS tests focus on the outcome of various testing
scenarios, rather than the exact policies applied. These new tests
list out the policies applied for each command type, including the
different paths through INSERT ... ON CONFLICT and MERGE.
---
 src/test/regress/expected/rowsecurity.out | 226 ++++++++++++++++++++++
 src/test/regress/sql/rowsecurity.sql      | 111 +++++++++++
 2 files changed, 337 insertions(+)

diff --git a/src/test/regress/expected/rowsecurity.out b/src/test/regress/expected/rowsecurity.out
index 87929191d06..ce80cbde938 100644
--- a/src/test/regress/expected/rowsecurity.out
+++ b/src/test/regress/expected/rowsecurity.out
@@ -31,6 +31,232 @@ CREATE OR REPLACE FUNCTION f_leak(text) RETURNS bool
     COST 0.0000001 LANGUAGE plpgsql
     AS 'BEGIN RAISE NOTICE ''f_leak => %'', $1; RETURN true; END';
 GRANT EXECUTE ON FUNCTION f_leak(text) TO public;
+-- Test policies applied by command type
+SET SESSION AUTHORIZATION regress_rls_alice;
+CREATE TABLE rls_test_src (a int PRIMARY KEY, b text);
+ALTER TABLE rls_test_src ENABLE ROW LEVEL SECURITY;
+INSERT INTO rls_test_src VALUES (1, 'src a');
+CREATE TABLE rls_test_tgt (a int PRIMARY KEY, b text, c text);
+ALTER TABLE rls_test_tgt ENABLE ROW LEVEL SECURITY;
+CREATE FUNCTION rls_test_tgt_set_c() RETURNS trigger AS
+  $$ BEGIN new.c = upper(new.b); RETURN new; END; $$
+  LANGUAGE plpgsql;
+CREATE TRIGGER rls_test_tgt_set_c BEFORE INSERT OR UPDATE ON rls_test_tgt
+  FOR EACH ROW EXECUTE FUNCTION rls_test_tgt_set_c();
+CREATE FUNCTION sel_using_fn(text, record) RETURNS bool AS
+  $$ BEGIN RAISE NOTICE 'SELECT USING on %.%', $1, $2; RETURN true; END; $$
+  LANGUAGE plpgsql;
+CREATE FUNCTION ins_check_fn(text, record) RETURNS bool AS
+  $$ BEGIN RAISE NOTICE 'INSERT CHECK on %.%', $1, $2; RETURN true; END; $$
+  LANGUAGE plpgsql;
+CREATE FUNCTION upd_using_fn(text, record) RETURNS bool AS
+  $$ BEGIN RAISE NOTICE 'UPDATE USING on %.%', $1, $2; RETURN true; END; $$
+  LANGUAGE plpgsql;
+CREATE FUNCTION upd_check_fn(text, record) RETURNS bool AS
+  $$ BEGIN RAISE NOTICE 'UPDATE CHECK on %.%', $1, $2; RETURN true; END; $$
+  LANGUAGE plpgsql;
+CREATE FUNCTION del_using_fn(text, record) RETURNS bool AS
+  $$ BEGIN RAISE NOTICE 'DELETE USING on %.%', $1, $2; RETURN true; END; $$
+  LANGUAGE plpgsql;
+CREATE POLICY sel_pol ON rls_test_src FOR SELECT
+  USING (sel_using_fn('rls_test_src', rls_test_src));
+CREATE POLICY upd_pol ON rls_test_src FOR UPDATE
+  USING (upd_using_fn('rls_test_src', rls_test_src))
+  WITH CHECK (upd_check_fn('rls_test_src', rls_test_src));
+CREATE POLICY sel_pol ON rls_test_tgt FOR SELECT
+  USING (sel_using_fn('rls_test_tgt', rls_test_tgt));
+CREATE POLICY ins_pol ON rls_test_tgt FOR INSERT
+  WITH CHECK (ins_check_fn('rls_test_tgt', rls_test_tgt));
+CREATE POLICY upd_pol ON rls_test_tgt FOR UPDATE
+  USING (upd_using_fn('rls_test_tgt', rls_test_tgt))
+  WITH CHECK (upd_check_fn('rls_test_tgt', rls_test_tgt));
+CREATE POLICY del_pol ON rls_test_tgt FOR DELETE
+  USING (del_using_fn('rls_test_tgt', rls_test_tgt));
+GRANT SELECT, UPDATE ON rls_test_src TO public;
+GRANT SELECT, INSERT, UPDATE, DELETE, TRUNCATE ON rls_test_tgt TO public;
+SET SESSION AUTHORIZATION regress_rls_bob;
+SELECT * FROM rls_test_src;
+NOTICE:  SELECT USING on rls_test_src.(1,"src a")
+ a |   b   
+---+-------
+ 1 | src a
+(1 row)
+
+SELECT * FROM rls_test_src FOR UPDATE;
+NOTICE:  UPDATE USING on rls_test_src.(1,"src a")
+NOTICE:  SELECT USING on rls_test_src.(1,"src a")
+ a |   b   
+---+-------
+ 1 | src a
+(1 row)
+
+SELECT * FROM rls_test_src FOR NO KEY UPDATE;
+NOTICE:  UPDATE USING on rls_test_src.(1,"src a")
+NOTICE:  SELECT USING on rls_test_src.(1,"src a")
+ a |   b   
+---+-------
+ 1 | src a
+(1 row)
+
+SELECT * FROM rls_test_src FOR SHARE;
+NOTICE:  UPDATE USING on rls_test_src.(1,"src a")
+NOTICE:  SELECT USING on rls_test_src.(1,"src a")
+ a |   b   
+---+-------
+ 1 | src a
+(1 row)
+
+SELECT * FROM rls_test_src FOR KEY SHARE;
+NOTICE:  UPDATE USING on rls_test_src.(1,"src a")
+NOTICE:  SELECT USING on rls_test_src.(1,"src a")
+ a |   b   
+---+-------
+ 1 | src a
+(1 row)
+
+INSERT INTO rls_test_tgt VALUES (1, 'tgt a');
+NOTICE:  INSERT CHECK on rls_test_tgt.(1,"tgt a","TGT A")
+TRUNCATE rls_test_tgt;
+INSERT INTO rls_test_tgt VALUES (1, 'tgt a') RETURNING *;
+NOTICE:  INSERT CHECK on rls_test_tgt.(1,"tgt a","TGT A")
+NOTICE:  SELECT USING on rls_test_tgt.(1,"tgt a","TGT A")
+ a |   b   |   c   
+---+-------+-------
+ 1 | tgt a | TGT A
+(1 row)
+
+UPDATE rls_test_tgt SET b = 'tgt b';
+NOTICE:  UPDATE USING on rls_test_tgt.(1,"tgt a","TGT A")
+NOTICE:  UPDATE CHECK on rls_test_tgt.(1,"tgt b","TGT B")
+UPDATE rls_test_tgt SET b = 'tgt c' WHERE a = 1;
+NOTICE:  UPDATE USING on rls_test_tgt.(1,"tgt b","TGT B")
+NOTICE:  SELECT USING on rls_test_tgt.(1,"tgt b","TGT B")
+NOTICE:  UPDATE CHECK on rls_test_tgt.(1,"tgt c","TGT C")
+NOTICE:  SELECT USING on rls_test_tgt.(1,"tgt c","TGT C")
+UPDATE rls_test_tgt SET b = 'tgt d' RETURNING *;
+NOTICE:  UPDATE USING on rls_test_tgt.(1,"tgt c","TGT C")
+NOTICE:  SELECT USING on rls_test_tgt.(1,"tgt c","TGT C")
+NOTICE:  UPDATE CHECK on rls_test_tgt.(1,"tgt d","TGT D")
+NOTICE:  SELECT USING on rls_test_tgt.(1,"tgt d","TGT D")
+ a |   b   |   c   
+---+-------+-------
+ 1 | tgt d | TGT D
+(1 row)
+
+BEGIN; DELETE FROM rls_test_tgt; ROLLBACK;
+NOTICE:  DELETE USING on rls_test_tgt.(1,"tgt d","TGT D")
+BEGIN; DELETE FROM rls_test_tgt WHERE a = 1; ROLLBACK;
+NOTICE:  DELETE USING on rls_test_tgt.(1,"tgt d","TGT D")
+NOTICE:  SELECT USING on rls_test_tgt.(1,"tgt d","TGT D")
+DELETE FROM rls_test_tgt RETURNING *;
+NOTICE:  DELETE USING on rls_test_tgt.(1,"tgt d","TGT D")
+NOTICE:  SELECT USING on rls_test_tgt.(1,"tgt d","TGT D")
+ a |   b   |   c   
+---+-------+-------
+ 1 | tgt d | TGT D
+(1 row)
+
+INSERT INTO rls_test_tgt VALUES (1, 'tgt a') ON CONFLICT (a) DO NOTHING;
+NOTICE:  INSERT CHECK on rls_test_tgt.(1,"tgt a","TGT A")
+NOTICE:  SELECT USING on rls_test_tgt.(1,"tgt a","TGT A")
+INSERT INTO rls_test_tgt VALUES (1, 'tgt b') ON CONFLICT (a) DO NOTHING;
+NOTICE:  INSERT CHECK on rls_test_tgt.(1,"tgt b","TGT B")
+NOTICE:  SELECT USING on rls_test_tgt.(1,"tgt b","TGT B")
+BEGIN;
+INSERT INTO rls_test_tgt VALUES (2, 'tgt a') ON CONFLICT (a) DO UPDATE SET b = 'tgt b';
+NOTICE:  INSERT CHECK on rls_test_tgt.(2,"tgt a","TGT A")
+NOTICE:  SELECT USING on rls_test_tgt.(2,"tgt a","TGT A")
+INSERT INTO rls_test_tgt VALUES (2, 'tgt c') ON CONFLICT (a) DO UPDATE SET b = 'tgt d';
+NOTICE:  INSERT CHECK on rls_test_tgt.(2,"tgt c","TGT C")
+NOTICE:  SELECT USING on rls_test_tgt.(2,"tgt c","TGT C")
+NOTICE:  UPDATE USING on rls_test_tgt.(2,"tgt a","TGT A")
+NOTICE:  SELECT USING on rls_test_tgt.(2,"tgt a","TGT A")
+NOTICE:  UPDATE CHECK on rls_test_tgt.(2,"tgt d","TGT D")
+NOTICE:  SELECT USING on rls_test_tgt.(2,"tgt d","TGT D")
+INSERT INTO rls_test_tgt VALUES (3, 'tgt a') ON CONFLICT (a) DO UPDATE SET b = 'tgt b' RETURNING *;
+NOTICE:  INSERT CHECK on rls_test_tgt.(3,"tgt a","TGT A")
+NOTICE:  SELECT USING on rls_test_tgt.(3,"tgt a","TGT A")
+ a |   b   |   c   
+---+-------+-------
+ 3 | tgt a | TGT A
+(1 row)
+
+INSERT INTO rls_test_tgt VALUES (3, 'tgt c') ON CONFLICT (a) DO UPDATE SET b = 'tgt d' RETURNING *;
+NOTICE:  INSERT CHECK on rls_test_tgt.(3,"tgt c","TGT C")
+NOTICE:  SELECT USING on rls_test_tgt.(3,"tgt c","TGT C")
+NOTICE:  UPDATE USING on rls_test_tgt.(3,"tgt a","TGT A")
+NOTICE:  SELECT USING on rls_test_tgt.(3,"tgt a","TGT A")
+NOTICE:  UPDATE CHECK on rls_test_tgt.(3,"tgt d","TGT D")
+NOTICE:  SELECT USING on rls_test_tgt.(3,"tgt d","TGT D")
+ a |   b   |   c   
+---+-------+-------
+ 3 | tgt d | TGT D
+(1 row)
+
+ROLLBACK;
+MERGE INTO rls_test_tgt t USING rls_test_src s ON t.a = s.a
+  WHEN NOT MATCHED THEN DO NOTHING;
+NOTICE:  SELECT USING on rls_test_tgt.(1,"tgt a","TGT A")
+NOTICE:  SELECT USING on rls_test_src.(1,"src a")
+TRUNCATE rls_test_tgt;
+MERGE INTO rls_test_tgt t USING rls_test_src s ON t.a = s.a
+  WHEN NOT MATCHED THEN INSERT VALUES (1, 'tgt a');
+NOTICE:  SELECT USING on rls_test_src.(1,"src a")
+NOTICE:  INSERT CHECK on rls_test_tgt.(1,"tgt a","TGT A")
+TRUNCATE rls_test_tgt;
+MERGE INTO rls_test_tgt t USING rls_test_src s ON t.a = s.a
+  WHEN NOT MATCHED THEN INSERT VALUES (1, 'tgt a')
+  RETURNING *;
+NOTICE:  SELECT USING on rls_test_src.(1,"src a")
+NOTICE:  INSERT CHECK on rls_test_tgt.(1,"tgt a","TGT A")
+NOTICE:  SELECT USING on rls_test_tgt.(1,"tgt a","TGT A")
+ a |   b   | a |   b   |   c   
+---+-------+---+-------+-------
+ 1 | src a | 1 | tgt a | TGT A
+(1 row)
+
+MERGE INTO rls_test_tgt t USING rls_test_src s ON t.a = s.a
+  WHEN MATCHED THEN UPDATE SET b = 'tgt b';
+NOTICE:  SELECT USING on rls_test_tgt.(1,"tgt a","TGT A")
+NOTICE:  SELECT USING on rls_test_src.(1,"src a")
+NOTICE:  UPDATE USING on rls_test_tgt.(1,"tgt a","TGT A")
+NOTICE:  UPDATE CHECK on rls_test_tgt.(1,"tgt b","TGT B")
+NOTICE:  SELECT USING on rls_test_tgt.(1,"tgt b","TGT B")
+MERGE INTO rls_test_tgt t USING rls_test_src s ON t.a = s.a
+  WHEN MATCHED THEN UPDATE SET b = 'tgt c'
+  RETURNING *;
+NOTICE:  SELECT USING on rls_test_tgt.(1,"tgt b","TGT B")
+NOTICE:  SELECT USING on rls_test_src.(1,"src a")
+NOTICE:  UPDATE USING on rls_test_tgt.(1,"tgt b","TGT B")
+NOTICE:  UPDATE CHECK on rls_test_tgt.(1,"tgt c","TGT C")
+NOTICE:  SELECT USING on rls_test_tgt.(1,"tgt c","TGT C")
+ a |   b   | a |   b   |   c   
+---+-------+---+-------+-------
+ 1 | src a | 1 | tgt c | TGT C
+(1 row)
+
+BEGIN;
+MERGE INTO rls_test_tgt t USING rls_test_src s ON t.a = s.a
+  WHEN MATCHED THEN DELETE;
+NOTICE:  SELECT USING on rls_test_tgt.(1,"tgt c","TGT C")
+NOTICE:  SELECT USING on rls_test_src.(1,"src a")
+NOTICE:  DELETE USING on rls_test_tgt.(1,"tgt c","TGT C")
+ROLLBACK;
+MERGE INTO rls_test_tgt t USING rls_test_src s ON t.a = s.a
+  WHEN MATCHED THEN DELETE
+  RETURNING *;
+NOTICE:  SELECT USING on rls_test_tgt.(1,"tgt c","TGT C")
+NOTICE:  SELECT USING on rls_test_src.(1,"src a")
+NOTICE:  DELETE USING on rls_test_tgt.(1,"tgt c","TGT C")
+ a |   b   | a |   b   |   c   
+---+-------+---+-------+-------
+ 1 | src a | 1 | tgt c | TGT C
+(1 row)
+
+RESET SESSION AUTHORIZATION;
+DROP TABLE rls_test_src, rls_test_tgt;
+DROP FUNCTION rls_test_tgt_set_c;
+DROP FUNCTION sel_using_fn, ins_check_fn, upd_using_fn, upd_check_fn, del_using_fn;
 -- BASIC Row-Level Security Scenario
 SET SESSION AUTHORIZATION regress_rls_alice;
 CREATE TABLE uaccount (
diff --git a/src/test/regress/sql/rowsecurity.sql b/src/test/regress/sql/rowsecurity.sql
index f61dbbf9581..4bec2672a4f 100644
--- a/src/test/regress/sql/rowsecurity.sql
+++ b/src/test/regress/sql/rowsecurity.sql
@@ -41,6 +41,117 @@ CREATE OR REPLACE FUNCTION f_leak(text) RETURNS bool
     AS 'BEGIN RAISE NOTICE ''f_leak => %'', $1; RETURN true; END';
 GRANT EXECUTE ON FUNCTION f_leak(text) TO public;
 
+-- Test policies applied by command type
+SET SESSION AUTHORIZATION regress_rls_alice;
+
+CREATE TABLE rls_test_src (a int PRIMARY KEY, b text);
+ALTER TABLE rls_test_src ENABLE ROW LEVEL SECURITY;
+INSERT INTO rls_test_src VALUES (1, 'src a');
+
+CREATE TABLE rls_test_tgt (a int PRIMARY KEY, b text, c text);
+ALTER TABLE rls_test_tgt ENABLE ROW LEVEL SECURITY;
+
+CREATE FUNCTION rls_test_tgt_set_c() RETURNS trigger AS
+  $$ BEGIN new.c = upper(new.b); RETURN new; END; $$
+  LANGUAGE plpgsql;
+CREATE TRIGGER rls_test_tgt_set_c BEFORE INSERT OR UPDATE ON rls_test_tgt
+  FOR EACH ROW EXECUTE FUNCTION rls_test_tgt_set_c();
+
+CREATE FUNCTION sel_using_fn(text, record) RETURNS bool AS
+  $$ BEGIN RAISE NOTICE 'SELECT USING on %.%', $1, $2; RETURN true; END; $$
+  LANGUAGE plpgsql;
+CREATE FUNCTION ins_check_fn(text, record) RETURNS bool AS
+  $$ BEGIN RAISE NOTICE 'INSERT CHECK on %.%', $1, $2; RETURN true; END; $$
+  LANGUAGE plpgsql;
+CREATE FUNCTION upd_using_fn(text, record) RETURNS bool AS
+  $$ BEGIN RAISE NOTICE 'UPDATE USING on %.%', $1, $2; RETURN true; END; $$
+  LANGUAGE plpgsql;
+CREATE FUNCTION upd_check_fn(text, record) RETURNS bool AS
+  $$ BEGIN RAISE NOTICE 'UPDATE CHECK on %.%', $1, $2; RETURN true; END; $$
+  LANGUAGE plpgsql;
+CREATE FUNCTION del_using_fn(text, record) RETURNS bool AS
+  $$ BEGIN RAISE NOTICE 'DELETE USING on %.%', $1, $2; RETURN true; END; $$
+  LANGUAGE plpgsql;
+
+CREATE POLICY sel_pol ON rls_test_src FOR SELECT
+  USING (sel_using_fn('rls_test_src', rls_test_src));
+CREATE POLICY upd_pol ON rls_test_src FOR UPDATE
+  USING (upd_using_fn('rls_test_src', rls_test_src))
+  WITH CHECK (upd_check_fn('rls_test_src', rls_test_src));
+
+CREATE POLICY sel_pol ON rls_test_tgt FOR SELECT
+  USING (sel_using_fn('rls_test_tgt', rls_test_tgt));
+CREATE POLICY ins_pol ON rls_test_tgt FOR INSERT
+  WITH CHECK (ins_check_fn('rls_test_tgt', rls_test_tgt));
+CREATE POLICY upd_pol ON rls_test_tgt FOR UPDATE
+  USING (upd_using_fn('rls_test_tgt', rls_test_tgt))
+  WITH CHECK (upd_check_fn('rls_test_tgt', rls_test_tgt));
+CREATE POLICY del_pol ON rls_test_tgt FOR DELETE
+  USING (del_using_fn('rls_test_tgt', rls_test_tgt));
+
+GRANT SELECT, UPDATE ON rls_test_src TO public;
+GRANT SELECT, INSERT, UPDATE, DELETE, TRUNCATE ON rls_test_tgt TO public;
+
+SET SESSION AUTHORIZATION regress_rls_bob;
+
+SELECT * FROM rls_test_src;
+SELECT * FROM rls_test_src FOR UPDATE;
+SELECT * FROM rls_test_src FOR NO KEY UPDATE;
+SELECT * FROM rls_test_src FOR SHARE;
+SELECT * FROM rls_test_src FOR KEY SHARE;
+
+INSERT INTO rls_test_tgt VALUES (1, 'tgt a');
+TRUNCATE rls_test_tgt;
+INSERT INTO rls_test_tgt VALUES (1, 'tgt a') RETURNING *;
+
+UPDATE rls_test_tgt SET b = 'tgt b';
+UPDATE rls_test_tgt SET b = 'tgt c' WHERE a = 1;
+UPDATE rls_test_tgt SET b = 'tgt d' RETURNING *;
+
+BEGIN; DELETE FROM rls_test_tgt; ROLLBACK;
+BEGIN; DELETE FROM rls_test_tgt WHERE a = 1; ROLLBACK;
+DELETE FROM rls_test_tgt RETURNING *;
+
+INSERT INTO rls_test_tgt VALUES (1, 'tgt a') ON CONFLICT (a) DO NOTHING;
+INSERT INTO rls_test_tgt VALUES (1, 'tgt b') ON CONFLICT (a) DO NOTHING;
+
+BEGIN;
+INSERT INTO rls_test_tgt VALUES (2, 'tgt a') ON CONFLICT (a) DO UPDATE SET b = 'tgt b';
+INSERT INTO rls_test_tgt VALUES (2, 'tgt c') ON CONFLICT (a) DO UPDATE SET b = 'tgt d';
+INSERT INTO rls_test_tgt VALUES (3, 'tgt a') ON CONFLICT (a) DO UPDATE SET b = 'tgt b' RETURNING *;
+INSERT INTO rls_test_tgt VALUES (3, 'tgt c') ON CONFLICT (a) DO UPDATE SET b = 'tgt d' RETURNING *;
+ROLLBACK;
+
+MERGE INTO rls_test_tgt t USING rls_test_src s ON t.a = s.a
+  WHEN NOT MATCHED THEN DO NOTHING;
+
+TRUNCATE rls_test_tgt;
+MERGE INTO rls_test_tgt t USING rls_test_src s ON t.a = s.a
+  WHEN NOT MATCHED THEN INSERT VALUES (1, 'tgt a');
+TRUNCATE rls_test_tgt;
+MERGE INTO rls_test_tgt t USING rls_test_src s ON t.a = s.a
+  WHEN NOT MATCHED THEN INSERT VALUES (1, 'tgt a')
+  RETURNING *;
+
+MERGE INTO rls_test_tgt t USING rls_test_src s ON t.a = s.a
+  WHEN MATCHED THEN UPDATE SET b = 'tgt b';
+MERGE INTO rls_test_tgt t USING rls_test_src s ON t.a = s.a
+  WHEN MATCHED THEN UPDATE SET b = 'tgt c'
+  RETURNING *;
+
+BEGIN;
+MERGE INTO rls_test_tgt t USING rls_test_src s ON t.a = s.a
+  WHEN MATCHED THEN DELETE;
+ROLLBACK;
+MERGE INTO rls_test_tgt t USING rls_test_src s ON t.a = s.a
+  WHEN MATCHED THEN DELETE
+  RETURNING *;
+
+RESET SESSION AUTHORIZATION;
+DROP TABLE rls_test_src, rls_test_tgt;
+DROP FUNCTION rls_test_tgt_set_c;
+DROP FUNCTION sel_using_fn, ins_check_fn, upd_using_fn, upd_check_fn, del_using_fn;
+
 -- BASIC Row-Level Security Scenario
 
 SET SESSION AUTHORIZATION regress_rls_alice;
-- 
2.43.0

