# Generated by Django 2.2.24 on 2022-03-04 16:41

import re

from django.db import migrations

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

AGGREGATE_PATTERN = r"^(\w+)\((.*)?\)$"
AGGREGATE_BASE = r".*(\w+)\((.*)?\)"
EQUATION_PREFIX = "equation|"


def is_equation(field: str) -> bool:
    """check if a public alias is an equation, which start with the equation prefix
    eg. `equation|5 + 5`
    """
    return field.startswith(EQUATION_PREFIX)


def is_aggregate(field: str) -> bool:
    field_match = re.match(AGGREGATE_PATTERN, field)
    if field_match:
        return True

    equation_match = re.match(AGGREGATE_BASE, field)
    if equation_match and is_equation(field):
        return True

    return False


def backfill_columns_aggregates(apps, schema_editor):
    DashboardWidgetQuery = apps.get_model("sentry", "DashboardWidgetQuery")

    for widget_query in RangeQuerySetWrapperWithProgressBar(DashboardWidgetQuery.objects.all()):
        if widget_query.columns or widget_query.aggregates:
            continue

        fields = widget_query.fields or []
        columns = []
        aggregates = []

        for field in fields:
            if is_aggregate(field):
                aggregates.append(field)
            else:
                columns.append(field)

        widget_query.columns = columns
        widget_query.aggregates = aggregates

        widget_query.save()


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", "0276_rulefirehistory_date_added_index"),
    ]

    operations = [
        migrations.RunPython(
            backfill_columns_aggregates,
            migrations.RunPython.noop,
            hints={"tables": ["sentry_dashboardwidgetquery"]},
        )
    ]
