Search code examples
python-3.xdjangodjango-modelsunique-constraint

Django unique together constraint in two directions


So I have some models that look like this

class Person(BaseModel):
    name = models.CharField(max_length=50)
    # other fields declared here...
    friends = models.ManyToManyField(
        to="self",
        through="Friendship",
        related_name="friends_to",
        symmetrical=True
    )

class Friendship(BaseModel):
    friend_from = models.ForeignKey(
        Person, on_delete=models.CASCADE, related_name="friendships_from")
    friend_to = models.ForeignKey(
        Person, on_delete=models.CASCADE, related_name="friendships_to")
    state = models.CharField(
        max_length=20, choices=FriendshipState.choices, default=FriendshipState.pending)

So basically, I'm trying to model a Facebook-like friends situation, where there are different persons and one can ask any other person to be friends. That relationship is expressed through this last model Friendship.

So far so good. But there are three situations I'd like to avoid:

    1. A Friendship can't have the same person in the friend_from and friend_to fields
    1. Only one Friendship for a set of two Friends should be allowed.

The closest I've got to that, is adding this under the Friendship model:

class Meta:
    constraints = [
        constraints.UniqueConstraint(
            fields=['friend_from', 'friend_to'], name="unique_friendship_reverse"
        ),
        models.CheckConstraint(
            name="prevent_self_follow",
            check=~models.Q(friend_from=models.F("friend_to")),
        )
    ]

This totally solves situation 1, avoiding someone to befriend with himself, using the CheckConstraint. And partially solves situation number 2, because it avoids having two Friendships like this:

p1 = Person.objects.create(name="Foo")
p2 = Person.objects.create(name="Bar")
Friendship.objects.create(friend_from=p1, friend_to=p2) # This one gets created OK
Friendship.objects.create(friend_from=p1, friend_to=p2) # This one fails and raises an IntegrityError, which is perfect 

Now there's one case that'd like to avoid that still can happen:

 Friendship.objects.create(friend_from=p1, friend_to=p2) # This one gets created OK
 Friendship.objects.create(friend_from=p2, friend_to=p1) # This one won't fail, but I'd want to

How would I make the UniqueConstraint work in this "two directions"? Or how could I add another constraint to cover this case?

Of course, I could overwrite the save method for the model or enforce this in some other way, but I'm curious about how this should be done at the database level.


Solution

  • So, just to sum up the discussion in the comments and to provide an example for anyone else looking into the same problem:

    • Contrary to my belief, this can not be achieved as of today, with Django 3.2.3, as pointed out by @Abdul Aziz Barkat. The condition kwarg that UniqueConstraint supports today isn't enough to make this work, because it just makes the constraint conditional, but it can't extend it to other cases.

    • The way of doing this in the future probably will be with UniqueConstraint's support for expressions, as commented by @Abdul Aziz Barkat too.

    • Finally, one way of solving this with a custom save method in the model could be:

    Having this situation as posted in the question:

    class Person(BaseModel):
        name = models.CharField(max_length=50)
        # other fields declared here...
        friends = models.ManyToManyField(
            to="self",
            through="Friendship",
            related_name="friends_to",
            symmetrical=True
        )
    
    class Friendship(BaseModel):
        friend_from = models.ForeignKey(
            Person, on_delete=models.CASCADE, related_name="friendships_from")
        friend_to = models.ForeignKey(
            Person, on_delete=models.CASCADE, related_name="friendships_to")
        state = models.CharField(
            max_length=20, choices=FriendshipState.choices, default=FriendshipState.pending)
    
        class Meta:
            constraints = [
                constraints.UniqueConstraint(
                    fields=['friend_from', 'friend_to'], name="unique_friendship_reverse"
                 ),
                models.CheckConstraint(
                    name="prevent_self_follow",
                    check=~models.Q(friend_from=models.F("friend_to")),
                )
            ]
    
    

    Add this to the Friendship class (that is, the "through" table of the M2M relationship):

        def save(self, *args, symmetric=True, **kwargs):
            if not self.pk:
                if symmetric:
                    f = Friendship(friend_from=self.friend_to,
                                   friend_to=self.friend_from, state=self.state)
                    f.save(symmetric=False)
            else:
                if symmetric:
                    f = Friendship.objects.get(
                        friend_from=self.friend_to, friend_to=self.friend_from)
                    f.state = self.state
                    f.save(symmetric=False)
            return super().save(*args, **kwargs)
    

    A couple of notes on that last snippet:

    • I'm not sure that using the save method from the Model class is the best way to achieve this because there are some cases where save isn't even called, notably when using bulk_create.

    • Notice that I'm first checking for self.pk. This is to identify when we are creating a record as opposed to updating a record.

    • If we are updating, then we will have to perform the same changes in the inverse relationship than in this one, to keep them in sync.

    • If we are creating, notice how we are not doing Friendship.objects.create() because that would trigger a RecursionError - maximum recursion depth exceeded. That's because, when creating the inverse relationship, it will also try to create its inverse relationship, and that one also will try, and so on. To solve this, we added the kwarg symmetric to the save method. So when we are calling it manually to create the inverse relationship, it doesn't trigger any more creations. That's why we have to first create a Friendship object and then separately call the save method passing symmetric=False.