Move request object and webhook queue to TLS

Author: jeremystretch

netbox/extras/context_managers.py

@@ -2,8 +2,9 @@
 
 from django.db.models.signals import m2m_changed, pre_delete, post_save
 
-from extras.signals import clear_webhooks, _clear_webhook_queue, _handle_changed_object, _handle_deleted_object
-from utilities.utils import curry
+from extras.signals import clear_webhooks, clear_webhook_queue, handle_changed_object, handle_deleted_object
+from netbox import thread_locals
+from netbox.request_context import set_request
 from .webhooks import flush_webhooks
 
 
@@ -15,12 +16,8 @@ def change_logging(request):
 
     :param request: WSGIRequest object with a unique `id` set
     """
-    webhook_queue = []
-
-    # Curry signals receivers to pass the current request
-    handle_changed_object = curry(_handle_changed_object, request, webhook_queue)
-    handle_deleted_object = curry(_handle_deleted_object, request, webhook_queue)
-    clear_webhook_queue = curry(_clear_webhook_queue, webhook_queue)
+    set_request(request)
+    thread_locals.webhook_queue = []
 
     # Connect our receivers to the post_save and post_delete signals.
     post_save.connect(handle_changed_object, dispatch_uid='handle_changed_object')
@@ -38,5 +35,8 @@ def change_logging(request):
     clear_webhooks.disconnect(clear_webhook_queue, dispatch_uid='clear_webhook_queue')
 
     # Flush queued webhooks to RQ
-    flush_webhooks(webhook_queue)
-    del webhook_queue
+    flush_webhooks(thread_locals.webhook_queue)
+    del thread_locals.webhook_queue
+
+    # Clear the request from thread-local storage
+    set_request(None)

netbox/extras/signals.py

@@ -6,6 +6,8 @@
 from django.dispatch import receiver, Signal
 from django_prometheus.models import model_deletes, model_inserts, model_updates
 
+from netbox import thread_locals
+from netbox.request_context import get_request
 from netbox.signals import post_clean
 from .choices import ObjectChangeActionChoices
 from .models import CustomField, ObjectChange
@@ -20,22 +22,23 @@
 clear_webhooks = Signal()
 
 
-def _handle_changed_object(request, webhook_queue, sender, instance, **kwargs):
+def handle_changed_object(sender, instance, **kwargs):
     """
     Fires when an object is created or updated.
     """
+    if not hasattr(instance, 'to_objectchange'):
+        return
+
+    request = get_request()
+    m2m_changed = False
+
     def is_same_object(instance, webhook_data):
         return (
             ContentType.objects.get_for_model(instance) == webhook_data['content_type'] and
             instance.pk == webhook_data['object_id'] and
             request.id == webhook_data['request_id']
         )
 
-    if not hasattr(instance, 'to_objectchange'):
-        return
-
-    m2m_changed = False
-
     # Determine the type of change being made
     if kwargs.get('created'):
         action = ObjectChangeActionChoices.ACTION_CREATE
@@ -65,6 +68,7 @@ def is_same_object(instance, webhook_data):
             objectchange.save()
 
     # If this is an M2M change, update the previously queued webhook (from post_save)
+    webhook_queue = thread_locals.webhook_queue
     if m2m_changed and webhook_queue and is_same_object(instance, webhook_queue[-1]):
         instance.refresh_from_db()  # Ensure that we're working with fresh M2M assignments
         webhook_queue[-1]['data'] = serialize_for_webhook(instance)
@@ -79,13 +83,15 @@ def is_same_object(instance, webhook_data):
         model_updates.labels(instance._meta.model_name).inc()
 
 
-def _handle_deleted_object(request, webhook_queue, sender, instance, **kwargs):
+def handle_deleted_object(sender, instance, **kwargs):
     """
     Fires when an object is deleted.
     """
     if not hasattr(instance, 'to_objectchange'):
         return
 
+    request = get_request()
+
     # Record an ObjectChange if applicable
     if hasattr(instance, 'to_objectchange'):
         objectchange = instance.to_objectchange(ObjectChangeActionChoices.ACTION_DELETE)
@@ -94,19 +100,21 @@ def _handle_deleted_object(request, webhook_queue, sender, instance, **kwargs):
         objectchange.save()
 
     # Enqueue webhooks
+    webhook_queue = thread_locals.webhook_queue
     enqueue_object(webhook_queue, instance, request.user, request.id, ObjectChangeActionChoices.ACTION_DELETE)
 
     # Increment metric counters
     model_deletes.labels(instance._meta.model_name).inc()
 
 
-def _clear_webhook_queue(webhook_queue, sender, **kwargs):
+def clear_webhook_queue(sender, **kwargs):
     """
     Delete any queued webhooks (e.g. because of an aborted bulk transaction)
     """
     logger = logging.getLogger('webhooks')
-    logger.info(f"Clearing {len(webhook_queue)} queued webhooks ({sender})")
+    webhook_queue = thread_locals.webhook_queue
 
+    logger.info(f"Clearing {len(webhook_queue)} queued webhooks ({sender})")
     webhook_queue.clear()
 
 

netbox/netbox/__init__.py

@@ -0,0 +1,3 @@
+import threading
+
+thread_locals = threading.local()

netbox/netbox/middleware.py

@@ -1,10 +1,10 @@
+import logging
 import uuid
 from urllib import parse
-import logging
 
 from django.conf import settings
-from django.contrib.auth.middleware import RemoteUserMiddleware as RemoteUserMiddleware_
 from django.contrib import auth
+from django.contrib.auth.middleware import RemoteUserMiddleware as RemoteUserMiddleware_
 from django.core.exceptions import ImproperlyConfigured
 from django.db import ProgrammingError
 from django.http import Http404, HttpResponseRedirect
@@ -114,7 +114,7 @@ def _get_groups(self, request):
         return groups
 
 
-class ObjectChangeMiddleware(object):
+class ObjectChangeMiddleware:
     """
     This middleware performs three functions in response to an object being created, updated, or deleted:
 

netbox/netbox/request_context.py

@@ -0,0 +1,9 @@
+from netbox import thread_locals
+
+
+def set_request(request):
+    thread_locals.request = request
+
+
+def get_request():
+    return getattr(thread_locals, 'request', None)

netbox/utilities/utils.py

@@ -327,13 +327,6 @@ def decode_value(value: Any, _decode_keys: bool) -> Any:
     return {urllib.parse.unquote(k): decode_value(v, decode_keys) for k, v in encoded_dict.items()}
 
 
-# Taken from django.utils.functional (<3.0)
-def curry(_curried_func, *args, **kwargs):
-    def _curried(*moreargs, **morekwargs):
-        return _curried_func(*args, *moreargs, **{**kwargs, **morekwargs})
-    return _curried
-
-
 def array_to_string(array):
     """
     Generate an efficient, human-friendly string from a set of integers. Intended for use with ArrayField.