upstream-zulip/zerver/migrations/0595_add_realmexport_table_and_backfill.py

# Generated by Django 5.0.9 on 2024-10-04 09:57

from datetime import datetime, timezone

import django.db.models.deletion
from django.conf import settings
from django.db import migrations, models
from django.db.backends.base.schema import BaseDatabaseSchemaEditor
from django.db.migrations.state import StateApps


def timestamp_to_datetime(timestamp: float | None) -> datetime | None:
    if timestamp is None:
        return None
    return datetime.fromtimestamp(float(timestamp), tz=timezone.utc)


def datetime_to_timestamp(dt: datetime) -> int:
    return int(dt.timestamp())


def backfill_realm_export(apps: StateApps, schema_editor: BaseDatabaseSchemaEditor) -> None:
    REALM_EXPORTED = 206
    REQUESTED = 1
    STARTED = 2
    SUCCEEDED = 3
    FAILED = 4
    DELETED = 5
    EXPORT_PUBLIC = 1

    RealmAuditLog = apps.get_model("zerver", "RealmAuditLog")
    RealmExport = apps.get_model("zerver", "RealmExport")

    for audit_log in RealmAuditLog.objects.filter(event_type=REALM_EXPORTED).order_by("id"):
        if audit_log.acting_user is None:
            # This audit_log is for an export made through shell.
            # We don't have enough data to determine if the export was
            # successful or failed.
            continue

        extra_data = audit_log.extra_data
        started_timestamp = extra_data.get("started_timestamp")
        failed_timestamp = extra_data.get("failed_timestamp")
        deleted_timestamp = extra_data.get("deleted_timestamp")
        export_path = extra_data.get("export_path")

        status = REQUESTED
        if deleted_timestamp is not None:
            status = DELETED
        elif failed_timestamp is not None:
            status = FAILED
        elif export_path is not None:
            status = SUCCEEDED
        elif started_timestamp is not None:
            status = STARTED

        # We don't have historical data to set `date_succeeded`.
        RealmExport.objects.create(
            realm=audit_log.realm,
            acting_user=audit_log.acting_user,
            type=EXPORT_PUBLIC,
            status=status,
            date_requested=audit_log.event_time,
            date_started=timestamp_to_datetime(started_timestamp),
            date_failed=timestamp_to_datetime(failed_timestamp),
            date_deleted=timestamp_to_datetime(deleted_timestamp),
            export_path=export_path,
        )


def reverse_code(apps: StateApps, schema_editor: BaseDatabaseSchemaEditor) -> None:
    RealmExport = apps.get_model("zerver", "RealmExport")

    RealmExport.objects.all().delete()


class Migration(migrations.Migration):
    dependencies = [
        ("zerver", "0594_remove_realm_user_group_edit_policy"),
    ]

    operations = [
        migrations.CreateModel(
            name="RealmExport",
            fields=[
                (
                    "id",
                    models.BigAutoField(
                        auto_created=True, primary_key=True, serialize=False, verbose_name="ID"
                    ),
                ),
                ("type", models.PositiveSmallIntegerField(default=1)),
                ("status", models.PositiveSmallIntegerField(default=1)),
                ("date_requested", models.DateTimeField()),
                ("date_started", models.DateTimeField(default=None, null=True)),
                ("date_succeeded", models.DateTimeField(default=None, null=True)),
                ("date_failed", models.DateTimeField(default=None, null=True)),
                ("date_deleted", models.DateTimeField(default=None, null=True)),
                ("export_path", models.TextField(default=None, null=True)),
                ("sha256sum_hex", models.CharField(default=None, max_length=64, null=True)),
                ("tarball_size_bytes", models.PositiveIntegerField(default=None, null=True)),
                ("stats", models.JSONField(default=None, null=True)),
                (
                    "acting_user",
                    models.ForeignKey(
                        null=True,
                        on_delete=django.db.models.deletion.SET_NULL,
                        to=settings.AUTH_USER_MODEL,
                    ),
                ),
                (
                    "realm",
                    models.ForeignKey(
                        on_delete=django.db.models.deletion.CASCADE, to="zerver.realm"
                    ),
                ),
            ],
        ),
        migrations.RunPython(backfill_realm_export, reverse_code=reverse_code, elidable=True),
    ]