Convert Rule label/annotation storage to JsonField

When Promgen was written, there was not a default JsonField, so labels
and annotations were implemented as a class. Now that JsonField has been
in Django since 3.1, it's a good time to change the implementation.

Previously, our audit log did *not* store change to labels/annotations
(because they were separate objects), so one benefit of this migration
is that those changes now show up in the edit log.

We also expect some queries to be a bit faster now that we do not need
to query the other table relations for labels/annotations.

promgen/admin.py

@@ -104,20 +104,11 @@ def has_add_permission(self, request):
     readonly_fields = ("project",)
 
 
-class RuleLabelInline(admin.TabularInline):
-    model = models.RuleLabel
-
-
-class RuleAnnotationInline(admin.TabularInline):
-    model = models.RuleAnnotation
-
-
 @admin.register(models.Rule)
 class RuleAdmin(admin.ModelAdmin):
     list_display = ("name", "clause", "duration", "content_object")
     list_filter = ("duration",)
     list_select_related = ("content_type",)
-    inlines = [RuleLabelInline, RuleAnnotationInline]
 
     def get_queryset(self, request):
         qs = super().get_queryset(request)

promgen/fixtures/extras.yaml

@@ -9,18 +9,10 @@
     duration: 1s
     content_type: ["promgen", "site"]
     object_id: 1
-- model: promgen.ruleannotation
-  pk: 1
-  fields:
-    name: summary
-    value: Example rule summary
-    rule_id: 1
-- model: promgen.rulelabel
-  pk: 1
-  fields:
-    name: severity
-    value: high
-    rule_id: 1
+    labels:
+      severity: "high"
+    annotations:
+      summary: "Example rule summary"
 - model: promgen.alert
   pk: 1
   fields:

promgen/forms.py

@@ -2,7 +2,7 @@
 # These sources are released under the terms of the MIT license: see LICENSE
 
 import re
-
+from functools import partial
 from dateutil import parser
 
 from promgen import models, plugins, prometheus, validators
@@ -119,14 +119,16 @@ class Meta:
 class AlertRuleForm(forms.ModelForm):
     class Meta:
         model = models.Rule
-        exclude = ["parent", "content_type", "object_id"]
         widgets = {
             "name": forms.TextInput(attrs={"class": "form-control"}),
             "duration": forms.TextInput(attrs={"class": "form-control"}),
             "clause": forms.Textarea(attrs={"rows": 5, "class": "form-control"}),
             "enabled": forms.CheckboxInput(attrs={"data-toggle": "toggle", "data-size": "mini"}),
             "description": forms.Textarea(attrs={"rows": 5, "class": "form-control"}),
         }
+        # We define a custom widget for each of our fields, so we just take the
+        # keys here to avoid manually updating a list of fields.
+        fields = widgets.keys()
 
     def clean(self):
         # Check our cleaned data then let Prometheus check our rule
@@ -142,6 +144,45 @@ def clean(self):
         prometheus.check_rules([rule])
 
 
+class _KeyValueForm(forms.Form):
+    key = forms.CharField(widget=forms.TextInput(attrs={"class": "form-control"}))
+    value = forms.CharField(widget=forms.TextInput(attrs={"class": "form-control"}))
+
+# We need a custom KeyValueSet because we need to be able to convert between the single dictionary
+# form saved to our models, and the list of models used by
+class _KeyValueSet(forms.BaseFormSet):
+    def __init__(self, initial=None, **kwargs):
+        if initial:
+            kwargs["initial"] = [{"key": key, "value": initial[key]} for key in initial]
+        super().__init__(**kwargs, form_kwargs={"empty_permitted": True})
+
+    def to_dict(self):
+        return {x["key"]: x["value"] for x in self.cleaned_data if x and not x["DELETE"]}
+
+# For both LabelFormSet and AnnotationFormSet we always want to have a prefix assigned, but it's
+# awkward if we need to specify it in multiple places. We use a partial here, so that it is the same
+# as always passing prefix as part of our __init__ call.
+LabelFormSet = partial(
+    forms.formset_factory(
+        form=_KeyValueForm,
+        formset=_KeyValueSet,
+        can_delete=True,
+        extra=1,
+    ),
+    prefix="labels",
+)
+
+AnnotationFormSet = partial(
+    forms.formset_factory(
+        form=_KeyValueForm,
+        formset=_KeyValueSet,
+        can_delete=True,
+        extra=1,
+    ),
+    prefix="annotations",
+)
+
+
 class RuleCopyForm(forms.Form):
     content_type = forms.ChoiceField(choices=[(x, x) for x in ["service", "project"]])
     object_id = forms.IntegerField()
@@ -185,25 +226,3 @@ def clean(self):
         if not hosts:
             raise ValidationError("No valid hosts")
         self.cleaned_data["hosts"] = list(hosts)
-
-
-LabelFormset = forms.inlineformset_factory(
-    models.Rule,
-    models.RuleLabel,
-    fields=("name", "value"),
-    widgets={
-        "name": forms.TextInput(attrs={"class": "form-control"}),
-        "value": forms.TextInput(attrs={"rows": 5, "class": "form-control"}),
-    },
-)
-
-
-AnnotationFormset = forms.inlineformset_factory(
-    models.Rule,
-    models.RuleAnnotation,
-    fields=("name", "value"),
-    widgets={
-        "name": forms.TextInput(attrs={"class": "form-control"}),
-        "value": forms.Textarea(attrs={"rows": 2, "class": "form-control"}),
-    },
-)

promgen/migrations/0022_rule_labels_annotations.py

@@ -0,0 +1,56 @@
+# Generated by Django 3.2.13 on 2023-07-11 03:02
+
+from django.db import migrations, models
+
+
+def forward(apps, schema_editor):
+    Rule = apps.get_model("promgen", "Rule")
+    Label = apps.get_model("promgen", "RuleLabel")
+    Annotation = apps.get_model("promgen", "RuleAnnotation")
+
+    # For our forward migration, we want to loop through all the old Label and Annotation entries
+    # and convert them to a dictionary property on our Rule model
+    for rule in Rule.objects.all():
+        rule.labels = {l.name: l.value for l in Label.objects.filter(rule_id=rule.id)}
+        rule.annotations = {l.name: l.value for l in Annotation.objects.filter(rule_id=rule.id)}
+        rule.save()
+
+
+def reverse(apps, schema_editor):
+    Rule = apps.get_model("promgen", "Rule")
+    Label = apps.get_model("promgen", "RuleLabel")
+    Annotation = apps.get_model("promgen", "RuleAnnotation")
+    for rule in Rule.objects.all():
+        Label.objects.bulk_create(
+            [Label(rule_id=rule.id, name=x, value=rule.labels[x]) for x in rule.labels]
+        )
+        Annotation.objects.bulk_create(
+            [Annotation(rule_id=rule.id, name=x, value=rule.annotations[x]) for x in rule.annotations]
+        )
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("promgen", "0021_shard_load"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="rule",
+            name="annotations",
+            field=models.JSONField(default=dict),
+        ),
+        migrations.AddField(
+            model_name="rule",
+            name="labels",
+            field=models.JSONField(default=dict),
+        ),
+        migrations.RunPython(forward, reverse),
+        migrations.DeleteModel(
+            name="RuleAnnotation",
+        ),
+        migrations.DeleteModel(
+            name="RuleLabel",
+        ),
+    ]

promgen/models.py

@@ -415,27 +415,12 @@ class Rule(models.Model):
     content_object = GenericForeignKey("content_type", "object_id", for_concrete_model=False)
     description = models.TextField(blank=True)
 
+    labels = models.JSONField(default=dict)
+    annotations = models.JSONField(default=dict)
+
     class Meta:
         ordering = ["content_type", "object_id", "name"]
 
-    @cached_property
-    def labels(self):
-        return {obj.name: obj.value for obj in self.rulelabel_set.all()}
-
-    def add_label(self, name, value):
-        return RuleLabel.objects.get_or_create(rule=self, name=name, value=value)
-
-    def add_annotation(self, name, value):
-        return RuleAnnotation.objects.get_or_create(rule=self, name=name, value=value)
-
-    @cached_property
-    def annotations(self):
-        _annotations = {obj.name: obj.value for obj in self.ruleannotation_set.all()}
-        # Skip when pk is not set, such as when test rendering a rule
-        if self.pk and "rule" not in _annotations:
-            _annotations["rule"] = resolve_domain("rule-detail", pk=self.pk)
-        return _annotations
-
     def __str__(self):
         return f"{self.name} [{self.content_object.name}]"
 
@@ -481,41 +466,14 @@ def copy_to(self, content_type, object_id):
                 macro.EXCLUSION_MACRO,
                 f'{content_type.model}="{content_object.name}",{macro.EXCLUSION_MACRO}',
             )
-            self.save()
 
             # Add a label to our new rule by default, to help ensure notifications
             # get routed to the notifier we expect
-            self.add_label(content_type.model, content_object.name)
-
-            for label in RuleLabel.objects.filter(rule_id=orig_pk):
-                # Skip service labels from our previous rule
-                if label.name in ["service", "project"]:
-                    logger.debug("Skipping %s: %s", label.name, label.value)
-                    continue
-                logger.debug("Copying %s to %s", label, self)
-                label.pk = None
-                label.rule = self
-                label.save()
-
-            for annotation in RuleAnnotation.objects.filter(rule_id=orig_pk):
-                logger.debug("Copying %s to %s", annotation, self)
-                annotation.pk = None
-                annotation.rule = self
-                annotation.save()
-
-        return self
-
-
-class RuleLabel(models.Model):
-    name = models.CharField(max_length=128)
-    value = models.CharField(max_length=128)
-    rule = models.ForeignKey("Rule", on_delete=models.CASCADE)
+            self.labels[content_type.model] = content_object.name
 
+            self.save()
 
-class RuleAnnotation(models.Model):
-    name = models.CharField(max_length=128)
-    value = models.TextField()
-    rule = models.ForeignKey("Rule", on_delete=models.CASCADE)
+        return self
 
 
 class AlertLabel(models.Model):

promgen/prometheus.py

@@ -160,33 +160,29 @@ def import_rules_v2(config, content_object=None):
     counters = collections.defaultdict(int)
     for group in config["groups"]:
         for r in group["rules"]:
-            labels = r.get("labels", {})
-            annotations = r.get("annotations", {})
 
             defaults = {
                 "clause": r["expr"],
                 "duration": r["for"],
+                "labels": r.get("labels", {}),
+                "annotations": r.get("annotations", {}),
             }
 
             # Check our labels to see if we have a project or service
             # label set and if not, default it to a global rule
             if content_object:
                 defaults["obj"] = content_object
-            elif "project" in labels:
-                defaults["obj"] = models.Project.objects.get(name=labels["project"])
-            elif "service" in labels:
-                defaults["obj"] = models.Service.objects.get(name=labels["service"])
+            elif "project" in defaults["labels"]:
+                defaults["obj"] = models.Project.objects.get(name=defaults["labels"]["project"])
+            elif "service" in defaults["labels"]:
+                defaults["obj"] = models.Service.objects.get(name=defaults["labels"]["service"])
             else:
                 defaults["obj"] = models.Site.objects.get_current()
 
-            rule, created = models.Rule.objects.get_or_create(name=r["alert"], defaults=defaults)
+            _, created = models.Rule.objects.get_or_create(name=r["alert"], defaults=defaults)
 
             if created:
                 counters["Rules"] += 1
-            for k, v in labels.items():
-                rule.add_label(k, v)
-            for k, v in annotations.items():
-                rule.add_annotation(k, v)
 
     return dict(counters)
 

promgen/serializers.py

@@ -6,6 +6,7 @@
 
 import promgen.templatetags.promgen as macro
 from promgen import models, shortcuts
+from promgen.shortcuts import resolve_domain
 
 
 class WebLinkField(serializers.Field):
@@ -88,16 +89,17 @@ def many_init(cls, queryset, *args, **kwargs):
             "content_type",
             "overrides__content_object",
             "overrides__content_type",
-            "ruleannotation_set",
-            "rulelabel_set",
         )
         return AlertRuleList(queryset, *args, **kwargs)
 
     def to_representation(self, obj):
+        annotations = obj.annotations
+        annotations["rule"] = resolve_domain("rule-detail", pk=obj.pk if obj.pk else 0)
+
         return {
             "alert": obj.name,
             "expr": macro.rulemacro(obj),
             "for": obj.duration,
             "labels": obj.labels,
-            "annotations": obj.annotations,
+            "annotations": annotations,
         }

promgen/templates/promgen/rule_formset_table.html

@@ -5,16 +5,16 @@
         <th>Value</th>
         <th>Delete?</th>
     </tr>
-    {% for row in formset %}
+    {% for form in formset %}
     <tr>
-        <td>{{ row.id }} {{ row.name }}</td>
-        <td>{{ row.value }}</td>
-        <td>{{ row.DELETE }}</td>
+        <td>{{ form.id }} {{ form.key }}</td>
+        <td>{{ form.value }}</td>
+        <td>{{ form.DELETE }}</td>
     </tr>
-    {% for k,v in row.errors.items %}
+    {% if form.errors %}
     <tr>
-        <td colspan="3">{{v}}</td>
+        <td colspan="3">{{form.errors}}</td>
     </tr>
-    {% endfor %}
+    {% endif %}
     {% endfor %}
 </table>