Misc cleanup

netbox/extras/choices.py

@@ -195,7 +195,7 @@ class ChangeActionChoices(ChoiceSet):
     ACTION_DELETE = 'delete'
 
     CHOICES = (
-        (ACTION_CREATE, 'Created'),
+        (ACTION_CREATE, 'Create'),
-        (ACTION_UPDATE, 'Updated'),
+        (ACTION_UPDATE, 'Update'),
-        (ACTION_DELETE, 'Deleted'),
+        (ACTION_DELETE, 'Delete'),
     )

netbox/extras/models/staging.py

@@ -84,21 +84,28 @@ class Change(ChangeLoggedModel):
             ),
         )
 
-    def apply(self):
+    def __str__(self):
-        model = self.object_type.model_class()
+        return f"{self.get_action_display()} {self.model}"
-        pk = self.object_id
 
+    @property
+    def model(self):
+        return self.object_type.model_class()
+
+    def apply(self):
+        """
+        Apply the staged create/update/delete action to the database.
+        """
         if self.action == ChangeActionChoices.ACTION_CREATE:
-            instance = deserialize_object(model, self.data, pk=pk)
+            instance = deserialize_object(self.model, self.data, pk=self.object_id)
-            logger.info(f'Creating {model._meta.verbose_name} {instance}')
+            logger.info(f'Creating {self.model._meta.verbose_name} {instance}')
             instance.save()
 
         if self.action == ChangeActionChoices.ACTION_UPDATE:
-            instance = deserialize_object(model, self.data, pk=pk)
+            instance = deserialize_object(self.model, self.data, pk=self.object_id)
-            logger.info(f'Updating {model._meta.verbose_name} {instance}')
+            logger.info(f'Updating {self.model._meta.verbose_name} {instance}')
             instance.save()
 
         if self.action == ChangeActionChoices.ACTION_DELETE:
-            instance = model.objects.get(pk=self.object_id)
+            instance = self.model.objects.get(pk=self.object_id)
-            logger.info(f'Deleting {model._meta.verbose_name} {instance}')
+            logger.info(f'Deleting {self.model._meta.verbose_name} {instance}')
             instance.delete()

netbox/netbox/staging.py

@@ -6,28 +6,24 @@ from django.db.models.signals import pre_delete, post_save
 
 from extras.choices import ChangeActionChoices
 from extras.models import Change
-from utilities.utils import serialize_object, shallow_compare_dict
+from utilities.utils import serialize_object
 
 logger = logging.getLogger('netbox.staging')
 
 
-def get_changed_fields(instance):
-    model = instance._meta.model
-    original = model.objects.get(pk=instance.pk)
-    return shallow_compare_dict(
-        serialize_object(original),
-        serialize_object(instance),
-        exclude=('last_updated',)
-    )
-
-
-def get_key_for_instance(instance):
-    object_type = ContentType.objects.get_for_model(instance)
-    return object_type, instance.pk
-
-
 class checkout:
+    """
+    Context manager for staging changes to NetBox objects. Staged changes are saved out-of-band
+    (as Change instances) for application at a later time, without modifying the production
+    database.
 
+        branch = Branch.objects.create(name='my-branch')
+        with checkout(branch):
+            # All changes made herein will be rolled back and stored for later
+
+    Note that invoking the context disabled transaction autocommit to facilitate manual rollbacks,
+    and restores its original value upon exit.
+    """
     def __init__(self, branch):
         self.branch = branch
         self.queue = {}
@@ -37,13 +33,12 @@ class checkout:
         # Disable autocommit to effect a new transaction
         logger.debug(f"Entering transaction for {self.branch}")
         self._autocommit = transaction.get_autocommit()
-
         transaction.set_autocommit(False)
 
         # Apply any existing Changes assigned to this Branch
         changes = self.branch.changes.all()
-        if changes.exists():
+        if change_count := changes.count():
-            logger.debug(f"Applying {changes.count()} pre-staged changes...")
+            logger.debug(f"Applying {change_count} pre-staged changes...")
             for change in changes:
                 change.apply()
         else:
@@ -56,42 +51,63 @@ class checkout:
 
     def __exit__(self, exc_type, exc_val, exc_tb):
 
-        # Roll back the transaction to return the database to its original state
-        logger.debug("Rolling back transaction")
-        transaction.rollback()
-        logger.debug(f"Restoring autocommit state {self._autocommit}")
-        transaction.set_autocommit(self._autocommit)
-
         # Disconnect signal handlers
         logger.debug("Disconnecting signal handlers")
         post_save.disconnect(self.post_save_handler)
         pre_delete.disconnect(self.pre_delete_handler)
 
+        # Roll back the transaction to return the database to its original state
+        logger.debug("Rolling back database transaction")
+        transaction.rollback()
+        logger.debug(f"Restoring autocommit state ({self._autocommit})")
+        transaction.set_autocommit(self._autocommit)
+
         # Process queued changes
+        self.process_queue()
+
+    def process_queue(self):
+        """
+        Create Change instances for all actions stored in the queue.
+        """
         changes = []
-        logger.debug(f"Processing {len(self.queue)} queued changes:")
+        if not self.queue:
+            logger.debug(f"No queued changes; aborting")
+            return
+        logger.debug(f"Processing {len(self.queue)} queued changes")
+
+        # Iterate through the in-memory queue, creating Change instances
         for key, change in self.queue.items():
             logger.debug(f'  {key}: {change}')
             object_type, pk = key
             action, instance = change
+            data = None
             if action in (ChangeActionChoices.ACTION_CREATE, ChangeActionChoices.ACTION_UPDATE):
                 data = serialize_object(instance)
-            else:
-                data = None
 
-            change = Change(
+            changes.append(Change(
                 branch=self.branch,
                 action=action,
                 object_type=object_type,
                 object_id=pk,
                 data=data
-            )
+            ))
-            changes.append(change)
 
+        # Save all Change instances to the database
         Change.objects.bulk_create(changes)
 
+    @staticmethod
+    def get_key_for_instance(instance):
+        return ContentType.objects.get_for_model(instance), instance.pk
+
+    #
+    # Signal handlers
+    #
+
     def post_save_handler(self, sender, instance, created, **kwargs):
-        key = get_key_for_instance(instance)
+        """
+        Hooks to the post_save signal when a branch is active to queue create and update actions.
+        """
+        key = self.get_key_for_instance(instance)
         object_type = instance._meta.verbose_name
 
         if created:
@@ -108,7 +124,10 @@ class checkout:
             self.queue[key] = (ChangeActionChoices.ACTION_UPDATE, instance)
 
     def pre_delete_handler(self, sender, instance, **kwargs):
-        key = get_key_for_instance(instance)
+        """
+        Hooks to the pre_delete signal when a branch is active to queue delete actions.
+        """
+        key = self.get_key_for_instance(instance)
         if key in self.queue and self.queue[key][0] == 'create':
             # Cancel the creation of a new object
             logger.debug(f"[{self.branch}] Removing staged deletion of {instance} (PK: {instance.pk})")