Closes #181: Introduce a configuration parameter to define the default schema name (#254)

Author: jeremystretch

docs/configuration.md

@@ -28,6 +28,14 @@ exempt_models = (
 
 ---
 
+## `main_schema`
+
+Default: `"public"`
+
+The name of the main (primary) PostgreSQL schema. (Use the `\dn` command in the PostgreSQL CLI to list all schemas.)
+
+---
+
 ## `max_working_branches`
 
 Default: None
@@ -46,7 +54,7 @@ The maximum total number of branches that can exist simultaneously, including me
 
 ## `schema_prefix`
 
-Default: `branch_`
+Default: `"branch_"`
 
 The string to prefix to the unique branch ID when provisioning the PostgreSQL schema for a branch. Per [the PostgreSQL documentation](https://www.postgresql.org/docs/16/sql-syntax-lexical.html#SQL-SYNTAX-IDENTIFIERS), this string must begin with a letter or underscore.
 

netbox_branching/__init__.py

@@ -28,6 +28,9 @@ class AppConfig(PluginConfig):
         # Models from other plugins which should be excluded from branching support
         'exempt_models': [],
 
+        # The name of the main schema
+        'main_schema': 'public',
+
         # This string is prefixed to the name of each new branch schema during provisioning
         'schema_prefix': 'branch_',
 

netbox_branching/constants.py

@@ -1,6 +1,3 @@
-# Name of the main (non-branch) PostgreSQL schema
-MAIN_SCHEMA = 'public'
-
 # HTTP cookie
 COOKIE_NAME = 'active_branch'
 

netbox_branching/database.py

@@ -37,5 +37,5 @@ def db_for_write(self, model, **hints):
         return self._get_db(model, **hints)
 
     def allow_relation(self, obj1, obj2, **hints):
-        # Permit relations from the branch schema to the main (public) schema
+        # Permit relations from the branch schema to the main schema
         return True

netbox_branching/models/branches.py

@@ -540,6 +540,7 @@ def provision(self, user):
         """
         logger = logging.getLogger('netbox_branching.branch.provision')
         logger.info(f'Provisioning branch {self} ({self.schema_name})')
+        main_schema = get_plugin_config('netbox_branching', 'main_schema')
 
         # Emit pre-provision signal
         pre_provision.send(sender=self.__class__, branch=self, user=user)
@@ -572,21 +573,21 @@ def provision(self, user):
                 # Create an empty copy of the global change log. Share the ID sequence from the main table to avoid
                 # reusing change record IDs.
                 table = ObjectChange_._meta.db_table
-                main_table = f'public.{table}'
+                main_table = f'{main_schema}.{table}'
                 schema_table = f'{schema}.{table}'
                 logger.debug(f'Creating table {schema_table}')
                 cursor.execute(
                     f"CREATE TABLE {schema_table} ( LIKE {main_table} INCLUDING INDEXES )"
                 )
                 # Set the default value for the ID column to the sequence associated with the source table
-                sequence_name = f'public.{table}_id_seq'
+                sequence_name = f'{main_schema}.{table}_id_seq'
                 cursor.execute(
                     f"ALTER TABLE {schema_table} ALTER COLUMN id SET DEFAULT nextval(%s)", [sequence_name]
                 )
 
                 # Replicate relevant tables from the main schema
                 for table in get_tables_to_replicate():
-                    main_table = f'public.{table}'
+                    main_table = f'{main_schema}.{table}'
                     schema_table = f'{schema}.{table}'
                     logger.debug(f'Creating table {schema_table}')
                     # Create the table in the new schema

netbox_branching/tests/test_branches.py

@@ -6,8 +6,8 @@
 from django.test import TransactionTestCase, override_settings
 from django.utils import timezone
 
+from netbox.plugins import get_plugin_config
 from netbox_branching.choices import BranchStatusChoices
-from netbox_branching.constants import MAIN_SCHEMA
 from netbox_branching.models import Branch
 from netbox_branching.utilities import get_tables_to_replicate
 from .utils import fetchall, fetchone
@@ -21,6 +21,7 @@ def test_create_branch(self):
         branch.save(provision=False)
         branch.provision(user=None)
 
+        main_schema = get_plugin_config('netbox_branching', 'main_schema')
         tables_to_replicate = get_tables_to_replicate()
 
         with connection.cursor() as cursor:
@@ -44,7 +45,7 @@ def test_create_branch(self):
 
             # Check that object counts match the main schema for each table
             for table_name in tables_to_replicate:
-                cursor.execute(f"SELECT COUNT(id) FROM {MAIN_SCHEMA}.{table_name}")
+                cursor.execute(f"SELECT COUNT(id) FROM {main_schema}.{table_name}")
                 main_count = fetchone(cursor).count
                 cursor.execute(f"SELECT COUNT(id) FROM {branch.schema_name}.{table_name}")
                 branch_count = fetchone(cursor).count

netbox_branching/utilities.py

@@ -3,6 +3,7 @@
 from collections import defaultdict
 from contextlib import contextmanager, nullcontext
 from dataclasses import dataclass
+from functools import cached_property
 
 from django.contrib import messages
 from django.db.models import ForeignKey, ManyToManyField
@@ -39,14 +40,18 @@ class DynamicSchemaDict(dict):
     Behaves like a normal dictionary, except for keys beginning with "schema_". Any lookup for
     "schema_*" will return the default configuration extended to include the search_path option.
     """
+    @cached_property
+    def main_schema(self):
+        return get_plugin_config('netbox_branching', 'main_schema')
+
     def __getitem__(self, item):
         if type(item) is str and item.startswith('schema_'):
             if schema := item.removeprefix('schema_'):
                 default_config = super().__getitem__('default')
                 return {
                     **default_config,
                     'OPTIONS': {
-                        'options': f'-c search_path={schema},public'
+                        'options': f'-c search_path={schema},{self.main_schema}'
                     }
                 }
         return super().__getitem__(item)