Flask Snippets

Server-side sessions with Pickle.

Posted by Robert Bruce Park on 2015-08-04 @ 19:41 and filed in Sessions

This snippet is based on the sqlite snippet here:

http://flask.pocoo.org/snippets/86/

I found that sqlite was giving strange IOErrors some of the time, in a way I wasn't able to troubleshoot, so I thought instead of storing pickled data in sqlite, I'd just store pickled data directly to disk.

Also this is only tested in python3 but can probably be made to work in python2 with some massaging.

import os

from uuid import uuid1
from pickle import UnpicklingError, dumps, loads
from contextlib import suppress
from collections import MutableMapping
from flask.sessions import SessionInterface, SessionMixin


class PickleSession(MutableMapping, SessionMixin):
    """Server-side session implementation.

    Uses pickle to achieve a disk-backed session such that multiple
    worker processes can access the same session data.
    """

    def __init__(self, directory, sid, *args, **kwargs):
        self.path = os.path.join(directory, sid)
        self.directory = directory
        self.sid = sid
        self.read()

    def __getitem__(self, key):
        self.read()
        return self.data[key]

    def __setitem__(self, key, value):
        self.data[key] = value
        self.save()

    def __delitem__(self, key):
        del self.data[key]
        self.save()

    def __iter__(self):
        return iter(self.data)

    def __len__(self):
        return len(self.data)

    def read(self):
        """Load pickle from (ram)disk."""
        try:
            with open(self.path, 'rb') as blob:
                self.data = loads(blob.read())
        except (FileNotFoundError, ValueError, EOFError, UnpicklingError):
            self.data = {}

    def save(self):
        """Dump pickle to (ram)disk atomically."""
        new_name = '{}.new'.format(self.path)
        with open(new_name, 'wb') as blob:
            blob.write(dumps(self.data))
        os.rename(new_name, self.path)

    # Note: Newer versions of Flask no longer require 
    # CallableAttributeProxy and PersistedObjectProxy

class PickleSessionInterface(SessionInterface):
    """Basic SessionInterface which uses the PickleSession."""

    def __init__(self, directory):
        self.directory = os.path.abspath(directory)
        os.makedirs(self.directory, exist_ok=True)

    def open_session(self, app, request):
        sid = request.cookies.get(
            app.session_cookie_name) or '{}-{}'.format(uuid1(), os.getpid())
        return PickleSession(self.directory, sid)

    def save_session(self, app, session, response):
        domain = self.get_cookie_domain(app)
        if not session:
            with suppress(FileNotFoundError):
                os.unlink(session.path)
            response.delete_cookie(
                app.session_cookie_name, domain=domain)
            return
        cookie_exp = self.get_expiration_time(app, session)
        response.set_cookie(
            app.session_cookie_name, session.sid,
            expires=cookie_exp, httponly=True, domain=domain)

Can be used like so:

path = '/run/shm/app_session'
if not os.path.exists(path):
    os.mkdir(path)
    os.chmod(path, int('700', 8))
app.session_interface = PickleSessionInterface(path)

This snippet by Robert Bruce Park can be used freely for anything you like. Consider it public domain.