RLS works

src/main/java/com/example/demo/RLSConnectionManager.java

@@ -1,11 +1,13 @@
 package com.example.demo;
 
 import org.springframework.jdbc.core.JdbcTemplate;
+import org.springframework.jdbc.datasource.SingleConnectionDataSource;
 import org.springframework.stereotype.Component;
 
 import javax.sql.DataSource;
 import java.sql.Connection;
 import java.sql.SQLException;
+import java.sql.Statement;
 import java.util.function.Function;
 
 /**
@@ -18,11 +20,9 @@ import java.util.function.Function;
 public class RLSConnectionManager {
     
     private final DataSource dataSource;
-    private final JdbcTemplate jdbcTemplate;
 
-    public RLSConnectionManager(DataSource dataSource, JdbcTemplate jdbcTemplate) {
+    public RLSConnectionManager(DataSource dataSource) {
         this.dataSource = dataSource;
-        this.jdbcTemplate = jdbcTemplate;
     }
 
     /**
@@ -43,18 +43,27 @@ public class RLSConnectionManager {
             // CRITICAL: Set connection to manual commit to ensure atomicity
             connection.setAutoCommit(false);
             
-            // Create a JdbcTemplate bound to THIS specific connection
-            Connection finalConnection = connection;
-            JdbcTemplate scopedTemplate = new JdbcTemplate(dataSource) {
-                // @Override
-                public Connection getConnection() {
-                    return finalConnection;
-                }
-            };
-            
             try {
-                // Set the RLS context variable on THIS connection
-                scopedTemplate.execute("SET LOCAL app.current_user_id = " + userId);
+                // Set the RLS context variable on THIS connection using raw Statement
+                try (Statement stmt = connection.createStatement()) {
+                    stmt.execute("SET LOCAL app.current_user_id = '" + userId + "'");
+                    System.out.println("RLS context set for user: " + userId);
+                }
+                
+                // Verify it's set (for debugging)
+                try (Statement stmt = connection.createStatement()) {
+                    var rs = stmt.executeQuery("SELECT current_setting('app.current_user_id', true)");
+                    if (rs.next()) {
+                        String value = rs.getString(1);
+                        System.out.println("Verified context value: " + value);
+                    }
+                }
+                
+                // Create a JdbcTemplate bound to THIS specific connection
+                // Use SingleConnectionDataSource to ensure the template uses this exact connection
+                SingleConnectionDataSource singleConnectionDataSource = 
+                    new SingleConnectionDataSource(connection, true);
+                JdbcTemplate scopedTemplate = new JdbcTemplate(singleConnectionDataSource);
                 
                 // Execute the operation with the scoped template
                 T result = operation.apply(scopedTemplate);
@@ -70,11 +79,10 @@ public class RLSConnectionManager {
                 throw new RuntimeException("Error executing RLS operation", e);
             } finally {
                 // CRITICAL: Reset the context variable before returning connection to pool
-                // Using RESET is safer than SET because it works even if SET LOCAL wasn't used
-                try {
-                    scopedTemplate.execute("RESET app.current_user_id");
+                try (Statement stmt = connection.createStatement()) {
+                    stmt.execute("RESET app.current_user_id");
+                    System.out.println("RLS context reset");
                 } catch (Exception e) {
-                    // Log but don't throw - we still want to close the connection
                     System.err.println("Warning: Failed to reset RLS context: " + e.getMessage());
                 }
             }

src/main/java/com/example/demo/RLSTestController.java

@@ -69,17 +69,20 @@ public class RLSTestController {
                 "SELECT current_setting('app.current_user_id', true)", 
                 String.class
             );
-            result.put("user1Context", ctx);
+            result.put("user1Context", ctx.isEmpty() ? "EMPTY" : ctx);
             return null;
         });
         
         // Check if context leaked (should be null or empty)
-        String leakedContext = null;
+        String leakedContext;
         try {
             leakedContext = jdbcTemplate.queryForObject(
                 "SELECT current_setting('app.current_user_id', true)", 
                 String.class
             );
+            if (leakedContext == null || leakedContext.isEmpty()) {
+                leakedContext = "NOT_SET";
+            }
         } catch (Exception e) {
             // Expected - variable should not be set
             leakedContext = "NOT_SET";
@@ -92,7 +95,7 @@ public class RLSTestController {
                 "SELECT current_setting('app.current_user_id', true)", 
                 String.class
             );
-            result.put("user2Context", ctx);
+            result.put("user2Context", ctx.isEmpty() ? "EMPTY" : ctx);
             return null;
         });
         
@@ -102,6 +105,9 @@ public class RLSTestController {
                 "SELECT current_setting('app.current_user_id', true)", 
                 String.class
             );
+            if (leakedContext == null || leakedContext.isEmpty()) {
+                leakedContext = "NOT_SET";
+            }
         } catch (Exception e) {
             leakedContext = "NOT_SET";
         }
@@ -157,27 +163,43 @@ public class RLSTestController {
         // Enable RLS
         jdbcTemplate.execute("ALTER TABLE documents ENABLE ROW LEVEL SECURITY");
         
+        // CRITICAL: Force RLS even for table owner (postgres superuser)
+        // Without this, RLS policies are bypassed for the table owner
+        jdbcTemplate.execute("ALTER TABLE documents FORCE ROW LEVEL SECURITY");
+        
         // Create RLS policy
+        // USING clause: determines which rows are visible (for SELECT)
+        // WITH CHECK clause: determines which rows can be inserted/updated
+        // Using NULLIF to handle empty strings from current_setting when variable isn't set
         jdbcTemplate.execute("""
             CREATE POLICY user_documents_policy ON documents
             FOR ALL
-            USING (user_id = current_setting('app.current_user_id', true)::bigint)
+            USING (user_id = NULLIF(current_setting('app.current_user_id', true), '')::bigint)
+            WITH CHECK (user_id = NULLIF(current_setting('app.current_user_id', true), '')::bigint)
         """);
         
-        // Insert test data
-        jdbcTemplate.update(
-            "INSERT INTO documents (title, content, user_id) VALUES (?, ?, ?)",
-            "User 1 Document", "Private content for user 1", 1L
-        );
-        jdbcTemplate.update(
-            "INSERT INTO documents (title, content, user_id) VALUES (?, ?, ?)",
-            "User 2 Document", "Private content for user 2", 2L
-        );
-        jdbcTemplate.update(
-            "INSERT INTO documents (title, content, user_id) VALUES (?, ?, ?)",
-            "Another User 1 Doc", "More private content for user 1", 1L
-        );
+        // Insert test data WITH RLS context set
+        // Now that FORCE RLS is enabled, even our inserts must respect the policy
+        rlsManager.executeWithRLSContext(1L, scopedTemplate -> {
+            scopedTemplate.update(
+                "INSERT INTO documents (title, content, user_id) VALUES (?, ?, ?)",
+                "User 1 Document", "Private content for user 1", 1L
+            );
+            scopedTemplate.update(
+                "INSERT INTO documents (title, content, user_id) VALUES (?, ?, ?)",
+                "Another User 1 Doc", "More private content for user 1", 1L
+            );
+            return null;
+        });
         
-        return "Database setup complete with RLS enabled";
+        rlsManager.executeWithRLSContext(2L, scopedTemplate -> {
+            scopedTemplate.update(
+                "INSERT INTO documents (title, content, user_id) VALUES (?, ?, ?)",
+                "User 2 Document", "Private content for user 2", 2L
+            );
+            return null;
+        });
+        
+        return "Database setup complete with RLS enabled and FORCED for table owner";
     }
 }

src/main/resources/application.properties

@@ -2,6 +2,6 @@ spring.application.name=demo
 
 # PostgreSQL Database Configuration
 spring.datasource.url=jdbc:postgresql://localhost:5432/rls_test
-spring.datasource.username=postgres
-spring.datasource.password=postgres
+spring.datasource.username=rls_test
+spring.datasource.password=rls_test
 spring.datasource.hikari.maximum-pool-size=10