diff --git a/netbox_branching/migrations/0006_tag_object_types.py b/netbox_branching/migrations/0006_tag_object_types.py new file mode 100644 index 0000000..9f557e4 --- /dev/null +++ b/netbox_branching/migrations/0006_tag_object_types.py @@ -0,0 +1,73 @@ +from django.db import connection, migrations + +from netbox.plugins import get_plugin_config +from netbox_branching.choices import BranchStatusChoices + + +def copy_table(apps, schema_editor): + """ + Create a copy of the extras_tag_object_types table in each active branch. + """ + Branch = apps.get_model('netbox_branching', 'Branch') + + table = 'extras_tag_object_types' + 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 {table} for branch {branch.name} ({branch.schema_id})... ', end='') + schema_name = f'{schema_prefix}{branch.schema_id}' + schema_table = f'{schema_name}.{table}' + + # Abort if the table already exists (somehow) + cursor.execute( + f"SELECT EXISTS (SELECT 1 FROM pg_tables WHERE schemaname='{schema_name}' AND tablename='{table}')" + ) + if cursor.fetchone()[0]: + print('Skipping; table already exists.', end='') + continue + + # Copy the extras_tag_object_types 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}") + + # Set the default value for the ID column to the sequence associated with the source table + cursor.execute( + f"ALTER TABLE {schema_table} ALTER COLUMN id SET DEFAULT nextval('extras_tag_object_types_id_seq')" + ) + + # Rename indexes + cursor.execute( + f"ALTER INDEX {schema_name}.extras_tag_object_types_objecttype_id_idx " + f"RENAME TO extras_tag_object_types_contenttype_id_c1b220c3" + ) + cursor.execute( + f"ALTER INDEX {schema_name}.extras_tag_object_types_tag_id_idx " + f"RENAME TO extras_tag_object_types_tag_id_2e1aab29" + ) + cursor.execute( + f"ALTER INDEX {schema_name}.extras_tag_object_types_tag_id_objecttype_id_key " + f"RENAME TO extras_tag_object_types_tag_id_contenttype_id_2ff9910c_uniq" + ) + + print('Success.', end='') + + print('\n ', end='') # Padding for final "OK" + + +class Migration(migrations.Migration): + + dependencies = [ + ('netbox_branching', '0005_branch_applied_migrations'), + ] + + operations = [ + migrations.RunPython( + code=copy_table, + reverse_code=migrations.RunPython.noop + ), + ] diff --git a/netbox_branching/utilities.py b/netbox_branching/utilities.py index f7ca10e..a78a2a6 100644 --- a/netbox_branching/utilities.py +++ b/netbox_branching/utilities.py @@ -144,13 +144,12 @@ def get_tables_to_replicate(): # Capture any M2M fields which reference other replicated models for m2m_field in model._meta.local_many_to_many: - if m2m_field.related_model in branch_aware_models: - if hasattr(m2m_field, 'through'): - # Field is actually a manager - m2m_table = m2m_field.through._meta.db_table - else: - m2m_table = m2m_field._get_m2m_db_table(model._meta) - tables.add(m2m_table) + if hasattr(m2m_field, 'through'): + # Field is actually a manager + m2m_table = m2m_field.through._meta.db_table + else: + m2m_table = m2m_field._get_m2m_db_table(model._meta) + tables.add(m2m_table) return sorted(tables)