# combo - content management system
# Copyright (C) 2015-2021  Entr'ouvert
#
# This program is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

from django.apps import apps
from django.core.management.base import BaseCommand, CommandError
from django.db import connection
from django.db.models import JSONField


class Command(BaseCommand):
    help = 'Ensure all JSON fields are of type jsonb'

    def handle(self, **options):
        for app in apps.get_models():
            for field in app._meta.get_fields():
                if isinstance(field, JSONField):
                    table_name = app._meta.db_table
                    column_name = app._meta.get_field(field.name).column
                    with connection.cursor() as cursor:
                        query = '''SELECT table_schema
                                     FROM information_schema.columns
                                    WHERE table_name = %s AND column_name = %s AND data_type != %s'''
                        cursor.execute(query, [table_name, column_name, 'jsonb'])
                        for line in cursor.fetchall():
                            alter_query = '''ALTER TABLE "%(schema_name)s"."%(table_name)s"
                                            ALTER COLUMN "%(column_name)s"
                                                    TYPE jsonb USING "%(column_name)s"::jsonb'''
                            params = {
                                'schema_name': line[0],
                                'table_name': table_name,
                                'column_name': column_name,
                            }
                            try:
                                cursor.execute(alter_query % params)
                            except Exception as e:
                                raise CommandError(e)