diff --git a/netbox_branching/migrations/0004_copy_migrations.py b/netbox_branching/migrations/0004_copy_migrations.py new file mode 100644 index 0000000..3de6a3e --- /dev/null +++ b/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 + ), + ] diff --git a/netbox_branching/models/branches.py b/netbox_branching/models/branches.py index 964f777..ea16a5e 100644 --- a/netbox_branching/models/branches.py +++ b/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}' diff --git a/netbox_branching/tests/test_branches.py b/netbox_branching/tests/test_branches.py index 804ba0a..a7936d0 100644 --- a/netbox_branching/tests/test_branches.py +++ b/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)