# passerelle - uniform access to multiple data sources and services
# Copyright (C) 2017  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.contrib.postgres.fields import JSONField
from django.core.management.base import BaseCommand, CommandError
from django.db import connection


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)
