Repository: adamchainz/django-mysql

Issue:
Command to output sql for converting datetime to datetime(6)
Loop over all django models with `DateTimeField`s and see if they are `datetime` in the database. If so, output the `ALTER TABLE` statement required to convert them to `datetime(6)`.



Discussion:

Resolving patch:
diff --git a/django_mysql/management/commands/fix_datetime_columns.py b/django_mysql/management/commands/fix_datetime_columns.py
new file mode 100644
--- /dev/null
+++ b/django_mysql/management/commands/fix_datetime_columns.py
@@ -0,0 +1,125 @@
+# -*- coding: utf-8 -*-
+from __future__ import unicode_literals
+
+import django
+from django.apps import apps
+from django.core.management import BaseCommand, CommandError
+from django.db import DEFAULT_DB_ALIAS, connections
+from django.db.utils import ConnectionDoesNotExist
+
+from django_mysql.utils import collapse_spaces
+
+
+class Command(BaseCommand):
+    args = "<optional connection alias>"
+
+    help = collapse_spaces("""
+        Detects DateTimeFields with column type 'datetime' instead of
+        'datetime(6)' and outputs the SQL to fix them.
+    """)
+
+    if django.VERSION[:2] >= (1, 10):
+        def add_arguments(self, parser):
+            parser.add_argument(
+                'alias', metavar='alias', nargs='?',
+                default=DEFAULT_DB_ALIAS,
+                help='Specify the database connection alias to output '
+                     'parameters for.'
+            )
+
+    def handle(self, *args, **options):
+        if django.VERSION[:2] >= (1, 10):
+            alias = options['alias']
+        else:
+            if len(args) > 1:
+                raise CommandError('Cannot connect to more than one '
+                                   'connection.')
+            elif len(args) == 0:
+                alias = DEFAULT_DB_ALIAS
+            else:
+                alias = args[0]
+
+        try:
+            connection = connections[alias]
+        except ConnectionDoesNotExist:
+            raise CommandError("Connection '{}' does not exist".format(alias))
+
+        if connection.vendor != 'mysql':
+            raise CommandError("{} is not a MySQL database connection"
+                               .format(alias))
+
+        sqls = []
+        with connection.cursor() as cursor:
+            for table_name in self.all_table_names():
+                sql = self.datetime_fix_sql(connection, cursor, table_name)
+                if sql:
+                    sqls.append(sql)
+
+        for sql in sqls:
+            self.stdout.write(sql)
+
+    def all_table_names(self):
+        table_names = set()
+        for app_config in apps.get_app_configs():
+            for model in app_config.get_models():
+                table_names.add(model._meta.db_table)
+        return list(sorted(table_names))
+
+    def datetime_fix_sql(self, connection, cursor, table_name):
+        cursor.execute(
+            """
+            SELECT COLUMN_NAME
+            FROM INFORMATION_SCHEMA.COLUMNS
+            WHERE TABLE_SCHEMA = DATABASE() AND
+                  TABLE_NAME = %s AND
+                  DATA_TYPE = 'datetime' AND
+                  DATETIME_PRECISION = 0
+            ORDER BY COLUMN_NAME
+            """,
+            (table_name,)
+        )
+        bad_column_names = [r[0] for r in cursor.fetchall()]
+        if not bad_column_names:
+            return
+
+        qn = connection.ops.quote_name
+
+        cursor.execute("SHOW CREATE TABLE {}".format(qn(table_name)))
+        create_table = cursor.fetchone()[1]
+        column_specs = parse_create_table(create_table)
+
+        modify_columns = []
+
+        for column_name in bad_column_names:
+            column_spec = column_specs[column_name]
+
+            new_column_spec = column_spec.replace('datetime', 'datetime(6)', 1)
+            modify_columns.append(
+                'MODIFY COLUMN {} {}'.format(qn(column_name), new_column_spec)
+            )
+
+        return 'ALTER TABLE {table_name}\n    {columns};'.format(
+            table_name=qn(table_name),
+            columns=',\n    '.join(modify_columns),
+        )
+
+
+def parse_create_table(sql):
+    """
+    Split output of SHOW CREATE TABLE into {column: column_spec}
+    """
+    column_types = {}
+    for line in sql.splitlines()[1:]:  # first line = CREATE TABLE
+        sline = line.strip()
+
+        if not sline.startswith('`'):
+            # We've finished parsing the columns
+            break
+
+        bits = sline.split('`')
+        assert len(bits) == 3
+        column_name = bits[1]
+        column_spec = bits[2].lstrip().rstrip(',')
+
+        column_types[column_name] = column_spec
+    return column_types


Test patch:
diff --git a/tests/testapp/management/__init__.py b/tests/testapp/management/__init__.py
new file mode 100644
diff --git a/tests/testapp/management/commands/__init__.py b/tests/testapp/management/commands/__init__.py
new file mode 100644
diff --git a/tests/testapp/management/commands/test_fix_datetime_columns.py b/tests/testapp/management/commands/test_fix_datetime_columns.py
new file mode 100644
--- /dev/null
+++ b/tests/testapp/management/commands/test_fix_datetime_columns.py
@@ -0,0 +1,159 @@
+# -*- coding:utf-8 -*-
+from __future__ import unicode_literals
+
+import mock
+from textwrap import dedent
+from unittest import SkipTest
+
+import pytest
+from django.core.management import CommandError, call_command
+from django.db import connection
+from django.db.utils import ConnectionHandler
+from django.test import SimpleTestCase, TestCase, TransactionTestCase
+from django.utils.six.moves import StringIO
+
+from django_mysql.management.commands.fix_datetime_columns import \
+    parse_create_table
+
+# Can't use @override_settings to swap out DATABASES, instead just mock.patch
+# a new ConnectionHandler into the command module
+
+command_connections = (
+    'django_mysql.management.commands.fix_datetime_columns.connections'
+)
+
+sqlite = ConnectionHandler({
+    'default': {'ENGINE': 'django.db.backends.sqlite3'}
+})
+
+
+def run_it(*args, **kwargs):
+    run_args = ['fix_datetime_columns']
+    run_args.extend(args)
+
+    out = StringIO()
+    run_kwargs = {'stdout': out, 'skip_checks': True}
+    run_kwargs.update(kwargs)
+
+    call_command(*run_args, **run_kwargs)
+
+    return out.getvalue()
+
+
+class Datetime6TestMixin(object):
+
+    @classmethod
+    def setUpClass(cls):
+        if (
+            connection.is_mariadb or
+            connection.mysql_version[:2] < (5, 6)
+        ):
+            raise SkipTest(
+                "Django only uses datetime(6) columns on MySQL 5.6+"
+            )
+        super(Datetime6TestMixin, cls).setUpClass()
+
+
+class FexDatetimeColumnsTests(Datetime6TestMixin, TestCase):
+
+    def test_nothing_by_default(self):
+        assert run_it() == ''
+
+    def test_nothing_by_default_alternative_connection(self):
+        assert run_it('other') == ''
+
+    def test_invalid_database(self):
+        with pytest.raises(CommandError) as excinfo:
+            run_it('bla')
+
+        assert "does not exist" in str(excinfo.value)
+
+    def test_invalid_number_of_databases(self):
+        with pytest.raises(CommandError) as excinfo:
+            run_it('other', 'default')
+        assert "more than one connection" in str(excinfo.value)
+
+    @mock.patch(command_connections, sqlite)
+    def test_invalid_not_mysql(self):
+        with pytest.raises(CommandError) as excinfo:
+            run_it()
+        assert "not a MySQL database connection" in str(excinfo.value)
+
+
+class SlowFexDatetimeColumnsTests(Datetime6TestMixin, TransactionTestCase):
+
+    def test_with_one_column_to_migrate(self):
+        with connection.cursor() as cursor:
+            cursor.execute("""
+                ALTER TABLE testapp_author
+                    MODIFY COLUMN birthday datetime DEFAULT NULL
+            """)
+            try:
+                out = run_it()
+            finally:
+                cursor.execute("""
+                    ALTER TABLE testapp_author
+                        MODIFY COLUMN birthday datetime(6) DEFAULT NULL
+                """)
+        assert out == dedent('''\
+            ALTER TABLE `testapp_author`
+                MODIFY COLUMN `birthday` datetime(6) DEFAULT NULL;
+            ''')
+
+    def test_with_two_columns_to_migrate(self):
+        with connection.cursor() as cursor:
+            cursor.execute("""
+                ALTER TABLE testapp_author
+                    MODIFY COLUMN birthday datetime DEFAULT NULL,
+                    MODIFY COLUMN deathday datetime DEFAULT NULL
+            """)
+            try:
+                out = run_it()
+            finally:
+                cursor.execute("""
+                    ALTER TABLE testapp_author
+                        MODIFY COLUMN birthday datetime(6) DEFAULT NULL,
+                        MODIFY COLUMN deathday datetime(6) DEFAULT NULL
+                """)
+        assert out == dedent('''\
+            ALTER TABLE `testapp_author`
+                MODIFY COLUMN `birthday` datetime(6) DEFAULT NULL,
+                MODIFY COLUMN `deathday` datetime(6) DEFAULT NULL;
+            ''')
+
+
+class ParseCreateTableTests(SimpleTestCase):
+
+    def test_large_normalish_table(self):
+        sql = dedent("""\
+        CREATE TABLE `example` (
+          `id` int(11) NOT NULL AUTO_INCREMENT,
+          `varchary` varchar(255) NOT NULL,
+          `texty` longtext NOT NULL,
+          `inty` int(11) DEFAULT NULL,
+          `datetimey` datetime NOT NULL,
+          PRIMARY KEY (`id`),
+          KEY `example_abcdef` (`varchary`),
+          KEY `example_123456` (`inty`)
+          CONSTRAINT `example_constraint` FOREIGN KEY (`id`) REFERENCES `other` (`id`)
+        ) ENGINE=InnoDB AUTO_INCREMENT=999999 DEFAULT CHARSET=utf8mb4 ROW_FORMAT=COMPRESSED
+        """)  # noqa
+
+        assert parse_create_table(sql) == {
+            'id': 'int(11) NOT NULL AUTO_INCREMENT',
+            'varchary': 'varchar(255) NOT NULL',
+            'texty': 'longtext NOT NULL',
+            'inty': 'int(11) DEFAULT NULL',
+            'datetimey': 'datetime NOT NULL',
+        }
+
+    def test_single_column(self):
+        sql = dedent("""\
+        CREATE TABLE `example_table` (
+          `the_data` longtext COLLATE utf8mb4_unicode_ci NOT NULL
+        ) ENGINE=MyISAM DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci COMMENT='hi mum';
+        """)  # noqa
+
+        assert parse_create_table(sql) == {
+            'the_data': 'longtext COLLATE utf8mb4_unicode_ci NOT NULL'
+        }
diff --git a/tests/testapp/models.py b/tests/testapp/models.py
--- a/tests/testapp/models.py
+++ b/tests/testapp/models.py
@@ -150,6 +150,8 @@ class Author(Model):
     name = CharField(max_length=32, db_index=True)
     tutor = ForeignKey('self', null=True, related_name='tutees')
     bio = TextField()
+    birthday = DateTimeField(null=True)
+    deathday = DateTimeField(null=True)
 
     def __unicode__(self):
         return "{} {}".format(self.id, self.name)