From 20cb1f1f171916ff97ab1bb470af64ad5e7cfb94 Mon Sep 17 00:00:00 2001 From: hoang Date: Mon, 2 Dec 2024 16:23:32 +0700 Subject: [PATCH] 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 | 17 +++++++++++++++++ promgen/tests/test_models.py | 37 +++++++++++++++++++++++++++++++++++- 2 files changed, 53 insertions(+), 1 deletion(-) diff --git a/promgen/models.py b/promgen/models.py index 27330f296..d99467a82 100644 --- a/promgen/models.py +++ b/promgen/models.py @@ -485,6 +485,23 @@ 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) diff --git a/promgen/tests/test_models.py b/promgen/tests/test_models.py index 373146d68..33d423620 100644 --- a/promgen/tests/test_models.py +++ b/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,36 @@ def test_validators(self): validators.metricname( "[[this.$el.ownerDocument.defaultView.alert(1337)]]", ) + + @mock.patch("django.dispatch.dispatcher.Signal.send") + def test_rule_annotation(self, mock_post): + # 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"])