Turn prometheus access into an org feature and grant to all orgs with existing token

temba/orgs/migrations/0165_grant_prometheus_feature.py

@@ -0,0 +1,25 @@
+# Generated by Django 5.1.4 on 2025-01-16 18:54
+
+from django.db import migrations
+
+
+def grant_prometheus_feature(apps, schema_editor):
+    Org = apps.get_model("orgs", "Org")
+
+    num_updated = 0
+    for org in Org.objects.exclude(prometheus_token=None):
+        org.features += ["prometheus"]
+        org.save(update_fields=("features",))
+        num_updated += 1
+
+    if num_updated:
+        print(f"Granted prometheus feature to {num_updated} orgs with existing tokens")
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("orgs", "0164_remove_viewers"),
+    ]
+
+    operations = [migrations.RunPython(grant_prometheus_feature, migrations.RunPython.noop)]

temba/orgs/models.py

@@ -456,11 +456,13 @@ class Org(SmartModel):
     FEATURE_NEW_ORGS = "new_orgs"  # can create new workspace with same login
     FEATURE_CHILD_ORGS = "child_orgs"  # can create child workspaces of this org
     FEATURE_TEAMS = "teams"  # can create teams to organize agent users
+    FEATURE_PROMETHEUS = "prometheus"  # can create a prometheus token to access metrics
     FEATURES_CHOICES = (
         (FEATURE_USERS, _("Users")),
         (FEATURE_NEW_ORGS, _("New Orgs")),
         (FEATURE_CHILD_ORGS, _("Child Orgs")),
         (FEATURE_TEAMS, _("Teams")),
+        (FEATURE_PROMETHEUS, _("Prometheus")),
     )
 
     LIMIT_CHANNELS = "channels"

temba/orgs/tests/test_org.py

@@ -299,46 +299,6 @@ def start_flow_via_api():
 
         mock_async_start.assert_called_once()
 
-    def test_prometheus(self):
-        # visit as editor, no prometheus section
-        self.login(self.editor)
-        settings_url = reverse("orgs.org_workspace")
-        response = self.client.get(settings_url)
-
-        self.assertNotContains(response, "Prometheus")
-
-        # admin can see it though
-        self.login(self.admin)
-
-        response = self.client.get(settings_url)
-        self.assertContains(response, "Prometheus")
-        self.assertContains(response, "Enable")
-
-        # enable it
-        prometheus_url = reverse("orgs.org_prometheus")
-        response = self.client.post(prometheus_url, {}, follow=True)
-        self.assertContains(response, "Disable")
-
-        # make sure our token is set
-        self.org.refresh_from_db()
-        self.assertIsNotNone(self.org.prometheus_token)
-
-        # other admin sees it enabled too
-        self.other_admin = self.create_user("[email protected]")
-        self.org.add_user(self.other_admin, OrgRole.ADMINISTRATOR)
-        self.login(self.other_admin)
-
-        response = self.client.get(settings_url)
-        self.assertContains(response, "Prometheus")
-        self.assertContains(response, "Disable")
-
-        # now disable it
-        response = self.client.post(prometheus_url, {}, follow=True)
-        self.assertContains(response, "Enable")
-
-        self.org.refresh_from_db()
-        self.assertIsNone(self.org.prometheus_token)
-
     def test_resthooks(self):
         resthook_url = reverse("orgs.org_resthooks")
 

temba/orgs/views/tests/test_migrations.py

@@ -0,0 +1,18 @@
+from temba.tests import MigrationTest
+
+
+class GrantPrometheusFeatureTest(MigrationTest):
+    app = "orgs"
+    migrate_from = "0164_remove_viewers"
+    migrate_to = "0165_grant_prometheus_feature"
+
+    def setUpBeforeMigration(self, apps):
+        self.org.prometheus_token = "sesame"
+        self.org.save(update_fields=("prometheus_token",))
+
+    def test_migration(self):
+        self.org.refresh_from_db()
+        self.org2.refresh_from_db()
+
+        self.assertIn("prometheus", self.org.features)
+        self.assertNotIn("prometheus", self.org2.features)

temba/orgs/views/tests/test_orgcrudl.py

@@ -119,7 +119,7 @@ def test_workspace(self):
         response = self.assertListFetch(workspace_url, [self.editor, self.admin])
 
         # make sure we have the appropriate number of sections
-        self.assertEqual(6, len(response.context["formax"].sections))
+        self.assertEqual(5, len(response.context["formax"].sections))
 
         self.assertPageMenu(
             f"{reverse('orgs.org_menu')}settings/",
@@ -137,8 +137,8 @@ def test_workspace(self):
             ],
         )
 
-        # enable child workspaces, users and teams
-        self.org.features = [Org.FEATURE_USERS, Org.FEATURE_CHILD_ORGS, Org.FEATURE_TEAMS]
+        # enable more features..
+        self.org.features = [Org.FEATURE_USERS, Org.FEATURE_CHILD_ORGS, Org.FEATURE_TEAMS, Org.FEATURE_PROMETHEUS]
         self.org.save(update_fields=("features",))
 
         self.child_org = Org.objects.create(
@@ -153,6 +153,8 @@ def test_workspace(self):
         with self.assertNumQueries(9):
             response = self.client.get(workspace_url)
 
+        self.assertEqual(6, len(response.context["formax"].sections))
+
         # should have an extra menu options for workspaces and users
         self.assertPageMenu(
             f"{reverse('orgs.org_menu')}settings/",
@@ -1256,3 +1258,36 @@ def test_languages(self):
             )
 
         languages.reload()
+
+    def test_prometheus(self):
+        prometheus_url = reverse("orgs.org_prometheus")
+        workspace_url = reverse("orgs.org_workspace")
+
+        # by default, prometheus feature not enabled
+        response = self.requestView(workspace_url, self.admin)
+        self.assertNotContains(response, "Prometheus")
+
+        self.org.features = [Org.FEATURE_PROMETHEUS]
+        self.org.save(update_fields=("features",))
+
+        response = self.requestView(workspace_url, self.admin)
+        self.assertContains(response, "Prometheus")
+        self.assertContains(response, "Enable")
+
+        # not visible to editor users tho
+        response = self.requestView(workspace_url, self.editor)
+        self.assertNotContains(response, "Prometheus")
+
+        self.assertRequestDisallowed(prometheus_url, [self.editor, self.agent])
+
+        # enable it...
+        self.requestView(prometheus_url, self.admin, post_data={})
+
+        self.org.refresh_from_db()
+        self.assertIsNotNone(self.org.prometheus_token)
+
+        # and then disable it...
+        self.requestView(prometheus_url, self.admin, post_data={})
+
+        self.org.refresh_from_db()
+        self.assertIsNone(self.org.prometheus_token)

temba/orgs/views/views.py

@@ -2019,14 +2019,15 @@ def pre_save(self, obj):
 
             return super().pre_save(obj)
 
-    class Prometheus(FormaxSectionMixin, InferOrgMixin, OrgPermsMixin, SmartUpdateView):
+    class Prometheus(RequireFeatureMixin, FormaxSectionMixin, InferOrgMixin, OrgPermsMixin, SmartUpdateView):
         class Form(forms.ModelForm):
             class Meta:
                 model = Org
                 fields = ("id",)
 
         form_class = Form
         success_url = "@orgs.org_workspace"
+        require_feature = Org.FEATURE_PROMETHEUS
 
         def save(self, obj):
             org = self.request.org