Disconnect clients when updating RLS rules

Author: mamcx

crates/schema/src/auto_migrate.rs

@@ -1019,13 +1019,24 @@ fn auto_migrate_constraints(plan: &mut AutoMigratePlan, new_tables: &HashSet<&Id
 // Because we can refer to many tables and fields on the row level-security query, we need to remove all of them,
 // then add the new ones, instead of trying to track the graph of dependencies.
 fn auto_migrate_row_level_security(plan: &mut AutoMigratePlan) -> Result<()> {
+    // Track if any RLS rules were changed.
+    let mut old_rls = HashSet::new();
+    let mut new_rls = HashSet::new();
+
     for rls in plan.old.row_level_security() {
+        old_rls.insert(rls.key());
         plan.steps.push(AutoMigrateStep::RemoveRowLevelSecurity(rls.key()));
     }
     for rls in plan.new.row_level_security() {
+        new_rls.insert(rls.key());
         plan.steps.push(AutoMigrateStep::AddRowLevelSecurity(rls.key()));
     }
 
+    // We can force flush the cache by force disconnecting all clients if an RLS rule has been added, removed, or updated.
+    if old_rls != new_rls && !plan.disconnects_all_users() {
+        plan.steps.push(AutoMigrateStep::DisconnectAllUsers);
+    }
+
     Ok(())
 }
 
@@ -2245,4 +2256,52 @@ mod tests {
             );
         }
     }
+
+    #[test]
+    fn change_rls_disconnect_clients() {
+        let old_def = create_module_def(|_builder| {});
+
+        let new_def = create_module_def(|_builder| {});
+
+        let plan = ponder_auto_migrate(&old_def, &new_def).expect("auto migration should succeed");
+        assert!(!plan.disconnects_all_users(), "{plan:#?}");
+
+        let old_def = create_module_def(|builder| {
+            builder.add_row_level_security("SELECT true;");
+        });
+        let new_def = create_module_def(|builder| {
+            builder.add_row_level_security("SELECT false;");
+        });
+
+        let plan = ponder_auto_migrate(&old_def, &new_def).expect("auto migration should succeed");
+        assert!(plan.disconnects_all_users(), "{plan:#?}");
+
+        let old_def = create_module_def(|builder| {
+            builder.add_row_level_security("SELECT true;");
+        });
+
+        let new_def = create_module_def(|_builder| {
+            // Remove RLS
+        });
+        let plan = ponder_auto_migrate(&old_def, &new_def).expect("auto migration should succeed");
+        assert!(plan.disconnects_all_users(), "{plan:#?}");
+
+        let old_def = create_module_def(|_builder| {});
+
+        let new_def = create_module_def(|builder| {
+            builder.add_row_level_security("SELECT false;");
+        });
+        let plan = ponder_auto_migrate(&old_def, &new_def).expect("auto migration should succeed");
+        assert!(plan.disconnects_all_users(), "{plan:#?}");
+
+        let old_def = create_module_def(|builder| {
+            builder.add_row_level_security("SELECT true;");
+        });
+
+        let new_def = create_module_def(|builder| {
+            builder.add_row_level_security("SELECT true;");
+        });
+        let plan = ponder_auto_migrate(&old_def, &new_def).expect("auto migration should succeed");
+        assert!(!plan.disconnects_all_users(), "{plan:#?}");
+    }
 }

crates/schema/src/snapshots/spacetimedb_schema__auto_migrate__tests__empty_to_populated_migration.snap

@@ -61,3 +61,4 @@ expression: "plan.pretty_print(PrettyPrintStyle::AnsiColor).expect(\"should pret
 
 ▸ Created row level security policy:
     `SELECT * FROM Apples`
+!!! Warning: All clients will be disconnected due to breaking schema changes

smoketests/tests/auto_migration.py

@@ -146,14 +146,15 @@ def test_add_table_auto_migration(self):
         )
 
         self.write_module_code(self.MODULE_CODE_UPDATED)
-        self.publish_module(self.database_identity, clear=False)
+        self.publish_module(self.database_identity, clear=False, break_clients=True)
 
         logging.info("Updated")
         self.call("add_person", "Husserl", "Student")
 
-        # If subscription, we should get 4 rows corresponding to 4 reducer calls (including before and after update)
+        # If subscription, we should get 3 rows corresponding to 3 reducer calls (including before and after update)
+        # minus the `print` one
         sub = sub();
-        self.assertEqual(len(sub), 4)
+        self.assertEqual(len(sub), 3)
 
         # Check the row-level SQL filter is added correctly
         self.assertSql(

smoketests/tests/rls.py

@@ -1,3 +1,5 @@
+import logging
+
 from .. import Smoketest, random_string
 
 class Rls(Smoketest):
@@ -78,3 +80,71 @@ def test_publish_fails_for_rls_on_private_table(self):
 
         with self.assertRaises(Exception):
             self.publish_module(name)
+
+class DisconnectRls(Smoketest):
+    AUTOPUBLISH = False
+    
+    MODULE_CODE = """
+use spacetimedb::{Identity, ReducerContext, Table};
+
+#[spacetimedb::table(name = users, public)]
+pub struct Users {
+    name: String,
+    identity: Identity,
+}
+
+#[spacetimedb::reducer]
+pub fn add_user(ctx: &ReducerContext, name: String) {
+    ctx.db.users().insert(Users { name, identity: ctx.sender });
+}
+
+// RLS
+"""
+    
+    ADD_RLS = """ 
+    #[spacetimedb::client_visibility_filter]
+const USER_FILTER: spacetimedb::Filter = spacetimedb::Filter::Sql(
+    "SELECT * FROM users WHERE identity = :sender"
+);
+"""
+    
+    def test_rls_disconnect_if_change(self):
+        """This tests that changing the RLS rules disconnects existing clients"""
+        
+        name = random_string()
+        
+        self.write_module_code(self.MODULE_CODE)
+        
+        self.publish_module(name)
+        logging.info("Initial publish complete")
+        
+        # Now add the RLS rules
+        self.write_module_code(self.MODULE_CODE + self.ADD_RLS)
+        self.publish_module(name, clear=False, break_clients=True)
+        
+        logging.info("Re-publish with RLS complete")
+        
+        logs = self.logs(100)
+        
+        # Validate disconnect + schema migration logs
+        self.assertIn("Disconnecting all users", logs)
+    
+    def test_rls_disconnect_no(self):
+        """This tests that not changing the RLS rules does not disconnect existing clients"""
+        
+        name = random_string()
+        
+        self.write_module_code(self.MODULE_CODE + self.ADD_RLS)
+        
+        self.publish_module(name)
+        logging.info("Initial publish complete")
+        
+        # Now re-publish the same module code
+        self.publish_module(name, clear=False, break_clients=False)
+        
+        logging.info("Re-publish without RLS change complete")
+        
+        logs = self.logs(100)
+        
+        # Validate no disconnect logs
+        self.assertNotIn("Disconnecting all users", logs)