django18_session_database_migration_pickle_to_json.py

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

from django.contrib.sessions.backends.db import SessionStore as DjangoSessionStore
from django.db import migrations
from django.contrib.sessions.serializers import PickleSerializer, JSONSerializer

logger = logging.getLogger(__name__)


class SessionStore(DjangoSessionStore):

    def __init__(self, serializer, *args, **kwargs):
        super(SessionStore, self).__init__(*args, **kwargs)
        self.serializer = serializer


def decode_pickle(session_data):
    return SessionStore(PickleSerializer).decode(session_data)


def decode_json(session_data):
    return SessionStore(JSONSerializer).decode(session_data)


def encode_pickle(session_dict):
    return SessionStore(PickleSerializer).encode(session_dict)


def encode_json(session_dict):
    return SessionStore(JSONSerializer).encode(session_dict)


def pickle_to_json(apps, _):
    session_model = apps.get_model("sessions", "Session")
    for session in session_model.objects.all():
        session_dict = decode_pickle(session.session_data)
        try:
            session.session_data = encode_json(session_dict)
            session.save()
        except TypeError as e:
            logger.info('Session removed %s, error: %s', session, e)
            session.delete()


def json_to_pickle(apps, _):
    session_model = apps.get_model("sessions", "Session")
    for session in session_model.objects.all():
        session_dict = decode_json(session.session_data)
        session.session_data = encode_pickle(session_dict)
        session.save()


class Migration(migrations.Migration):

    dependencies = [
        ('util', '0001_initial'),
    ]

    operations = [
        migrations.RunPython(pickle_to_json, reverse_code=json_to_pickle),
    ]