feat: add metric button to enable disable metrics

Author: Kl0ven

controller/sentry/admin.py

@@ -14,7 +14,8 @@
 from django_json_widget.widgets import JSONEditorWidget
 from django_object_actions import DjangoObjectActions, takes_instance_or_queryset
 
-from controller.sentry.forms import BumpForm
+from controller.sentry.choices import MetricType
+from controller.sentry.forms import BumpForm, MetricForm
 from controller.sentry.models import App
 from controller.sentry.utils import invalidate_cache
 
@@ -63,32 +64,36 @@ class AppAdmin(
         [
             "Sentry",
             {
-                "classes": ("collapse", "open"),
+                "classes": ("collapse",),
                 "fields": ("sentry_project_slug",),
             },
         ],
         [
-            "WSGI",
+            "Metric - WSGI",
             {
-                "classes": ("collapse", "open"),
-                "fields": ("wsgi_ignore_path", "wsgi_collect_metrics", "wsgi_metrics"),
+                "classes": ("collapse",),
+                "fields": (
+                    "wsgi_collect_metrics",
+                    "wsgi_ignore_path",
+                    "wsgi_metrics",
+                ),
             },
         ],
         [
-            "Celery",
+            "Metric - Celery",
             {
-                "classes": ("collapse", "open"),
+                "classes": ("collapse",),
                 "fields": (
-                    "celery_ignore_task",
                     "celery_collect_metrics",
+                    "celery_ignore_task",
                     "celery_metrics",
                 ),
             },
         ],
     ]
     actions = ["bump_sample_rate"]
     changelist_actions = ["panic", "unpanic"]
-    change_actions = ["bump_sample_rate"]
+    change_actions = ["bump_sample_rate", "enable_disable_metrics"]
 
     def get_changelist_actions(self, request):
         allowed_actions = []
@@ -104,6 +109,7 @@ def get_change_actions(self, request, object_id, form_url):
                 allowed_actions.append(action)
         return allowed_actions
 
+    # ----- Bump Sample Rate
     @takes_instance_or_queryset
     @add_form_to_action(BumpForm)
     @confirm_action()
@@ -114,6 +120,8 @@ def bump_sample_rate(self, request, queryset, form: BumpForm = None):  # pylint:
             active_sample_rate=form.cleaned_data["new_sample_rate"],
             active_window_end=new_date,
         )
+        for app in queryset:
+            invalidate_cache(f"/sentry/apps/{app.reference}/")
 
     bump_sample_rate.allowed_permissions = ("bump_sample_rate",)
 
@@ -125,6 +133,29 @@ def has_bump_sample_rate_permission(self, request):
         panic = cache.get(settings.PANIC_KEY)
         return not panic and request.user.has_perm("%s.%s" % (opts.app_label, codename))
 
+    # ----- Update Metrics
+    @takes_instance_or_queryset
+    @add_form_to_action(MetricForm)
+    @admin.action(description="Enable/Disable Metrics Collection")
+    def enable_disable_metrics(self, request, queryset, form: MetricForm = None):  # pylint: disable=unused-argument
+        metrics = form.cleaned_data["metrics"]
+        app: App
+        for app in queryset:
+            for metric in MetricType:
+                app.set_metric(metric, metric.value in metrics)
+            invalidate_cache(f"/sentry/apps/{app.reference}/")
+            app.save()
+
+    enable_disable_metrics.allowed_permissions = ("enable_disable_metrics",)
+
+    def has_enable_disable_metrics_permission(self, request):
+        """Does the user have the enable_disable_metrics permission?"""
+        opts = self.opts
+        codename = get_permission_codename("enable_disable_metrics", opts)
+
+        return request.user.has_perm("%s.%s" % (opts.app_label, codename))
+
+    # ----- Panic / Unpanic
     @takes_instance_or_queryset
     @confirm_action(display_queryset=False)
     @admin.action(description="Panic")
@@ -157,6 +188,7 @@ def has_unpanic_permission(self, request):
         codename = get_permission_codename("panic", opts)
         return panic and request.user.has_perm("%s.%s" % (opts.app_label, codename))
 
+    # Save model
     def save_model(self, request, obj, form, change) -> None:
         invalidate_cache(f"/sentry/apps/{obj.reference}/")
         return super().save_model(request, obj, form, change)

controller/sentry/forms.py

@@ -1,14 +1,16 @@
 from django.conf import settings
 from django.core.exceptions import ValidationError
-from django.forms import DurationField, FloatField, Form
+from django.forms import DurationField, FloatField, Form, MultipleChoiceField
+from django.forms.widgets import CheckboxSelectMultiple
 from durationwidget.widgets import TimeDurationWidget
 
+from controller.sentry.choices import MetricType
+
 
 class BumpForm(Form):
-    new_sample_rate = FloatField(help_text="Sample rate between 0 and 1")
+    new_sample_rate = FloatField(help_text="Sample rate between 0 and 1", initial=0.5)
     duration = DurationField(
-        widget=TimeDurationWidget(show_days=False, show_seconds=False),
-        help_text="Duration",
+        widget=TimeDurationWidget(show_days=False, show_seconds=False), help_text="Duration", initial="0:30:00"
     )
 
     def clean_new_sample_rate(self):
@@ -22,3 +24,12 @@ def clean_duration(self):
         if data.total_seconds() < 0 or data.total_seconds() > settings.MAX_BUMP_TIME_SEC:
             raise ValidationError(f"duration must be between 0 and {settings.MAX_BUMP_TIME_SEC}")
         return data
+
+
+class MetricForm(Form):
+    metrics = MultipleChoiceField(
+        choices=MetricType.choices,
+        widget=CheckboxSelectMultiple,
+        required=False,
+        help_text="Disable or Enable metric gathering",
+    )

controller/sentry/migrations/0008_alter_app_options.py

@@ -0,0 +1,40 @@
+# Generated by Django 4.1.6 on 2023-02-03 14:36
+
+from django.core.management.sql import emit_post_migrate_signal
+from django.db import migrations
+
+
+def apply_migration(apps, schema_editor):
+    emit_post_migrate_signal(2, False, "default")
+    Group = apps.get_model("auth", "Group")
+    Permission = apps.get_model("auth", "Permission")
+    owner = Group.objects.get(name="Owner")
+    admin = Group.objects.get(name="Admin")
+    developer = Group.objects.get(name="Developer")
+
+    metric_permission = Permission.objects.get(codename="enable_disable_metrics_app")
+
+    owner.permissions.add(metric_permission)
+    admin.permissions.add(metric_permission)
+    developer.permissions.add(metric_permission)
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("sentry", "0007_app_sentry_project_slug"),
+    ]
+
+    operations = [
+        migrations.AlterModelOptions(
+            name="app",
+            options={
+                "permissions": [
+                    ("bump_sample_rate_app", "Can bump sample rate"),
+                    ("panic_app", "Panic! Set all sample rate to 0"),
+                    ("enable_disable_metrics_app", "Can Enable/Disable metrics"),
+                ]
+            },
+        ),
+        migrations.RunPython(apply_migration, migrations.RunPython.noop),
+    ]

controller/sentry/models.py

@@ -61,6 +61,10 @@ def get_metric(self, metric_type: MetricType):
         prefix = metric_type.value.lower()
         return getattr(self, f"{prefix}_collect_metrics"), getattr(self, f"{prefix}_metrics")
 
+    def set_metric(self, metric_type: MetricType, metric_state: bool):
+        prefix = metric_type.value.lower()
+        setattr(self, f"{prefix}_collect_metrics", metric_state)
+
     def get_sentry_id(self):
         res = self.reference.split("_")
         if len(res) != 3:
@@ -71,4 +75,5 @@ class Meta:
         permissions = [
             ("bump_sample_rate_app", "Can bump sample rate"),
             ("panic_app", "Panic! Set all sample rate to 0"),
+            ("enable_disable_metrics_app", "Can Enable/Disable metrics"),
         ]

controller/sentry/tests/test_admin.py

@@ -7,7 +7,7 @@
 from undecorated import undecorated
 
 from controller.sentry.admin import AppAdmin
-from controller.sentry.forms import BumpForm
+from controller.sentry.forms import BumpForm, MetricForm
 from controller.sentry.models import App
 
 
@@ -146,6 +146,41 @@ def test_app_admin_bump(request, admin_with_user):
     assert app.active_window_end is not None
 
 
+@pytest.mark.django_db
+@pytest.mark.parametrize("user_group", ["Developer"])
+@pytest.mark.admin_site(model_class=App)
+def test_app_admin_metrics(request, admin_with_user):
+    app = App(wsgi_collect_metrics=False, celery_collect_metrics=False)
+    app.save()
+    form = MetricForm({"metrics": ["WSGI", "CELERY"]})
+
+    assert form.is_valid()
+    site, request = admin_with_user
+    enable_disable_metrics = undecorated(site.enable_disable_metrics)
+    enable_disable_metrics(site, request, App.objects.filter(reference=app.reference), form=form)
+    app.refresh_from_db()
+    assert app.wsgi_collect_metrics
+    assert app.celery_collect_metrics
+
+    form = MetricForm({"metrics": []})
+    assert form.is_valid()
+    site, request = admin_with_user
+    enable_disable_metrics = undecorated(site.enable_disable_metrics)
+    enable_disable_metrics(site, request, App.objects.filter(reference=app.reference), form=form)
+    app.refresh_from_db()
+    assert not app.wsgi_collect_metrics
+    assert not app.celery_collect_metrics
+
+    form = MetricForm({"metrics": ["WSGI"]})
+    assert form.is_valid()
+    site, request = admin_with_user
+    enable_disable_metrics = undecorated(site.enable_disable_metrics)
+    enable_disable_metrics(site, request, App.objects.filter(reference=app.reference), form=form)
+    app.refresh_from_db()
+    assert app.wsgi_collect_metrics
+    assert not app.celery_collect_metrics
+
+
 @patch("controller.sentry.admin.cache")
 @pytest.mark.django_db
 @pytest.mark.parametrize("user_group", ["Developer"])
@@ -197,13 +232,13 @@ def test_app_admin_get_changelist_actions(cache: Mock, admin_with_user, result,
 @pytest.mark.parametrize(
     "user_group,panic,result",
     [
-        ("Owner", False, ["bump_sample_rate"]),
-        ("Admin", False, ["bump_sample_rate"]),
-        ("Developer", False, ["bump_sample_rate"]),
+        ("Owner", False, ["bump_sample_rate", "enable_disable_metrics"]),
+        ("Admin", False, ["bump_sample_rate", "enable_disable_metrics"]),
+        ("Developer", False, ["bump_sample_rate", "enable_disable_metrics"]),
         ("Viewer", False, []),
-        ("Owner", True, []),
-        ("Admin", True, []),
-        ("Developer", True, []),
+        ("Owner", True, ["enable_disable_metrics"]),
+        ("Admin", True, ["enable_disable_metrics"]),
+        ("Developer", True, ["enable_disable_metrics"]),
         ("Viewer", True, []),
     ],
 )

poetry.lock

@@ -9,7 +9,7 @@ packages = [{ include = "controller" }]
 [tool.poetry.dependencies]
 python = "^3.9"
 Django = "^4.1.3"
-django-admin-action-tools = "^1.1"
+django-admin-action-tools = "1.1.5"
 djangorestframework = "^3.14.0"
 gunicorn = "^20.1.0"
 psycopg2-binary = "^2.9.5"