# Generated by Django 2.2.28 on 2022-06-15 11:14
# Based on https://github.com/getsentry/getsentry/blob/89ff1453be755ddef31f2b99de09bd03badeb25e/getsentry/migrations/0141_migrate_sessions_subs_to_metrics.py

import logging
import re

from django.db import migrations

from sentry.new_migrations.migrations import CheckedMigration
from sentry.snuba.dataset import Dataset, EntityKey
from sentry.snuba.tasks import _create_in_snuba, _delete_from_snuba
from sentry.utils.query import RangeQuerySetWrapperWithProgressBar

CRASH_RATE_ALERT_AGGREGATE_RE = (
    r"^percentage\([ ]*(sessions_crashed|users_crashed)[ ]*\,[ ]*(sessions|users)[ ]*\)"
)


def create_subscription_in_snuba(subscription):
    subscription.subscription_id = _create_in_snuba(subscription)
    subscription.save()


def map_aggregate_to_entity_key(dataset: Dataset, aggregate: str) -> EntityKey:
    if dataset == Dataset.Events:
        entity_key = EntityKey.Events
    elif dataset == Dataset.Transactions:
        entity_key = EntityKey.Transactions
    elif dataset in [Dataset.Metrics, Dataset.Sessions]:
        match = re.match(CRASH_RATE_ALERT_AGGREGATE_RE, aggregate)
        if not match:
            raise Exception(
                f"Only crash free percentage queries are supported for subscriptions"
                f"over the {dataset.value} dataset"
            )
        if dataset == Dataset.Metrics:
            count_col_matched = match.group(2)
            if count_col_matched == "sessions":
                entity_key = EntityKey.MetricsCounters
            else:
                entity_key = EntityKey.MetricsSets
        else:
            entity_key = EntityKey.Sessions
    else:
        raise Exception(f"{dataset} dataset does not have an entity key mapped to it")
    return entity_key


def delete_subscription_from_snuba(subscription, query_dataset: Dataset):
    entity_key: EntityKey = map_aggregate_to_entity_key(
        query_dataset, subscription.snuba_query.aggregate
    )
    _delete_from_snuba(
        query_dataset,
        subscription.subscription_id,
        entity_key,
    )


def event_types(self):
    return [type.event_type for type in self.snubaqueryeventtype_set.all()]


def update_metrics_subscriptions(apps, schema_editor):
    QuerySubscription = apps.get_model("sentry", "QuerySubscription")
    for subscription in RangeQuerySetWrapperWithProgressBar(
        QuerySubscription.objects.filter(snuba_query__dataset=Dataset.Metrics.value).select_related(
            "snuba_query"
        )
    ):
        old_subscription_id = subscription.subscription_id
        if old_subscription_id is not None:
            try:
                # The migration apps don't build this property, so patch it here:
                subscription.snuba_query.event_types = property(event_types)
                create_subscription_in_snuba(subscription)
                delete_subscription_from_snuba(subscription, Dataset.Metrics)
            except Exception:
                logging.exception(
                    "Failed to recreate metrics subscription in snuba",
                    extra={
                        "project": subscription.project.slug,
                        "subscription_id": subscription.id,
                        "query": subscription.snuba_query.query,
                        "aggregate": subscription.snuba_query.aggregate,
                        "time_window": subscription.snuba_query.time_window,
                        "resolution": subscription.snuba_query.resolution,
                    },
                )


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", "0291_add_new_perf_indexer"),
    ]

    operations = [
        migrations.RunPython(
            update_metrics_subscriptions,
            migrations.RunPython.noop,
            hints={"tables": ["sentry_querysubscription", "sentry_snubaquery"]},
        ),
    ]
