rls-fix.patch.no-cfbot

diff --git a/src/backend/rewrite/rewriteHandler.c b/src/backend/rewrite/rewriteHandler.c
new file mode 100644
index ef86ded..33ef64c
--- a/src/backend/rewrite/rewriteHandler.c
+++ b/src/backend/rewrite/rewriteHandler.c
@@ -2137,15 +2137,6 @@ fireRIRrules(Query *parsetree, List *act
 			}
 		}
 
-		/*
-		 * Expand virtual generated columns of this table
-		 *
-		 * FIXME: This should be done after applying RLS policies below, since
-		 * those could also contain virtual columns.  But that currently makes
-		 * some tests fail, so it needs further investigation.
-		 */
-		parsetree = (Query *) expand_generated_columns_internal((Node *) parsetree, rel, rt_index, rte);
-
 		table_close(rel, NoLock);
 	}
 
@@ -2171,6 +2162,10 @@ fireRIRrules(Query *parsetree, List *act
 	 * requires special recursion detection if the new quals have sublink
 	 * subqueries, and if we did it in the loop above query_tree_walker would
 	 * then recurse into those quals a second time.
+	 *
+	 * Finally, we expand any virtual generated columns.  We do this after
+	 * each table's RLS policies are applied because the RLS policies might
+	 * also refer to the table's virtual generated columns.
 	 */
 	rt_index = 0;
 	foreach(lc, parsetree->rtable)
@@ -2184,10 +2179,11 @@ fireRIRrules(Query *parsetree, List *act
 
 		++rt_index;
 
-		/* Only normal relations can have RLS policies */
-		if (rte->rtekind != RTE_RELATION ||
-			(rte->relkind != RELKIND_RELATION &&
-			 rte->relkind != RELKIND_PARTITIONED_TABLE))
+		/*
+		 * Only normal relations can have RLS policies or virtual generated
+		 * columns.
+		 */
+		if (rte->rtekind != RTE_RELATION)
 			continue;
 
 		rel = table_open(rte->relid, NoLock);
@@ -2265,6 +2261,14 @@ fireRIRrules(Query *parsetree, List *act
 		if (hasSubLinks)
 			parsetree->hasSubLinks = true;
 
+		/*
+		 * Expand any references to virtual generated columns of this table.
+		 * Note that subqueries in virtual generated column expressions are
+		 * not currently supported, so this cannot add any more sublinks.
+		 */
+		parsetree = (Query *) expand_generated_columns_internal((Node *) parsetree,
+																rel, rt_index, rte);
+
 		table_close(rel, NoLock);
 	}
 
diff --git a/src/test/regress/expected/rowsecurity.out b/src/test/regress/expected/rowsecurity.out
new file mode 100644
index 3191908..0f31602
--- a/src/test/regress/expected/rowsecurity.out
+++ b/src/test/regress/expected/rowsecurity.out
@@ -4378,6 +4378,35 @@ INSERT INTO r1 VALUES (10)
     ON CONFLICT ON CONSTRAINT r1_pkey DO UPDATE SET a = 30;
 ERROR:  new row violates row-level security policy for table "r1"
 DROP TABLE r1;
+--
+-- Test policies using virtual generated columns
+--
+SET SESSION AUTHORIZATION regress_rls_alice;
+SET row_security = on;
+CREATE TABLE r1 (a int, b int GENERATED ALWAYS AS (a * 10) VIRTUAL);
+ALTER TABLE r1 ADD c int GENERATED ALWAYS AS (a * 100) VIRTUAL;
+INSERT INTO r1 VALUES (1), (2), (4);
+CREATE POLICY p0 ON r1 USING (b * 10 = c);
+CREATE POLICY p1 ON r1 AS RESTRICTIVE USING (b > 10);
+CREATE POLICY p2 ON r1 AS RESTRICTIVE USING ((SELECT c) < 400);
+ALTER TABLE r1 ENABLE ROW LEVEL SECURITY;
+ALTER TABLE r1 FORCE ROW LEVEL SECURITY;
+-- Should fail p1
+INSERT INTO r1 VALUES (0);
+ERROR:  new row violates row-level security policy "p1" for table "r1"
+-- Should fail p2
+INSERT INTO r1 VALUES (4);
+ERROR:  new row violates row-level security policy "p2" for table "r1"
+-- OK
+INSERT INTO r1 VALUES (3);
+SELECT * FROM r1;
+ a | b  |  c  
+---+----+-----
+ 2 | 20 | 200
+ 3 | 30 | 300
+(2 rows)
+
+DROP TABLE r1;
 -- Check dependency handling
 RESET SESSION AUTHORIZATION;
 CREATE TABLE dep1 (c1 int);
diff --git a/src/test/regress/sql/rowsecurity.sql b/src/test/regress/sql/rowsecurity.sql
new file mode 100644
index 3011d71..6e15a5e
--- a/src/test/regress/sql/rowsecurity.sql
+++ b/src/test/regress/sql/rowsecurity.sql
@@ -2072,6 +2072,33 @@ INSERT INTO r1 VALUES (10)
 
 DROP TABLE r1;
 
+--
+-- Test policies using virtual generated columns
+--
+SET SESSION AUTHORIZATION regress_rls_alice;
+SET row_security = on;
+CREATE TABLE r1 (a int, b int GENERATED ALWAYS AS (a * 10) VIRTUAL);
+ALTER TABLE r1 ADD c int GENERATED ALWAYS AS (a * 100) VIRTUAL;
+INSERT INTO r1 VALUES (1), (2), (4);
+
+CREATE POLICY p0 ON r1 USING (b * 10 = c);
+CREATE POLICY p1 ON r1 AS RESTRICTIVE USING (b > 10);
+CREATE POLICY p2 ON r1 AS RESTRICTIVE USING ((SELECT c) < 400);
+ALTER TABLE r1 ENABLE ROW LEVEL SECURITY;
+ALTER TABLE r1 FORCE ROW LEVEL SECURITY;
+
+-- Should fail p1
+INSERT INTO r1 VALUES (0);
+
+-- Should fail p2
+INSERT INTO r1 VALUES (4);
+
+-- OK
+INSERT INTO r1 VALUES (3);
+SELECT * FROM r1;
+
+DROP TABLE r1;
+
 -- Check dependency handling
 RESET SESSION AUTHORIZATION;
 CREATE TABLE dep1 (c1 int);