From fa655d162a186c353cb5d923ebe1bd805e3aa63d Mon Sep 17 00:00:00 2001
From: Paul Traylor <kungfudiscomonkey@gmail.com>
Date: Wed, 3 Jul 2024 16:55:33 +0900
Subject: [PATCH] Move proxy_error to util.proxy_error for reuse

We can move our proxy_error wrapper so that it can be reused in our
PromqlQuery View to keep things consistent.
---
 promgen/proxy.py | 27 +++++----------------------
 promgen/util.py  | 17 +++++++++++++++++
 promgen/views.py |  6 +-----
 3 files changed, 23 insertions(+), 27 deletions(-)

diff --git a/promgen/proxy.py b/promgen/proxy.py
index 6eaee9eb7..fb16bd88c 100644
--- a/promgen/proxy.py
+++ b/promgen/proxy.py
@@ -18,23 +18,6 @@
 logger = logging.getLogger(__name__)
 
 
-def proxy_error(response):
-    """
-    Return a wrapped proxy error
-
-    Taking a request.response object as input, return it slightly modified
-    with an extra header for debugging so that we can see where the request
-    failed
-    """
-    r = HttpResponse(
-        response.content,
-        content_type=response.headers["content-type"],
-        status=response.status_code,
-    )
-    r.setdefault("X-PROMGEN-PROXY", response.url)
-    return r
-
-
 class PrometheusProxy(View):
     proxy_headers = {"HTTP_REFERER": "Referer"}
 
@@ -91,7 +74,7 @@ def get(self, request):
                 data.update(_json["data"])
             except HTTPError:
                 logger.warning("Error with response")
-                return proxy_error(result)
+                return util.proxy_error(result)
 
         return JsonResponse({"status": "success", "data": sorted(data)})
 
@@ -108,7 +91,7 @@ def get(self, request, label):
                 data.update(_json["data"])
             except HTTPError:
                 logger.warning("Error with response")
-                return proxy_error(result)
+                return util.proxy_error(result)
 
         return JsonResponse({"status": "success", "data": sorted(data)})
 
@@ -125,7 +108,7 @@ def get(self, request):
                 data += _json["data"]
             except HTTPError:
                 logger.warning("Error with response")
-                return proxy_error(result)
+                return util.proxy_error(result)
 
         return JsonResponse({"status": "success", "data": data})
 
@@ -143,7 +126,7 @@ def get(self, request):
                 data += _json["data"]["result"]
                 resultType = _json["data"]["resultType"]
             except HTTPError:
-                return proxy_error(result)
+                return util.proxy_error(result)
 
         return JsonResponse(
             {"status": "success", "data": {"resultType": resultType, "result": data}}
@@ -163,7 +146,7 @@ def get(self, request):
                 data += _json["data"]["result"]
                 resultType = _json["data"]["resultType"]
             except HTTPError:
-                return proxy_error(result)
+                return util.proxy_error(result)
 
         return JsonResponse(
             {"status": "success", "data": {"resultType": resultType, "result": data}}
diff --git a/promgen/util.py b/promgen/util.py
index 20642d349..7e754f05e 100644
--- a/promgen/util.py
+++ b/promgen/util.py
@@ -7,6 +7,7 @@
 import requests
 from django.conf import settings
 from django.db.models import F
+from django.http import HttpResponse
 
 # Wrappers around request api to ensure we always attach our user agent
 # https://github.com/requests/requests/blob/master/requests/api.py
@@ -120,6 +121,22 @@ def wrapped(field):
 
     return wrapped
 
+def proxy_error(response: requests.Response) -> HttpResponse:
+    """
+    Return a wrapped proxy error
+
+    Taking a request.response object as input, return it slightly modified
+    with an extra header for debugging so that we can see where the request
+    failed
+    """
+    r = HttpResponse(
+        response.content,
+        content_type=response.headers["content-type"],
+        status=response.status_code,
+    )
+    r.setdefault("X-PROMGEN-PROXY", response.url)
+    return r
+
 
 # Comment wrappers to get the docstrings from the upstream functions
 get.__doc__ = requests.get.__doc__
diff --git a/promgen/views.py b/promgen/views.py
index 88dbd61a6..f92262587 100644
--- a/promgen/views.py
+++ b/promgen/views.py
@@ -1404,10 +1404,6 @@ def get(self, request):
             response = util.get(f"{shard.url}/api/v1/query", params=params, headers=headers)
             response.raise_for_status()
         except HTTPError:
-            return HttpResponse(
-                response.content,
-                content_type=response.headers["content-type"],
-                status=response.status_code,
-            )
+            return util.proxy_error(response)
 
         return HttpResponse(response.content, content_type="application/json")