Closes #262: Copy migrations table to branch (#263)

* Add migration to copy table to existing branches

* Extend provision() method on Branch to copy migrations table

* Logging tweak

* Fix test

Author: jeremystretch

netbox_branching/migrations/0004_copy_migrations.py

@@ -0,0 +1,52 @@
+from django.db import connection, migrations
+
+from netbox.plugins import get_plugin_config
+from netbox_branching.choices import BranchStatusChoices
+
+
+def copy_migrations(apps, schema_editor):
+    """
+    Create a copy of the migrations table in each active branch.
+    """
+    Branch = apps.get_model('netbox_branching', 'Branch')
+
+    table = 'django_migrations'
+    schema_prefix = get_plugin_config('netbox_branching', 'schema_prefix')
+
+    with connection.cursor() as cursor:
+        main_table = f'public.{table}'
+
+        for branch in Branch.objects.filter(status=BranchStatusChoices.READY):
+            print(f'\n    Copying migrations for branch {branch.name} ({branch.schema_id})...', end='')
+            schema_name = f'{schema_prefix}{branch.schema_id}'
+            schema_table = f'{schema_name}.{table}'
+
+            # Copy the migrations table to the branch schema
+            cursor.execute(f"CREATE TABLE {schema_table} ( LIKE {main_table} INCLUDING INDEXES )")
+
+            # Copy table data
+            cursor.execute(f"INSERT INTO {schema_table} SELECT * FROM {main_table}")
+
+            # Designate id as an identity column
+            cursor.execute(f"ALTER TABLE {schema_table} ALTER COLUMN id ADD GENERATED BY DEFAULT AS IDENTITY")
+
+            # Set the next value for the ID sequence
+            cursor.execute("SELECT MAX(id) from django_migrations")
+            starting_id = cursor.fetchone()[0] + 1
+            cursor.execute(f"ALTER SEQUENCE {schema_name}.django_migrations_id_seq RESTART WITH {starting_id}")
+
+    print('\n ', end='')  # Padding for final "OK"
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('netbox_branching', '0003_rename_indexes'),
+    ]
+
+    operations = [
+        migrations.RunPython(
+            code=copy_migrations,
+            reverse_code=migrations.RunPython.noop
+        ),
+    ]

netbox_branching/models/branches.py

@@ -540,6 +540,29 @@ def provision(self, user):
                     f"ALTER TABLE {schema_table} ALTER COLUMN id SET DEFAULT nextval(%s)", [sequence_name]
                 )
 
+                # Copy the migrations table
+                main_table = f'{MAIN_SCHEMA}.django_migrations'
+                schema_table = f'{schema}.django_migrations'
+                logger.debug(f'Creating table {schema_table}')
+                cursor.execute(
+                    f"CREATE TABLE {schema_table} ( LIKE {main_table} INCLUDING INDEXES )"
+                )
+                cursor.execute(
+                    f"INSERT INTO {schema_table} SELECT * FROM {main_table}"
+                )
+                # Designate id as an identity column
+                cursor.execute(
+                    f"ALTER TABLE {schema_table} ALTER COLUMN id ADD GENERATED BY DEFAULT AS IDENTITY"
+                )
+                # Set the next value for the ID sequence
+                cursor.execute(
+                    f"SELECT MAX(id) from {schema_table}"
+                )
+                starting_id = cursor.fetchone()[0] + 1
+                cursor.execute(
+                    f"ALTER SEQUENCE {schema}.django_migrations_id_seq RESTART WITH {starting_id}"
+                )
+
                 # Replicate relevant tables from the main schema
                 for table in get_tables_to_replicate():
                     main_table = f'{MAIN_SCHEMA}.{table}'

netbox_branching/tests/test_branches.py

@@ -38,7 +38,7 @@ def test_create_branch(self):
                 "SELECT * FROM information_schema.tables WHERE table_schema=%s",
                 [branch.schema_name]
             )
-            tables_expected = {*tables_to_replicate, 'core_objectchange'}
+            tables_expected = {*tables_to_replicate, 'core_objectchange', 'django_migrations'}
             tables_found = {row.table_name for row in fetchall(cursor)}
             self.assertSetEqual(tables_expected, tables_found)