Tuesday, June 14, 2022

Enforce row count limits on ManyToManyField rows in Django models

I had a requirement where ModelA needs to limit how many ModelBs it can have. They are linked via a ManyToManyField using a through model.

It looks like this:

 class ModelA(models.Model):  
   name = models.Charfield(max_length=255)  
   b_limit = models.PositiveSmallIntegerField(help_text="Max number of Bs we can have")  
   model_b_set = models.ManyToManyFields('ModelB', through="ABJoinTable")  

 class ModelB(models.Model):  
   some_field = models.Charfield(max_length=255)    
   model_a_set = models.ManyToManyFields('ModelA', through="ABJoinTable")  

 class ABJoinTable(models.Model):   
   model_a = models.ForeignKey("ModelA", on_delete=models.CASCADE)  
   model_b = models.ForeignKey("ModelB", on_delete=models.CASCADE)  
   is_suspended = models.BooleanField(default=False)  

The through model just contains extra information is NOT the key to enforcing the rule. Remember, we want to limit how many ModelB rows ModelA can have. How many ModelB's ModelA can have it controlled by the "b_limit" field. 

The solution is to use the m2m_changed signal

from django.db.models.signals import m2m_changed  
from django.core.exceptions import ValidationError

def enforce_modelA_B_limit(sender, **kwargs):  
   modelA = kwargs['instance']  
   if modelA.model_b_set.count() >= modelA.b_limit:  
    raise ValidationError("ModelA has too many ModelBs")  

m2m_changed.connect(enforce_modelA_B_limit, sender=ModelA.model_b_set.through)  

You can place this wherever a signal function is valid. I just placed mine in the same place where the models are declared.

Why a signal?

If your gut feeling was to override the `save()` or `clean()` methods then those won't work. You cannot check the count for the ManyToManyField until you have save the primary model - ModelA. You'll get an exception saying something like: `ModelA needs to have a value for field "id" before this many-to-many relationship can be used.`.

Testing Tip

In writing a test case for this, you can use the `self.assertRaises()` context to catch the error. It's bad practice if you do a `try-catch` block to see if the rule works. The test case would look like this:

 def test_modelA_B_limit_rule(self):  
   # we assume that a setUpTestData() method exist  
   # we also assume that ModelA is limited to 1 ModelB so adding a second ModelB should throw an
   # exception
   sample_b = ModelB.objects.create(some_field="test")  
   modelA_instance = ModelA.objects.get(pk=1)  
   with self.assertRaises(ValidationError):  
     modelA_instance.model_b_set.add(sample_b) # should throw  

There's nothing to change if you have Django Rest Framework project. Your API will throw a 400 error if they violate the limit.

Refs: https://stackoverflow.com/questions/20203806/limit-maximum-choices-of-manytomanyfield

No comments:

Post a Comment