privaterelay/migrations/0006_add_email_indexes.py

# Generated by Django 3.2.20 on 2023-11-02 14:22

from django.apps.registry import Apps
from django.db import migrations
from django.db.backends.base.schema import BaseDatabaseSchemaEditor


def add_index(
    schema_editor: BaseDatabaseSchemaEditor, index_name: str, table_name: str
) -> None:
    if schema_editor.connection.vendor.startswith("postgres"):
        engine = "postgres"
    elif schema_editor.connection.vendor.startswith("sqlite"):
        engine = "sqlite"
    else:
        raise Exception(f'Unknown database vendor "{schema_editor.connection.vendor}"')

    if_not_exists = "IF NOT EXISTS" if engine == "postgres" else ""
    schema_editor.execute(
        f"CREATE INDEX {if_not_exists} {index_name} ON {table_name} (upper(email));"
    )


def drop_index(schema_editor: BaseDatabaseSchemaEditor, index_name: str) -> None:
    schema_editor.execute(f"DROP INDEX {index_name};")


def add_account_email_index(
    apps: Apps, schema_editor: BaseDatabaseSchemaEditor
) -> None:
    add_index(schema_editor, "account_emailaddress_email_upper", "account_emailaddress")


def drop_account_email_index(
    apps: Apps, schema_editor: BaseDatabaseSchemaEditor
) -> None:
    drop_index(schema_editor, "account_emailaddress_email_upper")


def add_auth_email_index(apps: Apps, schema_editor: BaseDatabaseSchemaEditor) -> None:
    add_index(schema_editor, "auth_user_email_upper", "auth_user")


def drop_auth_email_index(apps: Apps, schema_editor: BaseDatabaseSchemaEditor) -> None:
    drop_index(schema_editor, "auth_user_email_upper")


class Migration(migrations.Migration):
    dependencies = [
        ("privaterelay", "0005_remove_flag_new_from_address"),
        ("auth", "0012_alter_user_first_name_max_length"),
        ("account", "0002_email_max_length"),
    ]

    operations = [
        migrations.RunPython(
            code=add_account_email_index,
            reverse_code=drop_account_email_index,
        ),
        migrations.RunPython(
            code=add_auth_email_index,
            reverse_code=drop_auth_email_index,
        ),
    ]