Switch to Custom image model in a Production Wagtail project

A software project is always evolving. New requirements can arise at any stage of the project. One example of a new feature required in a production Wagtail project is to store the source of the Image. This post will list the steps involved to make the migration from Wagtail default Image model to Custom Image model. Therefore it requires some basic understanding of Wagtail. If you heard about the Wagtail for the first time. You should check it at wagtail.io.

Brief description of the feature

Wagtail has a wonderful admin which also allows handling images. Internally wagtail has wagtailimages app which contains the default Image model. But the fields provided in the default Image model may not be sufficient for your project. Like we want to store the source of every image and display on our site. Source field is not present in Image model. Fortunately, Wagtail allows switching to custom Image model in which new fields can be added based on the project requirements. In our case, we want to add the source field.

Implementation

Like Django, Wagtail also has a great documentation. It also has documentation on the Custom image models. After following the docs, our project will be switched to Custom Image model with new fields like source. But when we open Images section in Wagtail admin, it will be empty. The reason is mentioned in docs:

When changing an existing site to use a custom image model, no images will be copied to the new model automatically. Copying old images to the new model would need to be done manually with a data migration.

So, how much progress we have made?

Switched to custom image model - Yes
Migrated old data from default image model to Custom image model - No

What are data migrations?

There are basically two types of migrations in a Django project.

Schema migrations: It is the most commonly used and often referred as migrations. These help to change database schema like creating database tables, adding new fields to tables and changing fields etc.
Data migrations: As the name suggests, migrations that alter data are called Data migrations. For example, we add a new field named full_name to the User model. Schema migrations will create the new field in the database. Data migrations help to combine first_name and last_name of the existing users and store it in the new field named full_name. You can read more in the Django documentation.

Why we need Data migrations to switch to Custom Image model?

In our Wagtail project, we changed the Image model. We created the new model but the image records data is present in old default Image model. To copy images from old to the new model we need to create a data migration which copies the image records. To create data migration:

Make an empty migration file

python manage.py makemigrations --empty yourappname

Copy the following code to the generated file.

# -*- coding: utf-8 -*-
from django.db import migrations, models


def forwards_func(apps, schema_editor):
    # We get the model from the versioned app registry;
    wagtail_image_model = apps.get_model('wagtailimages', 'Image')
    # Change the pages with your project's app name and also model name if different
    extended_image_model = apps.get_model('pages', 'ExtendedImage')
    tagged_item_model = apps.get_model('taggit', 'TaggedItem')
    django_content_type = apps.get_model('contenttypes', 'contenttype')

    db_alias = schema_editor.connection.alias

    # Get images stored in default wagtail image model
    images = wagtail_image_model.objects.using(db_alias).all()
    new_images = []
    for image in images:
        new_images.append(extended_image_model(
            id=image.id,
            title=image.title,
            file=image.file,
            width=image.width,
            height=image.height,
            created_at=image.created_at,
            focal_point_x=image.focal_point_x,
            focal_point_y=image.focal_point_y,
            focal_point_width=image.focal_point_width,
            focal_point_height=image.focal_point_height,
            file_size=image.file_size,
            collection=image.collection,
            uploaded_by_user=image.uploaded_by_user,
        ))

    # Create images in new model
    extended_image_model.objects.using(db_alias).bulk_create(new_images)
    # Leave all images in previous model

    # Move tags from old image to new image model. Moving tags is
    # a little different case. The lookup table taggit_taggeditem looks like this:
    # id   object_id   content_type_id   tag_id
    # 1    1           10                 1
    # 2    1           10                 2
    # 3    1           10                 3
    # 4    1           10                 4
    # In our case, the object_id will be same for old and new image model
    # objects. So, we have to only change the content_type_id
    ct_extended_model, created = django_content_type.objects.using(db_alias).get_or_create(
        app_label='pages',
        model='extendedimage'
    )
    ct_wagtail_model = django_content_type.objects.using(db_alias).get(
        app_label='wagtailimages',
        model='image'
    )

    tagged_item_model.objects.using(db_alias).filter(
        content_type_id=ct_wagtail_model.id).update(
            content_type_id=ct_extended_model.id
    )


def reverse_func(apps, schema_editor):
    # We get the model from the versioned app registry;
    extended_image_model = apps.get_model('pages', 'ExtendedImage')
    tagged_item_model = apps.get_model('taggit', 'TaggedItem')
    django_content_type = apps.get_model('contenttypes', 'contenttype')

    db_alias = schema_editor.connection.alias

    # Move tags from new image model to old wagtail model
    ct_extended_model = django_content_type.objects.using(db_alias).get(
        app_label='pages',
        model='extendedimage'
    )
    ct_wagtail_model = django_content_type.objects.using(db_alias).get(
        app_label='wagtailimages',
        model='image'
    )

    tagged_item_model.objects.using(db_alias).filter(
        content_type_id=ct_extended_model.id).update(
            content_type_id=ct_wagtail_model.id
    )

    # Delete all images created in the new model
    extended_image_model.objects.using(db_alias).all().delete()


class Migration(migrations.Migration):

    dependencies = [
        # Change the following according to your project
        # ('pages', '0021_auto_20180122_1100'),
        ('wagtailimages', '0019_delete_filter'),
    ]

    operations = [
        migrations.RunPython(forwards_func, reverse_func),
    ]

Details about the steps in this data migration

First, this data migration gets the default Image model objects and create the new Custom Image model objects using bulk_create.
Migrating Tags: Migrating tags is little different than other fields like title. Tags in wagtail work using django-taggit. So, let's use a simple example to understand the concept. This is not the actual model from Django taggit but simplified for this use case.

class TaggedItem(models.Model):
    tag = models.SlugField()
    content_type = models.ForeignKey(ContentType, on_delete=models.CASCADE)
    object_id = models.PositiveIntegerField()
    content_object = GenericForeignKey('content_type', 'object_id')

    def __str__(self):
        return self.tag

Django includes a contenttypes app which can track all of the models installed in your Django-powered project. In this case, the Tags objects point to old Image model using content_type. So, we can filter the TaggedItem model to get the tags used in images and can then update these with the new image model content_type.

    tagged_item_model.objects.using(db_alias).filter(
        content_type_id=ct_wagtail_model.id).update(
            content_type_id=ct_extended_model.id
    )

content types have various uses like Tags but diving into it further will deviate from the topic of this blog post. Will write a new post on content types.

You may notice we are using get_or_create to get the content type of the new Image model. New instances of ContentType are automatically created whenever new models are installed. But in our case, we have two migrations. The first one create the Schema for new Image model and second one migrate data. You may not encounter the issue running locally but when both the migrations are run during deployment new content types may not be created. In that case, we can create the content type for the new model if not exist.

Issues with PostgreSQL

One issue you will encounter if you are using PostgreSQL in your project while uploading new images:

psycopg2.IntegrityError: duplicate key value violates unique constraint "pages_extendedimage_pkey"
DETAIL:  Key (id)=(1) already exists.

The current key is set to 1 so we have to change it to Max. If your project is using a small number of images say < 10 then you can try uploading again and again and error will be removed when the current key equal to Max. But it is not feasible in case of large number of images. So, in that case, you can run SQL query on your database like

SELECT pg_catalog.setval(pg_get_serial_sequence('pages_extendedimage', 'id'), (SELECT MAX(id) FROM pages_extendedimage));

I have tried to cover all the possible cases encountered by me while migrating a production site to Custom Image model. If you face any issues, feel free to comment below. Thanks for reading.

References: