Add session interface using Elasticsearch

.gitignore

@@ -52,3 +52,4 @@ coverage.xml
 # Sphinx documentation
 docs/_build/
 
+.DS_Store

flask_session/__init__.py

@@ -15,7 +15,8 @@
 
 from .sessions import NullSessionInterface, RedisSessionInterface, \
     MemcachedSessionInterface, FileSystemSessionInterface, \
-    MongoDBSessionInterface, SqlAlchemySessionInterface
+    MongoDBSessionInterface, SqlAlchemySessionInterface, \
+    ElasticsearchSessionInterface
 
 
 class Session(object):
@@ -75,6 +76,9 @@ def _get_interface(self, app):
         config.setdefault('SESSION_MONGODB', None)
         config.setdefault('SESSION_MONGODB_DB', 'flask_session')
         config.setdefault('SESSION_MONGODB_COLLECT', 'sessions')
+        config.setdefault('SESSION_ELASTICSEARCH', None)
+        config.setdefault('SESSION_ELASTICSEARCH_HOST', 'http://localhost:9200')
+        config.setdefault('SESSION_ELASTICSEARCH_INDEX', 'sessions')
         config.setdefault('SESSION_SQLALCHEMY', None)
         config.setdefault('SESSION_SQLALCHEMY_TABLE', 'sessions')
 
@@ -97,6 +101,12 @@ def _get_interface(self, app):
                 config['SESSION_MONGODB_COLLECT'],
                 config['SESSION_KEY_PREFIX'], config['SESSION_USE_SIGNER'],
                 config['SESSION_PERMANENT'])
+        elif config['SESSION_TYPE'] == 'elasticsearch':
+            session_interface = ElasticsearchSessionInterface(
+                config['SESSION_ELASTICSEARCH'], config['SESSION_ELASTICSEARCH_HOST'],
+                config['SESSION_ELASTICSEARCH_INDEX'],
+                config['SESSION_KEY_PREFIX'], config['SESSION_USE_SIGNER'],
+                config['SESSION_PERMANENT'])
         elif config['SESSION_TYPE'] == 'sqlalchemy':
             session_interface = SqlAlchemySessionInterface(
                 app, config['SESSION_SQLALCHEMY'],

flask_session/sessions.py

@@ -63,6 +63,10 @@ class MongoDBSession(ServerSideSession):
     pass
 
 
+class ElasticsearchSession(ServerSideSession):
+    pass
+
+
 class SqlAlchemySession(ServerSideSession):
     pass
 
@@ -449,6 +453,102 @@ def save_session(self, app, session, response):
                             domain=domain, path=path, secure=secure)
 
 
+class ElasticsearchSessionInterface(SessionInterface):
+    """A Session interface that uses Elasticsearch as backend.
+
+    .. versionadded:: 0.X
+
+    :param client: A ``elasticsearch.Elasticsearch`` instance.
+    :param host: The elasticsearch host url you want to use.
+    :param index: The elasticsearch index you want to use.
+    :param key_prefix: A prefix that is added to all MongoDB store keys.
+    :param use_signer: Whether to sign the session id cookie or not.
+    :param permanent: Whether to use permanent session or not.
+    """
+
+    serializer = None
+    session_class = ElasticsearchSession
+
+    def __init__(self, client, host, index, key_prefix, use_signer=False,
+                    permanent=True):
+        if client is None:
+            from elasticsearch import Elasticsearch
+            client = Elasticsearch(host)
+        self.client = client
+        self.index = index
+        try:
+            self.client.indices.create(index=self.index)
+        except:
+            pass
+        self.key_prefix = key_prefix
+        self.use_signer = use_signer
+        self.permanent = permanent
+
+    def open_session(self, app, request):
+        sid = request.cookies.get(app.session_cookie_name)
+        if not sid:
+            sid = self._generate_sid()
+            return self.session_class(sid=sid, permanent=self.permanent)
+        if self.use_signer:
+            signer = self._get_signer(app)
+            if signer is None:
+                return None
+            try:
+                sid_as_bytes = signer.unsign(sid)
+                sid = sid_as_bytes.decode()
+            except BadSignature:
+                sid = self._generate_sid()
+                return self.session_class(sid=sid, permanent=self.permanent)
+
+        store_id = self.key_prefix + sid
+        document = self.client.get(index=self.index, id=store_id, ignore=404)
+        if document["found"]:
+            expiration = document["_source"]["expiration"]
+            expiration = datetime.strptime(expiration, "%Y-%m-%dT%H:%M:%S.%f")
+            if expiration <= datetime.utcnow():
+                # Delete expired session
+                self.client.delete(index=self.index, id=store_id)
+                document = None
+        if document is not None:
+            try:
+                val = document["_source"]['val']
+                return self.session_class(val, sid=sid)
+            except:
+                return self.session_class(sid=sid, permanent=self.permanent)
+        return self.session_class(sid=sid, permanent=self.permanent)
+
+    def save_session(self, app, session, response):
+        domain = self.get_cookie_domain(app)
+        path = self.get_cookie_path(app)
+        store_id = self.key_prefix + session.sid
+        if not session:
+            if session.modified:
+                self.client.delete(index=self.index, id=store_id)
+                response.delete_cookie(app.session_cookie_name,
+                                        domain=domain, path=path)
+            return
+
+        httponly = self.get_cookie_httponly(app)
+        secure = self.get_cookie_secure(app)
+        expires = self.get_expiration_time(app, session)
+        val = dict(session)
+        self.client.index(index=self.index,
+            id=store_id,
+            body={
+                'id': store_id,
+                'val': val,
+                'expiration': expires
+            }
+        )
+        if self.use_signer:
+            session_id = self._get_signer(app).sign(want_bytes(session.sid))
+        else:
+            session_id = session.sid
+        response.set_cookie(app.session_cookie_name, session_id,
+                            expires=expires, httponly=httponly,
+                            domain=domain, path=path, secure=secure)
+
+
 class SqlAlchemySessionInterface(SessionInterface):
     """Uses the Flask-SQLAlchemy from a flask app as a session backend.
 

test_session.py

@@ -110,6 +110,28 @@ def delete():
         self.assertEqual(c.get('/get').data, b'42')
         c.post('/delete')
 
+    def test_elasticsearch_session(self):
+        app = flask.Flask(__name__)
+        app.testing = True
+        app.config['SESSION_TYPE'] = 'elasticsearch'
+        Session(app)
+        @app.route('/set', methods=['POST'])
+        def set():
+            flask.session['value'] = flask.request.form['value']
+            return 'value set'
+        @app.route('/get')
+        def get():
+            return flask.session['value']
+        @app.route('/delete', methods=['POST'])
+        def delete():
+            del flask.session['value']
+            return 'value deleted'
+
+        c = app.test_client()
+        self.assertEqual(c.post('/set', data={'value': '42'}).data, b'value set')
+        self.assertEqual(c.get('/get').data, b'42')
+        c.post('/delete')
+
     def test_flasksqlalchemy_session(self):
         app = flask.Flask(__name__)
         app.debug = True