Add an admin action to remove stale plan_activated_users (#778)

codecov_auth/admin.py

@@ -1,14 +1,17 @@
 import logging
-from typing import Optional
+from datetime import timedelta
+from typing import Optional, Sequence
 
 import django.forms as forms
 from django.conf import settings
 from django.contrib import admin, messages
 from django.contrib.admin.models import LogEntry
+from django.db.models import OuterRef, Subquery
 from django.db.models.fields import BLANK_CHOICE_DASH
 from django.forms import CheckboxInput, Select
 from django.http import HttpRequest
 from django.shortcuts import redirect, render
+from django.utils import timezone
 from django.utils.html import format_html
 from shared.django_apps.codecov_auth.models import (
     Account,
@@ -20,7 +23,7 @@
 from codecov.admin import AdminMixin
 from codecov.commands.exceptions import ValidationError
 from codecov_auth.helpers import History
-from codecov_auth.models import OrganizationLevelToken, Owner, SentryUser, User
+from codecov_auth.models import OrganizationLevelToken, Owner, SentryUser, Session, User
 from codecov_auth.services.org_level_token_service import OrgLevelTokenService
 from plan.constants import USER_PLAN_REPRESENTATIONS
 from plan.service import PlanService
@@ -336,13 +339,75 @@ class OwnerOrgInline(admin.TabularInline):
     fields = [] + readonly_fields
 
 
+def find_and_remove_stale_users(
+    orgs: Sequence[Owner], date_threshold: timedelta | None = None
+) -> tuple[set[int], set[int]]:
+    """
+    This functions finds all the stale `plan_activated_users` in any of the given `orgs`.
+
+    It then removes all those stale users from the given `orgs`,
+    returning the set of stale users (`ownerid`), and the set of `orgs` that were updated (`ownerid`).
+
+    A user is considered stale if it had no API or login `Session` or any opened PR within `date_threshold`.
+    If no `date_threshold` is given, it defaults to *90 days*.
+    """
+
+    active_users = set()
+    for org in orgs:
+        active_users.update(set(org.plan_activated_users))
+
+    if not active_users:
+        return (set(), set())
+
+    # NOTE: the `annotate_last_pull_timestamp` manager/queryset method does the same `annotate` with `Subquery`.
+    sessions = Session.objects.filter(owner=OuterRef("pk")).order_by("-lastseen")
+    resolved_users = list(
+        Owner.objects.filter(ownerid__in=active_users)
+        .annotate(latest_session=Subquery(sessions.values("lastseen")[:1]))
+        .annotate_last_pull_timestamp()
+        .values_list("ownerid", "latest_session", "last_pull_timestamp", named=True)
+    )
+
+    threshold = timezone.now() - (date_threshold or timedelta(days=90))
+
+    def is_stale(user: dict) -> bool:
+        return (user.latest_session is None or user.latest_session < threshold) and (
+            # NOTE: `last_pull_timestamp` is not timezone-aware, so we explicitly compare without timezones here
+            user.last_pull_timestamp is None
+            or user.last_pull_timestamp.replace(tzinfo=None)
+            < threshold.replace(tzinfo=None)
+        )
+
+    stale_users = {user.ownerid for user in resolved_users if is_stale(user)}
+
+    # TODO: the existing stale user cleanup script clears the `oauth_token`, though the reason for that is not clear?
+    # Owner.objects.filter(ownerid__in=stale_users).update(oauth_token=None)
+
+    affected_orgs = {
+        org for org in orgs if stale_users.intersection(set(org.plan_activated_users))
+    }
+
+    if not affected_orgs:
+        return (set(), set())
+
+    # TODO: it might make sense to run all this within a transaction and locking the `affected_orgs` for update,
+    # as we have a slight chance of races between querying the `orgs` at the very beginning and updating them here:
+    for org in affected_orgs:
+        org.plan_activated_users = list(
+            set(org.plan_activated_users).difference(stale_users)
+        )
+    Owner.objects.bulk_update(affected_orgs, ["plan_activated_users"])
+
+    return (stale_users, {org.ownerid for org in affected_orgs})
+
+
 @admin.register(Account)
 class AccountAdmin(AdminMixin, admin.ModelAdmin):
     list_display = ("name", "is_active", "organizations_count", "all_user_count")
     search_fields = ("name__iregex", "id")
     search_help_text = "Search by name (can use regex), or id (exact)"
     inlines = [OwnerOrgInline, StripeBillingInline, InvoiceBillingInline]
-    actions = ["seat_check", "link_users_to_account"]
+    actions = ["seat_check", "link_users_to_account", "deactivate_stale_users"]
 
     readonly_fields = ["id", "created_at", "updated_at", "users"]
 
@@ -356,6 +421,24 @@ class AccountAdmin(AdminMixin, admin.ModelAdmin):
         "is_delinquent",
     ]
 
+    @admin.action(description="Deactivate all stale `plan_activated_users`")
+    def deactivate_stale_users(self, request, queryset):
+        orgs = [org for account in queryset for org in account.organizations.all()]
+        stale_users, updated_orgs = find_and_remove_stale_users(orgs)
+
+        if not stale_users or not updated_orgs:
+            self.message_user(
+                request,
+                "No stale users found in selected accounts / organizations.",
+                messages.INFO,
+            )
+        else:
+            self.message_user(
+                request,
+                f"Removed {len(stale_users)} stale users from {len(updated_orgs)} affected organizations.",
+                messages.SUCCESS,
+            )
+
     @admin.action(
         description="Count current plan_activated_users across all Organizations"
     )

codecov_auth/tests/test_admin.py

@@ -1,9 +1,12 @@
+from datetime import timedelta
 from unittest.mock import MagicMock, patch
 
+import pytest
 from django.contrib.admin.helpers import ACTION_CHECKBOX_NAME
 from django.contrib.admin.sites import AdminSite
 from django.test import RequestFactory, TestCase
 from django.urls import reverse
+from django.utils import timezone
 from shared.django_apps.codecov_auth.models import (
     Account,
     AccountsUsers,
@@ -24,14 +27,18 @@
     OwnerAdmin,
     StripeBillingAdmin,
     UserAdmin,
+    find_and_remove_stale_users,
 )
 from codecov_auth.models import OrganizationLevelToken, Owner, SentryUser, User
 from codecov_auth.tests.factories import (
     OrganizationLevelTokenFactory,
     OwnerFactory,
     SentryUserFactory,
+    SessionFactory,
     UserFactory,
 )
+from core.models import Pull
+from core.tests.factories import PullFactory, RepositoryFactory
 from plan.constants import (
     ENTERPRISE_CLOUD_USER_PLAN_REPRESENTATIONS,
     PlanName,
@@ -421,6 +428,84 @@ def test_user_admin_detail_page(self):
         assert sentry_user.refresh_token not in res.content.decode("utf-8")
 
 
+def create_stale_users(
+    account: Account | None = None,
+) -> tuple[list[Owner], list[Owner]]:
+    org_1 = OwnerFactory(account=account)
+    org_2 = OwnerFactory(account=account)
+
+    now = timezone.now()
+
+    user_1 = OwnerFactory(user=UserFactory())  # stale, neither session nor PR
+
+    user_2 = OwnerFactory(user=UserFactory())  # semi-stale, semi-old session
+    SessionFactory(owner=user_2, lastseen=now - timedelta(days=45))
+
+    user_3 = OwnerFactory(user=UserFactory())  # stale, old session
+    SessionFactory(owner=user_3, lastseen=now - timedelta(days=120))
+
+    user_4 = OwnerFactory(user=UserFactory())  # semi-stale, semi-old PR
+    pull = PullFactory(
+        repository=RepositoryFactory(),
+        author=user_4,
+    )
+    pull.updatestamp = now - timedelta(days=45)
+    super(Pull, pull).save()  # `Pull` overrides the `updatestamp` on each save
+
+    user_5 = OwnerFactory(user=UserFactory())  # stale, old PR
+    pull = PullFactory(
+        repository=RepositoryFactory(),
+        author=user_5,
+    )
+    pull.updatestamp = now - timedelta(days=120)
+    super(Pull, pull).save()  # `Pull` overrides the `updatestamp` on each save
+
+    org_1.plan_activated_users = [
+        user_1.ownerid,
+        user_2.ownerid,
+    ]
+    org_1.save()
+
+    org_2.plan_activated_users = [
+        user_3.ownerid,
+        user_4.ownerid,
+        user_5.ownerid,
+    ]
+    org_2.save()
+
+    return ([org_1, org_2], [user_1, user_2, user_3, user_4, user_5])
+
+
+@pytest.mark.django_db()
+def test_stale_user_cleanup():
+    orgs, users = create_stale_users()
+
+    # remove stale users with default > 90 days
+    removed_users, affected_orgs = find_and_remove_stale_users(orgs)
+    assert removed_users == set([users[0].ownerid, users[2].ownerid, users[4].ownerid])
+    assert affected_orgs == set([orgs[0].ownerid, orgs[1].ownerid])
+
+    orgs = list(
+        Owner.objects.filter(ownerid__in=[org.ownerid for org in orgs])
+    )  # re-fetch orgs
+    # all good, nothing to do
+    removed_users, affected_orgs = find_and_remove_stale_users(orgs)
+    assert removed_users == set()
+    assert affected_orgs == set()
+
+    # remove even more stale users
+    removed_users, affected_orgs = find_and_remove_stale_users(orgs, timedelta(days=30))
+    assert removed_users == set([users[1].ownerid, users[3].ownerid])
+    assert affected_orgs == set([orgs[0].ownerid, orgs[1].ownerid])
+
+    orgs = list(
+        Owner.objects.filter(ownerid__in=[org.ownerid for org in orgs])
+    )  # re-fetch orgs
+    # all the users have been deactivated by now
+    for org in orgs:
+        assert len(org.plan_activated_users) == 0
+
+
 class AccountAdminTest(TestCase):
     def setUp(self):
         staff_user = UserFactory(is_staff=True)
@@ -700,6 +785,35 @@ def test_link_users_to_account_remove_unneeded_account_users(self):
             AccountsUsers.objects.filter(user=self.owner_with_user_1.user).exists()
         )
 
+    def test_deactivate_stale_users(self):
+        account = AccountFactory()
+        orgs, users = create_stale_users(account)
+
+        res = self.client.post(
+            reverse("admin:codecov_auth_account_changelist"),
+            {
+                "action": "deactivate_stale_users",
+                ACTION_CHECKBOX_NAME: [account.pk],
+            },
+        )
+        messages = list(res.wsgi_request._messages)
+        self.assertEqual(
+            messages[-1].message, "Removed 3 stale users from 2 affected organizations."
+        )
+
+        res = self.client.post(
+            reverse("admin:codecov_auth_account_changelist"),
+            {
+                "action": "deactivate_stale_users",
+                ACTION_CHECKBOX_NAME: [account.pk],
+            },
+        )
+        messages = list(res.wsgi_request._messages)
+        self.assertEqual(
+            messages[-1].message,
+            "No stale users found in selected accounts / organizations.",
+        )
+
 
 class StripeBillingAdminTest(TestCase):
     def setUp(self):