from django.db.backends.sqlite3.schema import DatabaseSchemaEditor \
    as BaseDatabaseSchemaEditor
from django.db.transaction import atomic


class DatabaseSchemaEditor(BaseDatabaseSchemaEditor):
    def __enter__(self):
        self.deferred_sql = []
        if self.atomic_migration:
            self.atomic = atomic(self.connection.alias)
            self.atomic.__enter__()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        if exc_type is None:
            for sql in self.deferred_sql:
                self.execute(sql)
        if self.atomic_migration:
            self.atomic.__exit__(exc_type, exc_value, traceback)
