management: add command to ensure all JSONField fields have correct db type (#43501)

passerelle/base/management/commands/ensure_jsonb.py

@@ -0,0 +1,45 @@
+# 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.core.management.base import BaseCommand, CommandError
+from django.contrib.postgres.fields import JSONField
+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)

tests/test_ensure_jsonbfields.py

@@ -0,0 +1,51 @@
+# -*- coding: utf-8 -*-
+
+import pytest
+
+from django.core.files import File
+from django.core.management import call_command
+from django.db import connection
+from django.utils.six import BytesIO
+
+from passerelle.apps.csvdatasource.models import CsvDataSource
+from passerelle.contrib.teamnet_axel.models import TeamnetAxel
+
+pytestmark = pytest.mark.django_db
+
+@pytest.fixture
+def setup():
+
+    def maker(columns_keynames='fam,id,lname,fname,sex', filename='data.csv', sheet_name='Feuille2',
+              data=b''):
+        csv = CsvDataSource.objects.create(csv_file=File(BytesIO(data), filename),
+                                           sheet_name=sheet_name, columns_keynames=columns_keynames,
+                                           slug='test', title='a title',
+                                           description='a description')
+        teamnet = TeamnetAxel.objects.create(slug='test', billing_regies={},
+                                             wsdl_url='http://example.net/AXEL_WS/AxelWS.php?wsdl')
+        return csv, teamnet
+    return maker
+
+
+def test_ensure_jsonb_fields(setup):
+    with connection.cursor() as cursor:
+        query = "SELECT table_name, column_name, data_type FROM information_schema.columns WHERE column_name IN ('_dialect_options', 'billing_regies')"
+        cursor.execute(query)
+
+        # make sure the data_type is correct
+        for line in cursor.fetchall():
+            assert line[2] == 'jsonb'
+
+        # alter columns
+        cursor.execute('ALTER TABLE csvdatasource_csvdatasource ALTER COLUMN _dialect_options TYPE text USING _dialect_options::text')
+        cursor.execute('ALTER TABLE teamnet_axel_teamnetaxel ALTER COLUMN billing_regies TYPE text USING billing_regies::text')
+
+    call_command('ensure_jsonb')
+
+    with connection.cursor() as cursor:
+        query = "SELECT table_name, column_name, data_type FROM information_schema.columns WHERE column_name IN ('_dialect_options', 'billing_regies')"
+        cursor.execute(query)
+
+        # check the data_type is correct
+        for line in cursor.fetchall():
+            assert line[2] == 'jsonb'