Shard: Add support for Authorization header #518

promgen/migrations/0023_shard_authorization.py

@@ -0,0 +1,23 @@
+# Generated by Django 4.2.11 on 2024-06-28 09:57
+
+import django.core.validators
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("promgen", "0022_rule_labels_annotations"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="shard",
+            name="authorization",
+            field=models.CharField(
+                blank=True,
+                help_text="HTTP Authorization header for this shard's API",
+                max_length=4083,
+                null=True,
+            ),
+        ),
+    ]

promgen/models.py

@@ -172,6 +172,12 @@ class Meta:
 class Shard(models.Model):
     name = models.CharField(max_length=128, unique=True, validators=[validators.labelvalue])
     url = models.URLField(max_length=256)
+    authorization = models.CharField(
+        max_length=4083,  # 4KB - len("authorization")
+        blank=True,
+        null=True,
+        help_text="HTTP Authorization header for this shard's API",
+    )
     proxy = models.BooleanField(
         default=False,
         help_text="Queries can be proxied to these shards",

promgen/proxy.py

@@ -39,7 +39,7 @@ class PrometheusProxy(View):
     proxy_headers = {"HTTP_REFERER": "Referer"}
 
     @property
-    def headers(self):
+    def request_headers(self):
         # Loop through the headers from our request, and decide which ones
         # we should pass through upstream. Currently, our 'Referer' header is
         # the main one we are interested in, since this can help us debug which
@@ -53,12 +53,15 @@ def headers(self):
     def proxy(self, request):
         futures = []
         with concurrent.futures.ThreadPoolExecutor(max_workers=20) as executor:
-            for host in models.Shard.objects.filter(proxy=True):
+            for shard in models.Shard.objects.filter(proxy=True):
+                headers = self.request_headers
+                if shard.authorization:
+                    headers["Authorization"] = shard.authorization
                 futures.append(
                     executor.submit(
                         util.get,
-                        urljoin(host.url, request.get_full_path_info()),
-                        headers=self.headers,
+                        urljoin(shard.url, request.get_full_path_info()),
+                        headers=headers,
                     )
                 )
             yield from concurrent.futures.as_completed(futures)

promgen/static/js/promgen.vue.js

@@ -234,7 +234,7 @@ app.component('silence-modal', {
 
 app.component("promql-query", {
     delimiters: ['[[', ']]'],
-    props: ["href", "query", "max"],
+    props: ["shard", "query", "max"],
     data: function () {
         return {
             count: 0,
@@ -256,9 +256,11 @@ app.component("promql-query", {
     },
     template: '#promql-query-template',
     mounted() {
-        var url = new URL(this.href);
-        url.search = new URLSearchParams({ query: this.query });
-        fetch(url)
+        const params = new URLSearchParams({
+            shard: this.shard,
+            query: this.query,
+        });
+        fetch(`/promql-query?${params}`)
             .then(response => response.json())
             .then(result => this.count = Number.parseInt(result.data.result[0].value[1]))
             .finally(() => this.ready = true);

promgen/templates/promgen/project_form.html

@@ -87,10 +87,10 @@ <h1>Register new Project</h1>
         </td>
         {% if shard.proxy %}
         <td>
-          <promql-query class="label" href='{{shard.url}}/api/v1/query' query='sum(scrape_samples_scraped)' max="{{shard.samples}}">Samples: </promql-query>
+          <promql-query class="label" shard="{{shard.id}}" query="sum(scrape_samples_scraped)" max="{{shard.samples}}">Samples: </promql-query>
         </td>
         <td>
-          <promql-query class="label" href='{{shard.url}}/api/v1/query' query='count(up)' max="{{shard.targets}}">Exporters: </promql-query>
+          <promql-query class="label" shard="{{shard.id}}" query="count(up)" max="{{shard.targets}}">Exporters: </promql-query>
         </td>
         {% else %}
         <td>&nbsp;</td>

promgen/templates/promgen/shard_header.html

@@ -1,8 +1,8 @@
 {% load i18n %}
 <div class="well">
   {% if shard.proxy %}
-  <promql-query class="label" href='{{shard.url}}/api/v1/query' query='sum(scrape_samples_scraped)' max="{{shard.samples}}">Samples: </promql-query>
-  <promql-query class="label" href='{{shard.url}}/api/v1/query' query='count(up)' max="{{shard.targets}}">Exporters: </promql-query>
+  <promql-query class="label" shard="{{shard.id}}" query="sum(scrape_samples_scraped)" max="{{shard.samples}}">Samples: </promql-query>
+  <promql-query class="label" shard="{{shard.id}}" query="count(up)" max="{{shard.targets}}">Exporters: </promql-query>
   {% endif %}
 
   {% if user.is_superuser %}

promgen/urls.py

@@ -126,6 +126,8 @@
     path("proxy/v1/silences/<silence_id>", csrf_exempt(proxy.ProxyDeleteSilence.as_view()), name="proxy-silence-delete"),
     # Promgen rest API
     path("rest/", include((router.urls, "api"), namespace="api")),
+    # PromQL Query
+    path("promql-query", views.PromqlQuery.as_view(), name="promql-query"),
 ]
 
 try:

promgen/views.py

@@ -28,6 +28,7 @@
 from django.views.generic.detail import SingleObjectMixin
 from django.views.generic.edit import CreateView, DeleteView, FormView
 from prometheus_client.core import CounterMetricFamily, GaugeMetricFamily
+from requests.exceptions import HTTPError
 
 import promgen.templatetags.promgen as macro
 from promgen import (
@@ -1374,3 +1375,28 @@ def post(self, request, pk):
         return JsonResponse(
             {request.POST["target"]: render_to_string("promgen/ajax_clause_check.html", result)}
         )
+
+
+class PromqlQuery(View):
+    def get(self, request):
+        if not all(x in request.GET for x in ["shard", "query"]):
+            return HttpResponse("BAD REQUEST", status=400)
+
+        shard = get_object_or_404(models.Shard, pk=request.GET["shard"])
+        params = {"query": request.GET["query"]}
+        headers = {}
+
+        if shard.authorization:
+            headers["Authorization"] = shard.authorization
+
+        try:
+            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 HttpResponse(response.content, content_type="application/json")