zulip/zerver/migrations/0093_subscription_event_log_backfill.py

# -*- coding: utf-8 -*-

from typing import List

from django.db import migrations, models
from django.db.backends.postgresql_psycopg2.schema import DatabaseSchemaEditor
from django.db.migrations.state import StateApps
from django.db.models import Max
from django.utils.timezone import now as timezone_now

def backfill_subscription_log_events(apps: StateApps, schema_editor: DatabaseSchemaEditor) -> None:
    migration_time = timezone_now()
    RealmAuditLog = apps.get_model('zerver', 'RealmAuditLog')
    Subscription = apps.get_model('zerver', 'Subscription')
    Message = apps.get_model('zerver', 'Message')
    objects_to_create = []

    subs_query = Subscription.objects.select_related(
        "user_profile", "user_profile__realm", "recipient").filter(recipient__type=2)
    for sub in subs_query:
        entry = RealmAuditLog(
            realm=sub.user_profile.realm,
            modified_user=sub.user_profile,
            modified_stream_id=sub.recipient.type_id,
            event_last_message_id=0,
            event_type='subscription_created',
            event_time=migration_time,
            backfilled=True)
        objects_to_create.append(entry)
    RealmAuditLog.objects.bulk_create(objects_to_create)
    objects_to_create = []

    event_last_message_id = Message.objects.aggregate(Max('id'))['id__max']
    migration_time_for_deactivation = timezone_now()
    for sub in subs_query.filter(active=False):
        entry = RealmAuditLog(
            realm=sub.user_profile.realm,
            modified_user=sub.user_profile,
            modified_stream_id=sub.recipient.type_id,
            event_last_message_id=event_last_message_id,
            event_type='subscription_deactivated',
            event_time=migration_time_for_deactivation,
            backfilled=True)
        objects_to_create.append(entry)
    RealmAuditLog.objects.bulk_create(objects_to_create)
    objects_to_create = []

def reverse_code(apps: StateApps, schema_editor: DatabaseSchemaEditor) -> None:
    RealmAuditLog = apps.get_model('zerver', 'RealmAuditLog')
    RealmAuditLog.objects.filter(event_type='subscription_created').delete()
    RealmAuditLog.objects.filter(event_type='subscription_deactivated').delete()


class Migration(migrations.Migration):

    dependencies = [
        ('zerver', '0092_create_scheduledemail'),
    ]

    operations = [
        migrations.AddField(
            model_name='realmauditlog',
            name='event_last_message_id',
            field=models.IntegerField(null=True),
        ),
        migrations.RunPython(backfill_subscription_log_events,
                             reverse_code=reverse_code),
    ]