Commit manually in Django data migration

Question

I'd like to write a data migration where I modify all rows in a big table in smaller batches in order to avoid locking issues. However, I can't figure out how to commit manually in a Django migration. Everytime I try to run commit I get:

TransactionManagementError: This is forbidden when an 'atomic' block is active.

AFAICT, the database schema editor always wraps Postgres migrations in an atomic block.

Is there a sane way to break out of the transaction from within the migration?

My migration looks like this:

def modify_data(apps, schema_editor):
    counter = 0
    BigData = apps.get_model("app", "BigData")
    for row in BigData.objects.iterator():
        # Modify row [...]
        row.save()
        # Commit every 1000 rows
        counter += 1
        if counter % 1000 == 0:
            transaction.commit()
    transaction.commit()

class Migration(migrations.Migration):
    operations = [
        migrations.RunPython(modify_data),
    ]

I'm using Django 1.7 and Postgres 9.3. This used to work with South and older versions of Django.

Answer 1

The best workaround I found is manually exiting the atomic scope before running the data migration:

def modify_data(apps, schema_editor):
    schema_editor.atomic.__exit__(None, None, None)
    # [...]

In contrast to resetting connection.in_atomic_block manually this allows using atomic context manager inside the migration. There doesn't seem to be a much saner way.

One can contain the (admittedly messy) transaction break out logic in a decorator to be used with the RunPython operation:

def non_atomic_migration(func):
  """
  Close a transaction from within code that is marked atomic. This is
  required to break out of a transaction scope that is automatically wrapped
  around each migration by the schema editor. This should only be used when
  committing manually inside a data migration. Note that it doesn't re-enter
  the atomic block afterwards.
  """
  @wraps(func)
  def wrapper(apps, schema_editor):
      if schema_editor.connection.in_atomic_block:
          schema_editor.atomic.__exit__(None, None, None)
      return func(apps, schema_editor)
  return wrapper

Update

Django 1.10 will support non-atomic migrations.