Fixed #25833 -- Added support for non-atomic migrations.

Added the Migration.atomic attribute which can be set to False
for non-atomic migrations.

django/core/management/commands/sqlmigrate.py

@@ -51,6 +51,9 @@ class Command(BaseCommand):
                 migration_name, app_label))
         targets = [(app_label, migration.name)]
 
+        # Show begin/end around output only for atomic migrations
+        self.output_transaction = migration.atomic
+
         # Make a plan that represents just the requested migrations and show SQL
         # for it
         plan = [(executor.loader.graph.nodes[targets[0]], options['backwards'])]

django/db/backends/base/schema.py

@@ -69,17 +69,18 @@ class BaseDatabaseSchemaEditor(object):
     sql_create_pk = "ALTER TABLE %(table)s ADD CONSTRAINT %(name)s PRIMARY KEY (%(columns)s)"
     sql_delete_pk = "ALTER TABLE %(table)s DROP CONSTRAINT %(name)s"
 
-    def __init__(self, connection, collect_sql=False):
+    def __init__(self, connection, collect_sql=False, atomic=True):
         self.connection = connection
         self.collect_sql = collect_sql
         if self.collect_sql:
             self.collected_sql = []
+        self.atomic_migration = self.connection.features.can_rollback_ddl and atomic
 
     # State-managing methods
 
     def __enter__(self):
         self.deferred_sql = []
-        if self.connection.features.can_rollback_ddl:
+        if self.atomic_migration:
             self.atomic = atomic(self.connection.alias)
             self.atomic.__enter__()
         return self
@@ -88,7 +89,7 @@ class BaseDatabaseSchemaEditor(object):
         if exc_type is None:
             for sql in self.deferred_sql:
                 self.execute(sql)
-        if self.connection.features.can_rollback_ddl:
+        if self.atomic_migration:
             self.atomic.__exit__(exc_type, exc_value, traceback)
 
     # Core utility functions

django/db/migrations/executor.py

@@ -170,7 +170,7 @@ class MigrationExecutor(object):
         statements = []
         state = None
         for migration, backwards in plan:
-            with self.connection.schema_editor(collect_sql=True) as schema_editor:
+            with self.connection.schema_editor(collect_sql=True, atomic=migration.atomic) as schema_editor:
                 if state is None:
                     state = self.loader.project_state((migration.app_label, migration.name), at_end=False)
                 if not backwards:
@@ -194,7 +194,7 @@ class MigrationExecutor(object):
                     fake = True
             if not fake:
                 # Alright, do it normally
-                with self.connection.schema_editor() as schema_editor:
+                with self.connection.schema_editor(atomic=migration.atomic) as schema_editor:
                     state = migration.apply(state, schema_editor)
         # For replacement migrations, record individual statuses
         if migration.replaces:
@@ -214,7 +214,7 @@ class MigrationExecutor(object):
         if self.progress_callback:
             self.progress_callback("unapply_start", migration, fake)
         if not fake:
-            with self.connection.schema_editor() as schema_editor:
+            with self.connection.schema_editor(atomic=migration.atomic) as schema_editor:
                 state = migration.unapply(state, schema_editor)
         # For replacement migrations, record individual statuses
         if migration.replaces:

django/db/migrations/migration.py

@@ -48,6 +48,10 @@ class Migration(object):
     # introspection. If False, never perform introspection.
     initial = None
 
+    # Whether to wrap the whole migration in a transaction. Only has an effect
+    # on database backends which support transactional DDL.
+    atomic = True
+
     def __init__(self, name, app_label):
         self.name = name
         self.app_label = app_label
@@ -114,8 +118,10 @@ class Migration(object):
             old_state = project_state.clone()
             operation.state_forwards(self.app_label, project_state)
             # Run the operation
-            if not schema_editor.connection.features.can_rollback_ddl and operation.atomic:
-                # We're forcing a transaction on a non-transactional-DDL backend
+            atomic_operation = operation.atomic or (self.atomic and operation.atomic is not False)
+            if not schema_editor.atomic_migration and atomic_operation:
+                # Force a transaction on a non-transactional-DDL backend or an
+                # atomic operation inside a non-atomic migration.
                 with atomic(schema_editor.connection.alias):
                     operation.database_forwards(self.app_label, schema_editor, old_state, project_state)
             else:

django/db/migrations/operations/special.py

@@ -139,7 +139,7 @@ class RunPython(Operation):
 
     reduces_to_sql = False
 
-    def __init__(self, code, reverse_code=None, atomic=True, hints=None, elidable=False):
+    def __init__(self, code, reverse_code=None, atomic=None, hints=None, elidable=False):
         self.atomic = atomic
         # Forwards code
         if not callable(code):
@@ -161,7 +161,7 @@ class RunPython(Operation):
         }
         if self.reverse_code is not None:
             kwargs['reverse_code'] = self.reverse_code
-        if self.atomic is not True:
+        if self.atomic is not None:
             kwargs['atomic'] = self.atomic
         if self.hints:
             kwargs['hints'] = self.hints

docs/howto/writing-migrations.txt

@@ -184,6 +184,53 @@ the respective field according to your needs.
   migration is running. Objects created after the ``AddField`` and before
   ``RunPython`` will have their original ``uuid``’s overwritten.
 
+.. _non-atomic-migrations:
+
+Non-atomic migrations
+~~~~~~~~~~~~~~~~~~~~~
+
+.. versionadded:: 1.10
+
+On databases that support DDL transactions (SQLite and PostgreSQL), migrations
+will run inside a transaction by default. For use cases such as performing data
+migrations on large tables, you may want to prevent a migration from running in
+a transaction by setting the ``atomic`` attribute to ``False``::
+
+    from django.db import migrations
+
+    class Migration(migrations.Migration):
+        atomic = False
+
+Within such a migration, all operations are run without a transaction. It's
+possible to execute parts of the migration inside a transaction using
+:func:`~django.db.transaction.atomic()` or by passing ``atomic=True`` to
+``RunPython``.
+
+Here's an example of a non-atomic data migration that updates a large table in
+smaller batches::
+
+    import uuid
+
+    from django.db import migrations, transaction
+
+    def gen_uuid(apps, schema_editor):
+        MyModel = apps.get_model('myapp', 'MyModel')
+        while MyModel.objects.filter(uuid__isnull=True).exists():
+            with transaction.atomic():
+                for row in MyModel.objects.filter(uuid__isnull=True)[:1000]:
+                    row.uuid = uuid.uuid4()
+                    row.save()
+
+    class Migration(migrations.Migration):
+        atomic = False
+
+        operations = [
+            migrations.RunPython(gen_uuid),
+        ]
+
+The ``atomic`` attribute doesn't have an effect on databases that don't support
+DDL transactions (e.g. MySQL, Oracle).
+
 Controlling the order of migrations
 ===================================
 

docs/ref/migration-operations.txt

@@ -269,7 +269,7 @@ be removed (elided) when :ref:`squashing migrations <migration-squashing>`.
 ``RunPython``
 -------------
 
-.. class:: RunPython(code, reverse_code=None, atomic=True, hints=None, elidable=False)
+.. class:: RunPython(code, reverse_code=None, atomic=None, hints=None, elidable=False)
 
 Runs custom Python code in a historical context. ``code`` (and ``reverse_code``
 if supplied) should be callable objects that accept two arguments; the first is
@@ -354,16 +354,19 @@ the ``schema_editor`` provided on these backends; in this case, pass
 
 On databases that do support DDL transactions (SQLite and PostgreSQL),
 ``RunPython`` operations do not have any transactions automatically added
-besides the transactions created for each migration (the ``atomic`` parameter
-has no effect on these databases). Thus, on PostgreSQL, for example, you should
-avoid combining schema changes and ``RunPython`` operations in the same
-migration or you may hit errors like ``OperationalError: cannot ALTER TABLE
-"mytable" because it has pending trigger events``.
+besides the transactions created for each migration. Thus, on PostgreSQL, for
+example, you should avoid combining schema changes and ``RunPython`` operations
+in the same migration or you may hit errors like ``OperationalError: cannot
+ALTER TABLE "mytable" because it has pending trigger events``.
 
 If you have a different database and aren't sure if it supports DDL
 transactions, check the ``django.db.connection.features.can_rollback_ddl``
 attribute.
 
+If the ``RunPython`` operation is part of a :ref:`non-atomic migration
+<non-atomic-migrations>`, the operation will only be executed in a transaction
+if ``atomic=True`` is passed to the ``RunPython`` operation.
+
 .. warning::
 
     ``RunPython`` does not magically alter the connection of the models for you;
@@ -382,6 +385,11 @@ attribute.
 
     The ``elidable`` argument was added.
 
+.. versionchanged:: 1.10
+
+    The ``atomic`` argument default was changed to ``None``, indicating that
+    the atomicity is controlled by the ``atomic`` attribute of the migration.
+
 ``SeparateDatabaseAndState``
 ----------------------------
 

docs/releases/1.10.txt

@@ -258,6 +258,9 @@ Migrations
   :class:`~django.db.migrations.operations.RunPython` operations to allow them
   to be removed when squashing migrations.
 
+* Added support for :ref:`non-atomic migrations <non-atomic-migrations>` by
+  setting the ``atomic`` attribute on a ``Migration``.
+
 Models
 ~~~~~~
 

tests/migrations/test_commands.py

@@ -362,6 +362,19 @@ class MigrateTests(MigrationTestBase):
         # Cleanup by unmigrating everything
         call_command("migrate", "migrations", "zero", verbosity=0)
 
+    @override_settings(MIGRATION_MODULES={"migrations": "migrations.test_migrations_non_atomic"})
+    def test_sqlmigrate_for_non_atomic_migration(self):
+        """
+        Transaction wrappers aren't shown for non-atomic migrations.
+        """
+        out = six.StringIO()
+        call_command("sqlmigrate", "migrations", "0001", stdout=out)
+        output = out.getvalue().lower()
+        queries = [q.strip() for q in output.splitlines()]
+        if connection.ops.start_transaction_sql():
+            self.assertNotIn(connection.ops.start_transaction_sql().lower(), queries)
+        self.assertNotIn(connection.ops.end_transaction_sql().lower(), queries)
+
     @override_settings(
         INSTALLED_APPS=[
             "migrations.migrations_test_apps.migrated_app",

tests/migrations/test_executor.py

@@ -100,6 +100,33 @@ class ExecutorTests(MigrationTestBase):
         self.assertTableNotExists("migrations_author")
         self.assertTableNotExists("migrations_book")
 
+    @override_settings(MIGRATION_MODULES={"migrations": "migrations.test_migrations_non_atomic"})
+    def test_non_atomic_migration(self):
+        """
+        Applying a non-atomic migration works as expected.
+        """
+        executor = MigrationExecutor(connection)
+        with self.assertRaisesMessage(RuntimeError, "Abort migration"):
+            executor.migrate([("migrations", "0001_initial")])
+        self.assertTableExists("migrations_publisher")
+        migrations_apps = executor.loader.project_state(("migrations", "0001_initial")).apps
+        Publisher = migrations_apps.get_model("migrations", "Publisher")
+        self.assertTrue(Publisher.objects.exists())
+        self.assertTableNotExists("migrations_book")
+
+    @override_settings(MIGRATION_MODULES={"migrations": "migrations.test_migrations_atomic_operation"})
+    def test_atomic_operation_in_non_atomic_migration(self):
+        """
+        An atomic operation is properly rolled back inside a non-atomic
+        migration.
+        """
+        executor = MigrationExecutor(connection)
+        with self.assertRaisesMessage(RuntimeError, "Abort migration"):
+            executor.migrate([("migrations", "0001_initial")])
+        migrations_apps = executor.loader.project_state(("migrations", "0001_initial")).apps
+        Editor = migrations_apps.get_model("migrations", "Editor")
+        self.assertFalse(Editor.objects.exists())
+
     @override_settings(MIGRATION_MODULES={
         "migrations": "migrations.test_migrations",
         "migrations2": "migrations2.test_migrations_2",

tests/migrations/test_migrations_atomic_operation/0001_initial.py

@@ -0,0 +1,25 @@
+# -*- coding: utf-8 -*-
+from __future__ import unicode_literals
+
+from django.db import migrations, models
+
+
+def raise_error(apps, schema_editor):
+    # Test atomic operation in non-atomic migration is wrapped in transaction
+    Editor = apps.get_model('migrations', 'Editor')
+    Editor.objects.create(name='Test Editor')
+    raise RuntimeError('Abort migration')
+
+
+class Migration(migrations.Migration):
+    atomic = False
+
+    operations = [
+        migrations.CreateModel(
+            "Editor",
+            [
+                ("name", models.CharField(primary_key=True, max_length=255)),
+            ],
+        ),
+        migrations.RunPython(raise_error, atomic=True),
+    ]

tests/migrations/test_migrations_non_atomic/0001_initial.py

@@ -0,0 +1,32 @@
+# -*- coding: utf-8 -*-
+from __future__ import unicode_literals
+
+from django.db import migrations, models
+
+
+def raise_error(apps, schema_editor):
+    # Test operation in non-atomic migration is not wrapped in transaction
+    Publisher = apps.get_model('migrations', 'Publisher')
+    Publisher.objects.create(name='Test Publisher')
+    raise RuntimeError('Abort migration')
+
+
+class Migration(migrations.Migration):
+    atomic = False
+
+    operations = [
+        migrations.CreateModel(
+            "Publisher",
+            [
+                ("name", models.CharField(primary_key=True, max_length=255)),
+            ],
+        ),
+        migrations.RunPython(raise_error),
+        migrations.CreateModel(
+            "Book",
+            [
+                ("title", models.CharField(primary_key=True, max_length=255)),
+                ("publisher", models.ForeignKey("migrations.Publisher", models.SET_NULL, null=True)),
+            ],
+        ),
+    ]