Using Django Check Constraints to Prevent Self-Following

This puppy would follow you anywhere.

Another way to use Django’s CheckConstraint class to ensure your data is valid. Based on my answer to a question on the Django forum.

Imagine we have a user model that we’d like to introduce a social media “following” pattern to. Users can follow other users to receive updates on our site. We’d like to ensure that users do not follow themselves, since that would need special care in all our code.

We can block self-following in the UI layer, but there’s always the risk that we’ll accidentally enable it. For example, we might add an “import your contacts” feature that allows self-following.

As with all the other posts in this series - the best way to prevent bad data is to block it in the database.

Setting up the relationship

To add the followers relationship, we’ll be using ManyToManyField. By default, ManyToManyField creates a hidden model class to hold the relationships. Because we want to customize our model with add an extra constraint, we’ll need to use the through argument to define our own visible model class instead.

We can define a first version of this like so:

from django.db import models


class User(models.Model):
    ...
    followers = models.ManyToManyField(
        to="self",
        through="Follow",
        related_name="following",
        symmetrical=False,
    )


class Follow(models.Model):
    from_user = models.ForeignKey(User, on_delete=models.CASCADE, related_name="+")
    to_user = models.ForeignKey(User, on_delete=models.CASCADE, related_name="+")

    class Meta:
        constraints = [
            models.UniqueConstraint(
                name="%(app_label)s_%(class)s_unique_relationships",
                fields=["from_user", "to_user"],
            ),
        ]

Note:

After adding a migration and running it, we can try our models out on the shell:

In [1]: from example.core.models import User

In [2]: user1 = User.objects.create()

In [3]: user2 = User.objects.create()

In [4]: user1.following.add(user2)

In [5]: user1.following.add(user1)

In [6]: user1.following.all()
Out[6]: <QuerySet [<User: User object (2)>, <User: User object (1)>]>

Note we were still able to have user1 follow itself. Let’s fix that now.

Preventing self-follows

It’s time to set up our CheckConstraint.

First, we add the constraint in Meta.constraints:

class Follow(models.Model):
    ...

    class Meta:
        constraints = [
            ...,
            models.CheckConstraint(
                name="%(app_label)s_%(class)s_prevent_self_follow",
                check=~models.Q(from_user=models.F("to_user")),
            ),
        ]

Our check here declares that the from_user field is not equal to the to_user field. We use an F() object to represent the to_user field. We negate the condition with the easy-to-miss squiggle operator (~), which turns the == into !=. The condition is therefore from_user != to_user in Python syntax.

Second, we run makemigrations to generate a new migration:

$ ./manage.py makemigrations core
Migrations for 'core':
  example/core/migrations/0003_auto_20210225_0339.py
    - Create constraint core_follow_prevent_self_follow on model follow

We check the migration looks correct:

from django.db import migrations, models


class Migration(migrations.Migration):

    dependencies = [
        ("core", "0002_auto_20210225_0320"),
    ]

    operations = [
        migrations.AddConstraint(
            model_name="follow",
            constraint=models.CheckConstraint(
                check=models.Q(_negated=True, from_user=models.F("to_user")),
                name="core_follow_prevent_self_follow",
            ),
        ),
    ]

So far so good!

Third, we can add a test that checks our check constraint works as intended:

from django.db import IntegrityError
from django.test import TestCase

from example.core.models import Follow, User


class FollowTests(TestCase):
    def test_no_self_follow(self):
        user = User.objects.create()
        constraint_name = "core_follow_prevent_self_follow"
        with self.assertRaisesMessage(IntegrityError, constraint_name):
            Follow.objects.create(from_user=user, to_user=user)

The test tries to self-follow and ensures that this triggers a database error which contains the name of our constraint.

Fourth, we need to consider how our migration will work with existing bad data. If we try and migrate whilst self-follow relationships exist, the migration will fail with an IntegrityError:

$ ./manage.py migrate
Operations to perform:
  Apply all migrations: core
Running migrations:
  Applying core.0003_auto_20210225_0339...Traceback (most recent call last):
  File "/.../django/db/backends/utils.py", line 84, in _execute
  ...
  File "/.../django/db/backends/sqlite3/base.py", line 413, in execute
    return Database.Cursor.execute(self, query, params)
django.db.utils.IntegrityError: CHECK constraint failed: core_follow_prevent_self_follow

We can delete any self-follow relationships in the migration with a RunPython operation. This allows us to use the ORM to modify the database. We can add a module-level function for RunPython and then reference it in the operation:

from django.db import migrations, models


def forwards_func(apps, schema_editor):
    Follow = apps.get_model("core", "Follow")
    db_alias = schema_editor.connection.alias
    Follow.objects.using(db_alias).filter(from_user=models.F("to_user")).delete()


class Migration(migrations.Migration):

    dependencies = [
        ("core", "0002_auto_20210225_0320"),
    ]

    operations = [
        migrations.RunPython(
            code=forwards_func,
            reverse_code=migrations.RunPython.noop,
            elidable=True,
        ),
        migrations.AddConstraint(
            model_name="follow",
            constraint=models.CheckConstraint(
                check=models.Q(_negated=True, from_user=models.F("to_user")),
                name="core_follow_prevent_self_follow",
            ),
        ),
    ]

Note:

Now we can run the migration without problem:

$ ./manage.py migrate
Operations to perform:
  Apply all migrations: core
Running migrations:
  Applying core.0003_auto_20210225_0339... OK

Great!

Fifth, we should step through our UI and API and check the ability to self-follow is not exposed anywhere. If we have an API, we’d want to make sure it returns a sensible error message when attempting to self-follow, since otherwise it would just crash with an IntegrityError.

Fin

You can’t follow yourself, but you can follow me on Twitter,

—Adam

Learn how to make your tests run quickly in my book Speed Up Your Django Tests.

Subscribe via RSS, Twitter, Mastodon, or email:

One summary email a week, no spam, I pinky promise.

Related posts:

Tags: