upstream-zulip/zerver/migrations/0383_revoke_invitations_from_deactivated_users.py

from django.db import migrations
from django.db.backends.base.schema import BaseDatabaseSchemaEditor
from django.db.migrations.state import StateApps
from django.utils.timezone import now as timezone_now


def revoke_invitations(apps: StateApps, schema_editor: BaseDatabaseSchemaEditor) -> None:
    Realm = apps.get_model("zerver", "Realm")
    Confirmation = apps.get_model("confirmation", "Confirmation")
    Confirmation.INVITATION = 2
    Confirmation.MULTIUSE_INVITE = 6
    PreregistrationUser = apps.get_model("zerver", "PreregistrationUser")
    UserProfile = apps.get_model("zerver", "UserProfile")
    MultiuseInvite = apps.get_model("zerver", "MultiuseInvite")

    STATUS_REVOKED = 2

    def get_valid_invite_confirmations_generated_by_users(
        user_ids: list[int],
    ) -> list[int]:
        prereg_user_ids = (
            PreregistrationUser.objects.filter(referred_by_id__in=user_ids)
            .exclude(status=STATUS_REVOKED)
            .values_list("id", flat=True)
        )
        confirmation_ids = list(
            Confirmation.objects.filter(
                type=Confirmation.INVITATION,
                object_id__in=prereg_user_ids,
                expiry_date__gte=timezone_now(),
            ).values_list("id", flat=True)
        )

        multiuse_invite_ids = MultiuseInvite.objects.filter(
            referred_by_id__in=user_ids
        ).values_list("id", flat=True)
        confirmation_ids += Confirmation.objects.filter(
            type=Confirmation.MULTIUSE_INVITE,
            expiry_date__gte=timezone_now(),
            object_id__in=multiuse_invite_ids,
        ).values_list("id", flat=True)

        return confirmation_ids

    print()
    for realm_id in Realm.objects.values_list("id", flat=True):
        deactivated_user_ids = UserProfile.objects.filter(
            is_active=False, realm_id=realm_id
        ).values_list("id", flat=True)
        confirmation_ids = get_valid_invite_confirmations_generated_by_users(deactivated_user_ids)

        if len(confirmation_ids) > 0:
            print(
                f"Revoking invitations by deactivated users in realm {realm_id}: {confirmation_ids}"
            )

        Confirmation.objects.filter(id__in=confirmation_ids).update(expiry_date=timezone_now())


class Migration(migrations.Migration):
    """
    User deactivation used to *not* revoke invitations generated by the user.
    This has been fixed in the implementation, but this migration is still needed
    to ensure old invitations are revoked for users who were deactivated in the past.
    """

    atomic = False

    dependencies = [
        ("zerver", "0382_create_role_based_system_groups"),
    ]

    operations = [
        migrations.RunPython(
            revoke_invitations,
            reverse_code=migrations.RunPython.noop,
            elidable=True,
        )
    ]