# Generated by Django 2.2.28 on 2023-04-04 23:53

import logging

from django.db import migrations

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


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


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


def update_performance_subscriptions(apps, schema_editor):
    QuerySubscription = apps.get_model("sentry", "QuerySubscription")
    for subscription in RangeQuerySetWrapperWithProgressBar(
        QuerySubscription.objects.filter(
            snuba_query__dataset=Dataset.PerformanceMetrics.value,
            snuba_query__environment_id__isnull=False,
            status=0,
        ).select_related("snuba_query", "project")
    ):
        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)
                entity_key: EntityKey = get_entity_key_from_snuba_query(
                    subscription.snuba_query,
                    subscription.project.organization_id,
                    subscription.project_id,
                )
                # Delete the old subscription ID:
                _delete_from_snuba(
                    Dataset.PerformanceMetrics,
                    old_subscription_id,
                    entity_key,
                )
            except Exception:
                logging.exception(
                    "Failed to recreate performance 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

    dependencies = [
        ("sentry", "0406_monitor_cleanup"),
    ]

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