How do you use the migration framework to add Django REST Framework Authentication Tokens

Question

I am setting up a new API using the Django REST Framework, and I need to add Auth tokens to all the existing users. The docs say to do:

from django.contrib.auth.models import User
from rest_framework.authtoken.models import Token

for user in User.objects.all():
    Token.objects.get_or_create(user=user)

But ideally, this should happen using Django's new Migration Framework.

Is there an easy way to do this?

Answer 1

The trick here is knowing that Token uses a custom save() method to generate a unique token.key but that custom save() methods are not run inside migrations. So the first token will have a blank key and the second will fail with an IntegrityError because they key is also blank and not unique.

Instead, copy the generate_key() code into your migration like this:

# Generated with `manage.py makemigrations --empty YOUR-APP`.

import binascii
import os

from django.db import migrations

# Copied from rest_framework/authtoken/models.py.
def generate_key():
    return binascii.hexlify(os.urandom(20)).decode()

def create_tokens(apps, schema_editor):
    User = apps.get_model('auth', 'User')
    Token = apps.get_model('authtoken', 'Token')
    for user in User.objects.filter(auth_token__isnull=True):
        token, _ = Token.objects.get_or_create(user=user, key=generate_key())

class Migration(migrations.Migration):

    dependencies = [
        ('YOUR-APP', 'YOUR-PREVIOUS-MIGRATION'),
    ]

    operations = [
        migrations.RunPython(create_tokens),
    ]

You should avoid importing the rest_framework code directly into a migration or one day your migrations will fail to run because you decided to remove rest_framework or the library's interface changed. Migrations need to be frozen in time.

Answer 2

Start by creating an empty migration for the app you want it to be used with. In my case, I have an app called users where this kind of thing lives, so I ran:

manage.py makemigrations users --empty

That created a new file in my migrations directory that I was able to update with the following contents:

# -*- coding: utf-8 -*-
from __future__ import unicode_literals

from django.db import models, migrations
from rest_framework.authtoken.models import Token
from django.contrib.auth.models import User

def add_tokens(apps, schema_editor):
    print "Adding auth tokens for the API..."
    for user in User.objects.all():
        Token.objects.get_or_create(user=user)

def remove_tokens(apps, schema_editor):
    print "Deleting all auth tokens for the API..."
    Token.objects.all().delete()

class Migration(migrations.Migration):

    dependencies = [
        ('users', '0002_load_initial_data'),
    ]

    operations = [
        migrations.RunPython(add_tokens, reverse_code=remove_tokens),
    ]