# -*- coding: utf-8 -*-
# Generated by Django 1.11.14 on 2019-01-05 00:21
from __future__ import unicode_literals
from __future__ import absolute_import
from __future__ import print_function

import sys
import traceback

from django.core.management import call_command
from django.db import connections, migrations

from corehq.sql_db.operations import HqRunPython, noop_migration_fn
from corehq.sql_db.util import get_db_aliases_for_partitioned_query


BLOBMETAS_NOT_MIGRATED_ERROR = """
Blob metadata needs to be migrated before this environment can be upgraded to
the latest version of CommCareHQ. Instructions for running the migration can be
found at this link:

https://github.com/dimagi/commcare-cloud/blob/master/docs/changelog/0009-blob-metadata-part-2.md

If you are unable to run the management command because it has been deleted,
you will need to checkout an older version of CommCareHQ first:

git checkout acc40116a96a40c64efb8613fb2ba5933122b151
"""


def get_num_attachments():
    """Get the number of attachments that need to be migrated"""
    def count(dbname):
        with connections[dbname].cursor() as cursor:
            cursor.execute("SELECT COUNT(*) FROM form_processor_xformattachmentsql")
            return cursor.fetchone()[0]

    return sum(count(db) for db in get_db_aliases_for_partitioned_query())


def _assert_blobmetas_migrated(apps, schema_editor):
    """Check if migrated. Raises SystemExit if not migrated"""
    num_attachments = get_num_attachments()

    migrated = num_attachments == 0
    if migrated:
        return

    if num_attachments < 10000:
        try:
            call_command(
                "run_sql",
                "simple_move_form_attachments_to_blobmeta",
                yes=True,
            )
            migrated = get_num_attachments() == 0
            if not migrated:
                print("Automatic migration failed")
                migrated = False
        except Exception:
            traceback.print_exc()
    else:
        print("Found %s attachments." % num_attachments)
        print("Too many to migrate automatically.")

    if not migrated:
        print("")
        print(BLOBMETAS_NOT_MIGRATED_ERROR)
        sys.exit(1)


class Migration(migrations.Migration):

    dependencies = [
        ('blobs', '0006_restrict_form_attachments'),
    ]

    operations = [
        HqRunPython(_assert_blobmetas_migrated, noop_migration_fn)
    ]

Add blob metadata migration check
Add blob metadata migration check