Custom the save method of the Rule model to handle the annotation "rule"

When creating or updating a rule on Promgen, the corresponding rule data on Prometheus will be given an annotation rule that is the link to that rule on Promgen. However, the annotation rule on Promgen always has the URI "/rule/0" because it has not been intentionally set up.

To ensure data is synced between the two applications, we customize the save method of the Rule model.

promgen/models.py

@@ -485,6 +485,22 @@ def copy_to(self, content_type, object_id):
 
         return self
 
+    # Custom logic before saving Rule to control the value of the annotation "rule":
+    # Format: annotations["rule"] = {domain}/rule/{id}
+    def save(self, *args, **kwargs):
+        with transaction.atomic():
+            if self.pk:
+                # When updating rule, we already have the primary key.
+                # Just set annotations["rule"] before saving to database.
+                self.annotations["rule"] = resolve_domain("rule-detail", pk=self.pk)
+                super().save(*args, **kwargs)
+            else:
+                # When creating a new rule, the primary key is typically not available until the instance is saved to the database.
+                # Therefore, we save it first then set annotations["rule"] and save again.
+                super().save(*args, **kwargs)
+                self.annotations["rule"] = resolve_domain("rule-detail", pk=self.pk)
+                super().save(update_fields=['annotations'])
+
 
 class AlertLabel(models.Model):
     alert = models.ForeignKey("Alert", on_delete=models.CASCADE)

promgen/tests/test_models.py

@@ -1,10 +1,12 @@
 # Copyright (c) 2020 LINE Corporation
 # These sources are released under the terms of the MIT license: see LICENSE
+from unittest import mock
 
-
+from django.contrib.contenttypes.models import ContentType
 from django.core.exceptions import ValidationError
 
 from promgen import models, validators
+from promgen.shortcuts import resolve_domain
 from promgen.tests import PromgenTest
 
 
@@ -34,3 +36,32 @@ def test_validators(self):
             validators.metricname(
                 "[[this.$el.ownerDocument.defaultView.alert(1337)]]",
             )
+
+    @mock.patch("django.dispatch.dispatcher.Signal.send")
+    def test_custom_save_rule(self):
+        # Check if annotation["rule"] is automatically set to be {domain}/rule/{id} when creating a new rule
+        rule = models.Rule(
+            name="example-rule",
+            content_type=ContentType.objects.get_for_model(models.Site),
+            object_id=1,
+            clause="up==1",
+            duration="1s"
+        )
+        rule.save()
+        self.assertEqual(resolve_domain("rule-detail", rule.pk), rule.annotations["rule"])
+
+        # Check if annotation["rule"] is automatically set to be {domain}/rule/{id} when updating an existed rule
+        rule.name = "another-example-rule"
+        rule.annotations["rule"] = "another-annotation-value"
+        rule.save()
+        self.assertEqual(resolve_domain("rule-detail", rule.pk), rule.annotations["rule"])
+
+        # Check if annotation["rule"] is still set to be {domain}/rule/{id} when trying to remove annotation["rule"]
+        rule.annotations["rule"] = None
+        rule.save()
+        self.assertEqual(resolve_domain("rule-detail", rule.pk), rule.annotations["rule"])
+
+        # Check if annotation["rule"] of new rule is automatically set to be {domain}/rule/{id}  when cloning an existed rule
+        new_rule = rule.copy_to(content_type="service", object_id=2)
+        self.assertEqual(resolve_domain("rule-detail", rule.pk), rule.annotations["rule"])
+        self.assertEqual(resolve_domain("rule-detail", new_rule.pk), new_rule.annotations["rule"])