From b08bc0c7b08a0fed80b83ec94a9175afe098137d Mon Sep 17 00:00:00 2001 From: Arpad Borsos Date: Fri, 30 Aug 2024 11:54:43 +0200 Subject: [PATCH] Add an admin action to remove stale `plan_activated_users` (#778) --- codecov_auth/admin.py | 89 +++++++++++++++++++++++- codecov_auth/tests/test_admin.py | 114 +++++++++++++++++++++++++++++++ 2 files changed, 200 insertions(+), 3 deletions(-) diff --git a/codecov_auth/admin.py b/codecov_auth/admin.py index 01425bb0ae..80575e4fda 100644 --- a/codecov_auth/admin.py +++ b/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" ) diff --git a/codecov_auth/tests/test_admin.py b/codecov_auth/tests/test_admin.py index 61a42e8631..d61e013ea0 100644 --- a/codecov_auth/tests/test_admin.py +++ b/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):