How to restore the database in the case of an IntegrityError error in Django?

advertisements

The following static method (taken from flasky) is a way to populate the database with fake data, and I want to import it to Django.

models.py

class User(UserMixin, db.Model):

    # ......

    @staticmethod
    def generate_fake(count=100):
        from sqlalchemy.exc import IntegrityError
        from random import seed
        import forgery_py

        seed()
        for i in range(count):
            u = User(email=forgery_py.internet.email_address(),
                     username=forgery_py.internet.user_name(True),
                     password=forgery_py.lorem_ipsum.word(),
                     confirmed=True,
                     name=forgery_py.name.full_name(),
                     location=forgery_py.address.city(),
                     about_me=forgery_py.lorem_ipsum.sentence(),
                     member_since=forgery_py.date.date(True))
            db.session.add(u)
            try:
                db.session.commit()
            except IntegrityError:
                db.session.rollback()

The problem is, while I can do something like User.objects.create(...), I don't know how to rollback the database in case an IntegrityError happens (presumable due to a duplicate primary key).


By default, Django makes SQL queries in auto-commit mode, where every query is wrapped in a transaction, so you should not worry about an integrity error messing up your database. If an integrity error happens, Django won't insert the data.

To be extra-safe, you can wrap the creation code in a transaction, that in Django, is implemented as a python context handler:

from django.db import transaction

with transaction.atomic():
    User.objects.create(...)

But that would be unnecessary! You can read more about Transactions and how Django handles them in this documentation page.