This article explains how to convert an existing Django model into an MP_Node model in django-treebeard, handling migrations, data population, and field constraints.
Django-treebeard is a popular library for managing efficient tree structures in Django. It is used in Wagtail CMS, django-CMS, and many other projects.
This article focuses on converting an existing model to an MP_Node model (Materialized Path tree). The existing model is a simple structure with a ForeignKey to itself, representing a tree structure (Adjacency List Model).
The recommended approach on the project page suggests creating a new model, copying the data, deleting the old model, and renaming the new model. Alternatively, the dump_bulk and load_bulk methods can be used.
I have tried both methods, but they are not ideal for models with many relationships or large datasets.
Creating a standard Django migration, on the other hand, can be problematic because the MP_Node model has non-nullable fields with unique constraints. These fields include path, depth, and numchild.
The approach I used involves creating a standard migration, temporarily making the required fields nullable, filling in the necessary data within the same migration, and then reverting the fields back to non-nullable.
Here is an example of a model that we want to convert to an MP_Node model:
from django.db import models
class Category(models.Model):
name = models.CharField(max_length=100)
parent = models.ForeignKey(
"self", null=True, blank=True, related_name="children", on_delete=models.PROTECT
)Modify the model to inherit from MP_Node and create a migration:
from django.db import models
from treebeard.mp_tree import MP_Node
class Category(MP_Node):
name = models.CharField(max_length=100)
parent = models.ForeignKey(
"self", null=True, blank=True, related_name="children", on_delete=models.PROTECT
)During migration creation, Django will warn that it is impossible to add a non-nullable field to the model without specifying a default value. It will then prompt you to provide a one-off default. Choose "1) Provide a one-off default now" and enter 0 for depth and an empty string ("") for path.
This is the migration that Django generated. However, it is not complete yet and will fail if you try to run it as is.
from django.db import migrations, models
class Migration(migrations.Migration):
dependencies = [
('custom', '0001_initial'),
]
operations = [
migrations.AddField(
model_name='category',
name='depth',
field=models.PositiveIntegerField(default=0),
preserve_default=False,
),
migrations.AddField(
model_name='category',
name='numchild',
field=models.PositiveIntegerField(default=0),
),
migrations.AddField(
model_name='category',
name='path',
field=models.CharField(default='', max_length=255, unique=True),
preserve_default=False,
),
]We need to temporarily make the fields nullable and remove their unique constraints, fill them with data, and then revert them back to non-nullable.
The following code recursively loops through all categories and fills the required fields:
depth and numchild are straightforward—they represent the current node's depth and the number of its children.path is more complex. It is a string representing the node's path in the tree, constructed by concatenating the parent's path with the current node's index, formatted as a four-character string.from treebeard.mp_tree import MP_Node
def process_children(items, parent, depth):
for index, instance in enumerate(items, 1):
children_qs = instance.children.all()
instance.depth = depth
instance.numchild = children_qs.count()
instance.path = MP_Node._get_path(
parent.path if parent else "", instance.depth, index
)
instance.save()
process_children(children_qs, instance, depth + 1)
def forward(apps, schema_editor):
Category = apps.get_model("custom", "Category")
process_children(Category.objects.filter(parent=None), None, 1)This code should be added just below the imports in the migration file. It should run after the fields are added but before they are reverted to non-nullable.
The complete migration file would look like this:
from django.db import migrations, models
from treebeard.mp_tree import MP_Node
def process_children(items, parent, depth):
for index, instance in enumerate(items, 1):
children_qs = instance.children.all()
instance.depth = depth
instance.numchild = children_qs.count()
instance.path = MP_Node._get_path(
parent.path if parent else "", instance.depth, index
)
instance.save()
process_children(children_qs, instance, depth + 1)
def forward(apps, schema_editor):
Category = apps.get_model("custom", "Category")
process_children(Category.objects.filter(parent=None), None, 1)
class Migration(migrations.Migration):
dependencies = [
("custom", "0001_initial"),
]
operations = [
migrations.AddField(
model_name="category",
name="depth",
field=models.PositiveIntegerField(default=0),
preserve_default=False,
),
migrations.AddField(
model_name="category",
name="numchild",
field=models.PositiveIntegerField(default=0),
),
migrations.AddField(
model_name="category",
name="path",
# temporarily set to nullable and not unique
field=models.CharField(default="", max_length=255, null=True),
preserve_default=False,
),
# fill the data
migrations.RunPython(forward, migrations.RunPython.noop),
# set the path field back to not nullable and unique
migrations.AlterField(
model_name="category",
name="path",
field=models.CharField(max_length=255, unique=True),
),
]That’s all! For reference, below is the diff between the Django-generated migration and the updated working migration.
Happy coding!
from django.db import migrations, models
+from treebeard.mp_tree import MP_Node
+
+
+def process_children(items, parent, depth):
+ for index, instance in enumerate(items, 1):
+ children_qs = instance.children.all()
+ instance.depth = depth
+ instance.numchild = children_qs.count()
+ instance.path = MP_Node._get_path(
+ parent.path if parent else "", instance.depth, index
+ )
+ instance.save()
+ process_children(children_qs, instance, depth + 1)
+
+
+def forward(apps, schema_editor):
+ Category = apps.get_model("custom", "Category")
+ process_children(Category.objects.filter(parent=None), None, 1)
class Migration(migrations.Migration):
@@ -12,7 +28,7 @@ class Migration(migrations.Migration):
migrations.AddField(
model_name="category",
name="depth",
- field=models.PositiveIntegerField(default=-1),
+ field=models.PositiveIntegerField(default=0),
preserve_default=False,
),
migrations.AddField(
@@ -23,7 +39,16 @@ class Migration(migrations.Migration):
migrations.AddField(
model_name="category",
name="path",
- field=models.CharField(default="", max_length=255, unique=True),
+ # temporarily set to nullable and not unique
+ field=models.CharField(default="", max_length=255, null=True),
preserve_default=False,
),
+ # fill the data
+ migrations.RunPython(forward, migrations.RunPython.noop),
+ # set the path field back to not nullable and unique
+ migrations.AlterField(
+ model_name="category",
+ name="path",
+ field=models.CharField(max_length=255, unique=True),
+ ),
]