Fixes #728: Copy the tables of models made branchable by a resolver into new branches

docs/plugin-development.md

@@ -90,7 +90,7 @@ def my_resolver(model) -> bool | None:
 
 | Return value | Meaning |
 |---|---|
-| `True`  | Model is branchable; route queries to the active branch (still subject to the `exempt_models` filter). |
+| `True`  | Model is branchable; route queries to the active branch, and give new branches their own copy of its table and its M2M through tables (still subject to the `exempt_models` filter). |
 | `False` | Model is not branchable; always route to main. |
 | `None`  | Defer to the next resolver, or to the default `ChangeLoggingMixin` heuristic. |
 

netbox_branching/tests/test_branching_resolver.py

@@ -8,12 +8,16 @@
 
 from contextlib import contextmanager
 
-from django.contrib.auth.models import Group
-from django.test import TestCase
+from django.contrib.auth.models import Group, Permission
+from django.db import connection
+from django.test import TestCase, override_settings
 
 from ipam.models import Prefix
 from netbox_branching import utilities
+from netbox_branching.tests.utils import FastTeardownTransactionTestCase, provision_branch
 from netbox_branching.utilities import (
+    activate_branch,
+    get_tables_to_replicate,
     register_branching_resolver,
     supports_branching,
 )
@@ -195,3 +199,67 @@ def test_exempt_models_filter_still_runs_after_resolver(self):
                 }
             ):
                 self.assertFalse(supports_branching(Group))
+
+
+def _group_resolver(model):
+    return True if model is Group else None
+
+
+GroupPermissions = Group.permissions.through
+GROUP_PERMISSIONS_TABLE = GroupPermissions._meta.db_table
+
+
+class ResolverTableReplicationTestCase(TestCase):
+    """``get_tables_to_replicate`` includes the tables of models a resolver makes branchable."""
+
+    def test_resolver_model_table_is_replicated(self):
+        self.assertNotIn(Group._meta.db_table, get_tables_to_replicate())
+        with _isolated_resolvers():
+            register_branching_resolver(_group_resolver)
+            tables = get_tables_to_replicate()
+            self.assertIn(Group._meta.db_table, tables)
+            self.assertIn(GROUP_PERMISSIONS_TABLE, tables)
+
+    def test_exempt_resolver_model_table_is_not_replicated(self):
+        with _isolated_resolvers():
+            register_branching_resolver(_group_resolver)
+            with override_settings(PLUGINS_CONFIG={'netbox_branching': {'exempt_models': ['auth.group']}}):
+                tables = get_tables_to_replicate()
+                self.assertNotIn(Group._meta.db_table, tables)
+                self.assertNotIn(GROUP_PERMISSIONS_TABLE, tables)
+
+
+class ResolverModelBranchTestCase(FastTeardownTransactionTestCase):
+    """A branch provisioned with a resolver registered gets its own copy of the model's table."""
+
+    serialized_rollback = True
+
+    def test_changes_in_branch_stay_in_branch(self):
+        Group.objects.create(name='Group A')
+        permission = Permission.objects.first()
+        with _isolated_resolvers():
+            register_branching_resolver(_group_resolver)
+            branch = provision_branch(user=None, name='Resolver Branch')
+
+            with connection.cursor() as cursor:
+                cursor.execute(
+                    'SELECT 1 FROM information_schema.tables WHERE table_schema=%s AND table_name=%s',
+                    [branch.schema_name, Group._meta.db_table],
+                )
+                self.assertIsNotNone(cursor.fetchone())
+
+            with activate_branch(branch):
+                group = Group.objects.get(name='Group A')
+                group.name = 'Group B'
+                group.save()
+                group.permissions.add(permission)
+                self.assertTrue(Group.objects.filter(name='Group B').exists())
+                # Query the through table directly: NetBox's users.Group also has a permissions relation,
+                # and group.permissions.filter() joins users_group_permissions instead.
+                self.assertTrue(
+                    GroupPermissions.objects.using(group._state.db).filter(group=group, permission=permission).exists()
+                )
+
+        self.assertTrue(Group.objects.filter(name='Group A').exists())
+        self.assertFalse(Group.objects.filter(name='Group B').exists())
+        self.assertFalse(GroupPermissions.objects.filter(group__name='Group A').exists())

netbox_branching/utilities.py

@@ -332,6 +332,35 @@ def get_tables_to_replicate():
                 m2m_table = m2m_field._get_m2m_db_table(model._meta)
             tables.add(m2m_table)
 
+    # A model that a plugin's resolver makes branchable without change logging has no branching
+    # feature, but the router sends its queries to the branch, so the branch needs its table too.
+    if _branching_resolvers:
+        from django.apps import apps
+        from django.db import connection
+
+        from netbox.models.features import ChangeLoggingMixin
+
+        existing_tables = set(connection.introspection.table_names())
+        for model in apps.get_models(include_auto_created=True):
+            meta = model._meta
+            if (
+                meta.db_table not in tables
+                and meta.db_table in existing_tables
+                and meta.managed
+                and not meta.proxy
+                and not issubclass(model, ChangeLoggingMixin)
+                and supports_branching(model)
+            ):
+                tables.add(meta.db_table)
+
+                # M2M writes follow the instance's database, so the branch needs the through tables too
+                for m2m_field in meta.local_many_to_many:
+                    if hasattr(m2m_field, 'through'):
+                        m2m_table = m2m_field.through._meta.db_table
+                    else:
+                        m2m_table = m2m_field._get_m2m_db_table(meta)
+                    tables.add(m2m_table)
+
     return sorted(tables)