Repository: pypi/warehouse

Issue:
Re-enable 2FA metrics, speed up queries by denormalizing counts
In https://github.com/pypa/warehouse/pull/11699 we disabled the 2FA metrics temporarily due to DB performance issues.

We should improve the count performance with something similar to https://github.com/pypa/warehouse/pull/745/ and re-enable the metrics.


Discussion:

Resolving patch:
diff --git a/warehouse/accounts/models.py b/warehouse/accounts/models.py
--- a/warehouse/accounts/models.py
+++ b/warehouse/accounts/models.py
@@ -27,8 +27,6 @@
     String,
     Text,
     UniqueConstraint,
-    and_,
-    or_,
     orm,
     select,
     sql,
@@ -136,23 +134,10 @@ def email(self):
             .scalar_subquery()
         )
 
-    @hybrid_property
+    @property
     def has_two_factor(self):
         return self.totp_secret is not None or len(self.webauthn) > 0
 
-    @has_two_factor.expression  # type: ignore
-    def has_two_factor(self):
-        return (
-            select([True])
-            .where(
-                or_(
-                    WebAuthn.user_id == self.id,
-                    and_(User.id == self.id, User.totp_secret.is_not(None)),
-                )
-            )
-            .limit(1)
-        )
-
     @property
     def has_recovery_codes(self):
         return any(not code.burned for code in self.recovery_codes)
diff --git a/warehouse/packaging/__init__.py b/warehouse/packaging/__init__.py
--- a/warehouse/packaging/__init__.py
+++ b/warehouse/packaging/__init__.py
@@ -21,6 +21,7 @@
 from warehouse.packaging.models import File, Project, Release, Role
 from warehouse.packaging.tasks import (
     compute_2fa_mandate,
+    compute_2fa_metrics,
     compute_trending,
     update_description_html,
 )
@@ -107,7 +108,7 @@ def includeme(config):
         config.add_periodic_task(crontab(minute=0, hour=3), compute_2fa_mandate)
 
     # Add a periodic task to generate 2FA metrics
-    # config.add_periodic_task(crontab(minute="*/5"), compute_2fa_metrics)
+    config.add_periodic_task(crontab(minute="*/5"), compute_2fa_metrics)
 
     # Add a periodic task to compute trending once a day, assuming we have
     # been configured to be able to access BigQuery.
diff --git a/warehouse/packaging/tasks.py b/warehouse/packaging/tasks.py
--- a/warehouse/packaging/tasks.py
+++ b/warehouse/packaging/tasks.py
@@ -23,7 +23,7 @@
 from warehouse.cache.origin import IOriginCache
 from warehouse.email import send_two_factor_mandate_email
 from warehouse.metrics import IMetricsService
-from warehouse.packaging.models import Description, File, Project, Release
+from warehouse.packaging.models import Description, File, Project, Release, Role
 from warehouse.utils import readme
 
 
@@ -105,10 +105,45 @@ def compute_2fa_metrics(request):
         critical_maintainers.count(),
     )
 
+    # Number of critical project maintainers with TOTP enabled
+    total_critical_project_maintainers_with_totp_enabled = (
+        request.db.query(User.id)
+        .distinct()
+        .join(Role, Role.user_id == User.id)
+        .join(Project, Project.id == Role.project_id)
+        .where(Project.pypi_mandates_2fa)
+        .where(User.totp_secret.is_not(None))
+        .count()
+    )
+    metrics.gauge(
+        "warehouse.2fa.total_critical_maintainers_with_totp_enabled",
+        total_critical_project_maintainers_with_totp_enabled,
+    )
+
+    # Number of critical project maintainers with WebAuthn enabled
+    metrics.gauge(
+        "warehouse.2fa.total_critical_maintainers_with_webauthn_enabled",
+        request.db.query(User.id)
+        .distinct()
+        .join(Role.user)
+        .join(Role.project)
+        .join(WebAuthn, WebAuthn.user_id == User.id)
+        .where(Project.pypi_mandates_2fa)
+        .count(),
+    )
+
     # Number of critical project maintainers with 2FA enabled
     metrics.gauge(
         "warehouse.2fa.total_critical_maintainers_with_2fa_enabled",
-        critical_maintainers.where(User.has_two_factor).count(),
+        total_critical_project_maintainers_with_totp_enabled
+        + request.db.query(User.id)
+        .distinct()
+        .join(Role.user)
+        .join(Role.project)
+        .join(WebAuthn, WebAuthn.user_id == User.id)
+        .where(Project.pypi_mandates_2fa)
+        .where(User.totp_secret.is_(None))
+        .count(),
     )
 
     # Number of projects manually requiring 2FA
@@ -124,21 +159,32 @@ def compute_2fa_metrics(request):
     )
 
     # Total number of users with TOTP enabled
+    total_users_with_totp_enabled = (
+        request.db.query(User).where(User.totp_secret.is_not(None)).count()
+    )
     metrics.gauge(
         "warehouse.2fa.total_users_with_totp_enabled",
-        request.db.query(User).where(User.totp_secret.is_not(None)).count(),
+        total_users_with_totp_enabled,
     )
 
     # Total number of users with WebAuthn enabled
     metrics.gauge(
         "warehouse.2fa.total_users_with_webauthn_enabled",
-        request.db.query(WebAuthn.user_id).group_by(WebAuthn.user_id).count(),
+        request.db.query(User.id)
+        .distinct()
+        .join(WebAuthn, WebAuthn.user_id == User.id)
+        .count(),
     )
 
     # Total number of users with 2FA enabled
     metrics.gauge(
         "warehouse.2fa.total_users_with_two_factor_enabled",
-        request.db.query(User).where(User.has_two_factor).count(),
+        total_users_with_totp_enabled
+        + request.db.query(User.id)
+        .distinct()
+        .join(WebAuthn, WebAuthn.user_id == User.id)
+        .where(User.totp_secret.is_(None))
+        .count(),
     )
 
 


Test patch:
diff --git a/tests/unit/packaging/test_tasks.py b/tests/unit/packaging/test_tasks.py
--- a/tests/unit/packaging/test_tasks.py
+++ b/tests/unit/packaging/test_tasks.py
@@ -727,6 +727,10 @@ def test_compute_2fa_metrics(db_request, monkeypatch):
     assert gauge.calls == [
         pretend.call("warehouse.2fa.total_critical_projects", 1),
         pretend.call("warehouse.2fa.total_critical_maintainers", 3),
+        pretend.call("warehouse.2fa.total_critical_maintainers_with_totp_enabled", 1),
+        pretend.call(
+            "warehouse.2fa.total_critical_maintainers_with_webauthn_enabled", 1
+        ),
         pretend.call("warehouse.2fa.total_critical_maintainers_with_2fa_enabled", 2),
         pretend.call("warehouse.2fa.total_projects_with_2fa_opt_in", 1),
         pretend.call("warehouse.2fa.total_projects_with_two_factor_required", 2),