From 412c08230cb64025a22e1f66d03eae4ae1d8b472 Mon Sep 17 00:00:00 2001
From: Sjoerd Job Postmus <sjoerdjob@sjec.nl>
Date: Mon, 29 Jul 2019 16:51:04 +0200
Subject: [PATCH] Allow overriding context for Django Middleware.

By subclassing `django.HoneyMiddlewareBase`, it is now possible to
change the context that gets sent to Honeycomb.

See #72.
---
 beeline/middleware/django/__init__.py | 36 +++++++++++++++++----------
 1 file changed, 23 insertions(+), 13 deletions(-)

diff --git a/beeline/middleware/django/__init__.py b/beeline/middleware/django/__init__.py
index 3d819ae..bfbb984 100644
--- a/beeline/middleware/django/__init__.py
+++ b/beeline/middleware/django/__init__.py
@@ -55,14 +55,9 @@ def __call__(self, request):
         response = self.create_http_event(request)
         return response
 
-    def create_http_event(self, request):
-        # Code to be executed for each request before
-        # the view (and later middleware) are called.
-
-        trace_id, parent_id, context = _get_trace_context(request)
+    def get_context_from_request(self, request):
         trace_name = "django_http_%s" % request.method.lower()
-
-        trace = beeline.start_trace(context={
+        return {
             "name": trace_name,
             "type": "http_server",
             "request.host": request.get_host(),
@@ -75,19 +70,34 @@ def create_http_event(self, request):
             "request.secure": request.is_secure(),
             "request.query": request.GET.dict(),
             "request.xhr": request.is_ajax(),
-            "request.post": request.POST.dict()
-        }, trace_id=trace_id, parent_span_id=parent_id)
+            "request.post": request.POST.dict(),
+        }
+
+    def get_context_from_response(self, response):
+        return {
+            "response.status_code": response.status_code,
+        }
+
+    def create_http_event(self, request):
+        # Code to be executed for each request before
+        # the view (and later middleware) are called.
+
+        trace_id, parent_id, parent_context = _get_trace_context(request)
 
-        if isinstance(context, dict):
-            for k, v in context.items():
+        request_context = self.get_context_from_request(request)
+
+        trace = beeline.start_trace(context=request_context, trace_id=trace_id, parent_span_id=parent_id)
+
+        if isinstance(parent_context, dict):
+            for k, v in parent_context.items():
                 beeline.add_trace_field(k, v)
 
         response = self.get_response(request)
 
         # Code to be executed for each request/response after
         # the view is called.
-
-        beeline.add_context_field("response.status_code", response.status_code)
+        response_context = self.get_context_from_response(response)
+        beeline.add_context(response_context)
         beeline.finish_trace(trace)
 
         return response