From be6d678bc56a246ae75076a93b0697e66db7d4ca Mon Sep 17 00:00:00 2001
From: Serghei Mihai <smihai@entrouvert.com>
Date: Tue, 2 Jun 2020 17:00:17 +0200
Subject: [PATCH] management: add command to ensure all JSONField fields have
 correct db type (#43501)

---
 .../base/management/commands/ensure_jsonb.py  | 50 +++++++++++++++++
 tests/test_ensure_jsonbfields.py              | 56 +++++++++++++++++++
 2 files changed, 106 insertions(+)
 create mode 100644 passerelle/base/management/commands/ensure_jsonb.py
 create mode 100644 tests/test_ensure_jsonbfields.py

diff --git a/passerelle/base/management/commands/ensure_jsonb.py b/passerelle/base/management/commands/ensure_jsonb.py
new file mode 100644
index 00000000..d9148e91
--- /dev/null
+++ b/passerelle/base/management/commands/ensure_jsonb.py
@@ -0,0 +1,50 @@
+# 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.db import connection
+from django.apps import apps
+
+from django.core.management.base import BaseCommand, CommandError
+from django.contrib.postgres.fields 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, data_type FROM information_schema.columns WHERE table_name=%s AND column_name=%s'
+                        cursor.execute(query, [table_name, column_name])
+                        for schema_name, db_type in cursor.fetchall():
+                            if db_type == 'jsonb':
+                                continue
+                            alter = 'ALTER TABLE "%(schema_name)s"."%(table_name)s" ALTER COLUMN "%(column_name)s" TYPE jsonb USING "%(column_name)s"::jsonb'
+                            params = {
+                                "schema_name": schema_name,
+                                'table_name': table_name,
+                                'column_name': column_name
+                            }
+                            try:
+                                cursor.execute(alter % params)
+                            except Exception as e:
+                                raise CommandError(e)
diff --git a/tests/test_ensure_jsonbfields.py b/tests/test_ensure_jsonbfields.py
new file mode 100644
index 00000000..ed5b3aeb
--- /dev/null
+++ b/tests/test_ensure_jsonbfields.py
@@ -0,0 +1,56 @@
+# -*- coding: utf-8 -*-
+import pytest
+
+
+from django.core.files import File
+from django.utils.six import BytesIO
+
+from django.db import connection
+
+from django.core.management import call_command
+
+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_create_reference_column(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'
-- 
2.27.0

