# Generated by Django 2.2.28 on 2022-07-08 22:53

from django.db import migrations, models

from sentry.new_migrations.migrations import CheckedMigration
from sentry.utils.query import RangeQuerySetWrapperWithProgressBar

# Keys from `Dataset`, values from `SnubaQuery.Type`
dataset_to_type_map = {
    "events": 0,
    "transactions": 1,
    # There are some inactive alert rules that still have `sessions` as their dataset. We want to
    # still fill in type for these rows
    "sessions": 2,
    "metrics": 2,
}


def backfill_snubaquery_type(apps, schema_editor):
    SnubaQuery = apps.get_model("sentry", "SnubaQuery")

    for snuba_query in RangeQuerySetWrapperWithProgressBar(SnubaQuery.objects.all()):
        if snuba_query.type is not None:
            continue

        snuba_query.type = dataset_to_type_map[snuba_query.dataset]
        snuba_query.save(update_fields=["type"])


class Migration(CheckedMigration):
    # This flag is used to mark that a migration shouldn't be automatically run in production. For
    # the most part, this should only be used for operations where it's safe to run the migration
    # after your code has deployed. So this should not be used for most operations that alter the
    # schema of a table.
    # Here are some things that make sense to mark as dangerous:
    # - Large data migrations. Typically we want these to be run manually by ops so that they can
    #   be monitored and not block the deploy for a long period of time while they run.
    # - Adding indexes to large tables. Since this can take a long time, we'd generally prefer to
    #   have ops run this and not block the deploy. Note that while adding an index is a schema
    #   change, it's completely safe to run the operation after the code has deployed.
    is_dangerous = False

    # This flag is used to decide whether to run this migration in a transaction or not. Generally
    # we don't want to run in a transaction here, since for long running operations like data
    # back-fills this results in us locking an increasing number of rows until we finally commit.
    atomic = False

    dependencies = [
        ("sentry", "0301_mep_move_type_to_snuba_query"),
    ]

    operations = [
        migrations.RunPython(
            backfill_snubaquery_type,
            migrations.RunPython.noop,
            hints={"tables": ["sentry_snubaquery"]},
        ),
        migrations.AlterField(
            model_name="snubaquery",
            name="type",
            field=models.SmallIntegerField(),
        ),
    ]
