Add memcached storage (#224)

* Add memcached storage.

* Change .travis file to add memcached service and install aiomcache.

Author: Yaser-Amiri

.travis.yml

@@ -17,6 +17,7 @@ install:
 - pip install pytest-aiohttp
 - pip install pytest-cov
 - pip install codecov
+- pip install aiomcache
 - python setup.py develop
 
 script:
@@ -37,6 +38,8 @@ env:
   - PYTHONASYNCIODEBUG=
 services:
   - redis-server
+  - memcached
+
 sudo: false
 
 deploy:

aiohttp_session/memcached_storage.py

@@ -0,0 +1,65 @@
+import asyncio
+import json
+import uuid
+
+from . import AbstractStorage, Session
+
+
+class MemcachedStorage(AbstractStorage):
+    """Memcached storage"""
+
+    def __init__(self, memcached_conn, *, cookie_name="AIOHTTP_SESSION",
+                 domain=None, max_age=None, path='/',
+                 secure=None, httponly=True,
+                 encoder=json.dumps, decoder=json.loads,
+                 key_factory=lambda: uuid.uuid4().hex):
+        super().__init__(cookie_name=cookie_name, domain=domain,
+                         max_age=max_age, path=path, secure=secure,
+                         httponly=httponly)
+        self._encoder = encoder
+        self._decoder = decoder
+        self._key_factory = key_factory
+        self.conn = memcached_conn
+
+    @asyncio.coroutine
+    def load_session(self, request):
+        cookie = self.load_cookie(request)
+        if cookie is None:
+            return Session(None, data=None, new=True, max_age=self.max_age)
+        else:
+            key = str(cookie)
+            stored_key = (self.cookie_name + '_' + key).encode('utf-8')
+            data = yield from self.conn.get(stored_key)
+            if data is None:
+                return Session(None, data=None,
+                               new=True, max_age=self.max_age)
+            data = data.decode('utf-8')
+            try:
+                data = self._decoder(data)
+            except ValueError:
+                data = None
+            return Session(key, data=data, new=False, max_age=self.max_age)
+
+    @asyncio.coroutine
+    def save_session(self, request, response, session):
+        key = session.identity
+        if key is None:
+            key = self._key_factory()
+            self.save_cookie(response, key,
+                             max_age=session.max_age)
+        else:
+            if session.empty:
+                self.save_cookie(response, '',
+                                 max_age=session.max_age)
+            else:
+                key = str(key)
+                self.save_cookie(response, key,
+                                 max_age=session.max_age)
+
+        data = self._encoder(self._get_session_data(session))
+        max_age = session.max_age
+        expire = max_age if max_age is not None else 0
+        stored_key = (self.cookie_name + '_' + key).encode('utf-8')
+        yield from self.conn.set(
+                                stored_key, data.encode('utf-8'),
+                                exptime=expire)

demo/memcached_storage.py

@@ -0,0 +1,27 @@
+import time
+from aiohttp import web
+import asyncio
+import aiomcache
+from aiohttp_session import setup, get_session
+
+from aiohttp_session.memcached_storage import MemcachedStorage
+
+
+async def handler(request):
+    session = await get_session(request)
+    last_visit = session['last_visit'] if 'last_visit' in session else None
+    session['last_visit'] = time.time()
+    text = 'Last visited: {}'.format(last_visit)
+    return web.Response(text=text)
+
+
+async def make_app():
+    app = web.Application()
+    mc = aiomcache.Client("127.0.0.1", 11211, loop=loop)
+    setup(app, MemcachedStorage(mc))
+    app.router.add_get('/', handler)
+    return app
+
+loop = asyncio.get_event_loop()
+app = loop.run_until_complete(make_app())
+web.run_app(app)

docs/conf.py

@@ -332,5 +332,6 @@
 intersphinx_mapping = {'https://docs.python.org/3': None,
                        'https://aiohttp.readthedocs.io/en/stable': None,
                        'https://aioredis.readthedocs.io/en/latest': None,
+                       'https://github.com/aio-libs/aiomcache', None,
                        'http://cryptography.io/en/latest': None,
                        'https://pynacl.readthedocs.io/en/latest': None}

docs/glossary.rst

@@ -14,6 +14,12 @@
 
       https://aioredis.readthedocs.io/
 
+   aiomcache
+
+      :term:`asyncio` compatible Memcached client library
+
+      https://github.com/aio-libs/aiomcache
+
    asyncio
 
       The library for writing single-threaded concurrent code using

requirements-dev.txt

@@ -12,3 +12,4 @@ aiohttp==2.3.1
 multidict==3.3.0
 chardet==3.0.4
 yarl==0.13.0
+aiomcache==0.5.2

setup.py

@@ -18,6 +18,7 @@ def read(f):
 install_requires = ['aiohttp>=2.3.0']
 extras_require = {
     'aioredis': ['aioredis>=0.1.4'],
+    'aiomcache': ['aiomcache>=0.5.2'],
     'pycrypto': ['cryptography'],
     'secure': ['cryptography'],
     'pynacl': ['pynacl'],

tests/test_memcached_storage.py

@@ -0,0 +1,231 @@
+import asyncio
+import json
+import uuid
+import aiomcache
+import time
+import pytest
+
+from aiohttp import web
+from aiohttp_session import Session, session_middleware, get_session
+from aiohttp_session.memcached_storage import MemcachedStorage
+
+
+@pytest.yield_fixture
+def memcached(loop):
+    conn = aiomcache.Client("127.0.0.1", 11211, loop=loop)
+    yield conn
+    conn.close()
+
+
+def create_app(loop, handler, memcached, max_age=None,
+               key_factory=lambda: uuid.uuid4().hex):
+    middleware = session_middleware(
+        MemcachedStorage(memcached, max_age=max_age, key_factory=key_factory))
+    app = web.Application(middlewares=[middleware], loop=loop)
+    app.router.add_route('GET', '/', handler)
+    return app
+
+
+@asyncio.coroutine
+def make_cookie(client, memcached, data):
+    session_data = {
+        'session': data,
+        'created': int(time.time())
+    }
+    value = json.dumps(session_data)
+    key = uuid.uuid4().hex
+    storage_key = ('AIOHTTP_SESSION_' + key).encode('utf-8')
+    yield from memcached.set(storage_key, bytes(value, 'utf-8'))
+    client.session.cookie_jar.update_cookies({'AIOHTTP_SESSION': key})
+
+
+@asyncio.coroutine
+def make_cookie_with_bad_value(client, memcached):
+    key = uuid.uuid4().hex
+    storage_key = ('AIOHTTP_SESSION_' + key).encode('utf-8')
+    yield from memcached.set(storage_key, b'')
+    client.session.cookie_jar.update_cookies({'AIOHTTP_SESSION': key})
+
+
+@asyncio.coroutine
+def load_cookie(client, memcached):
+    cookies = client.session.cookie_jar.filter_cookies(client.make_url('/'))
+    key = cookies['AIOHTTP_SESSION']
+    storage_key = ('AIOHTTP_SESSION_' + key.value).encode('utf-8')
+    encoded = yield from memcached.get(storage_key)
+    s = encoded.decode('utf-8')
+    value = json.loads(s)
+    return value
+
+
+@asyncio.coroutine
+def test_create_new_sesssion(test_client, memcached):
+
+    @asyncio.coroutine
+    def handler(request):
+        session = yield from get_session(request)
+        assert isinstance(session, Session)
+        assert session.new
+        assert not session._changed
+        assert {} == session
+        return web.Response(body=b'OK')
+
+    client = yield from test_client(create_app, handler, memcached)
+    resp = yield from client.get('/')
+    assert resp.status == 200
+
+
+@asyncio.coroutine
+def test_load_existing_sesssion(test_client, memcached):
+
+    @asyncio.coroutine
+    def handler(request):
+        session = yield from get_session(request)
+        assert isinstance(session, Session)
+        assert not session.new
+        assert not session._changed
+        assert {'a': 1, 'b': 12} == session
+        return web.Response(body=b'OK')
+
+    client = yield from test_client(create_app, handler, memcached)
+    yield from make_cookie(client, memcached, {'a': 1, 'b': 12})
+    resp = yield from client.get('/')
+    assert resp.status == 200
+
+
+@asyncio.coroutine
+def test_load_bad_sesssion(test_client, memcached):
+    @asyncio.coroutine
+    def handler(request):
+        session = yield from get_session(request)
+        assert isinstance(session, Session)
+        assert not session.new
+        assert not session._changed
+        assert {} == session
+        return web.Response(body=b'OK')
+
+    client = yield from test_client(create_app, handler, memcached)
+    yield from make_cookie_with_bad_value(client, memcached)
+    resp = yield from client.get('/')
+    assert resp.status == 200
+
+
+@asyncio.coroutine
+def test_change_sesssion(test_client, memcached):
+
+    @asyncio.coroutine
+    def handler(request):
+        session = yield from get_session(request)
+        session['c'] = 3
+        return web.Response(body=b'OK')
+
+    client = yield from test_client(create_app, handler, memcached)
+    yield from make_cookie(client, memcached, {'a': 1, 'b': 2})
+    resp = yield from client.get('/')
+    assert resp.status == 200
+
+    value = yield from load_cookie(client, memcached)
+    assert 'session' in value
+    assert 'a' in value['session']
+    assert 'b' in value['session']
+    assert 'c' in value['session']
+    assert 'created' in value
+    assert value['session']['a'] == 1
+    assert value['session']['b'] == 2
+    assert value['session']['c'] == 3
+    morsel = resp.cookies['AIOHTTP_SESSION']
+    assert morsel['httponly']
+    assert '/' == morsel['path']
+
+
+@asyncio.coroutine
+def test_clear_cookie_on_sesssion_invalidation(test_client, memcached):
+
+    @asyncio.coroutine
+    def handler(request):
+        session = yield from get_session(request)
+        session.invalidate()
+        return web.Response(body=b'OK')
+
+    client = yield from test_client(create_app, handler, memcached)
+    yield from make_cookie(client, memcached, {'a': 1, 'b': 2})
+    resp = yield from client.get('/')
+    assert resp.status == 200
+
+    value = yield from load_cookie(client, memcached)
+    assert {} == value
+    morsel = resp.cookies['AIOHTTP_SESSION']
+    assert morsel['path'] == '/'
+    assert morsel['expires'] == "Thu, 01 Jan 1970 00:00:00 GMT"
+    assert morsel['max-age'] == "0"
+
+
+@asyncio.coroutine
+def test_create_cookie_in_handler(test_client, memcached):
+
+    @asyncio.coroutine
+    def handler(request):
+        session = yield from get_session(request)
+        session['a'] = 1
+        session['b'] = 2
+        return web.Response(body=b'OK', headers={'HOST': 'example.com'})
+
+    client = yield from test_client(create_app, handler, memcached)
+    resp = yield from client.get('/')
+    assert resp.status == 200
+
+    value = yield from load_cookie(client, memcached)
+    assert 'session' in value
+    assert 'a' in value['session']
+    assert 'b' in value['session']
+    assert 'created' in value
+    assert value['session']['a'] == 1
+    assert value['session']['b'] == 2
+    morsel = resp.cookies['AIOHTTP_SESSION']
+    assert morsel['httponly']
+    assert morsel['path'] == '/'
+    storage_key = ('AIOHTTP_SESSION_' + morsel.value).encode('utf-8')
+    exists = yield from memcached.get(storage_key)
+    assert exists
+
+
+@asyncio.coroutine
+def test_create_new_sesssion_if_key_doesnt_exists_in_memcached(
+        test_client, memcached):
+
+    @asyncio.coroutine
+    def handler(request):
+        session = yield from get_session(request)
+        assert session.new
+        return web.Response(body=b'OK')
+
+    client = yield from test_client(create_app, handler, memcached)
+    client.session.cookie_jar.update_cookies(
+        {'AIOHTTP_SESSION': 'invalid_key'})
+    resp = yield from client.get('/')
+    assert resp.status == 200
+
+
+@asyncio.coroutine
+def test_create_storate_with_custom_key_factory(test_client, memcached):
+
+    @asyncio.coroutine
+    def handler(request):
+        session = yield from get_session(request)
+        session['key'] = 'value'
+        assert session.new
+        return web.Response(body=b'OK')
+
+    def key_factory():
+        return 'test-key'
+
+    client = yield from test_client(create_app, handler, memcached, 8,
+                                    key_factory)
+    resp = yield from client.get('/')
+    assert resp.status == 200
+
+    assert resp.cookies['AIOHTTP_SESSION'].value == 'test-key'
+
+    value = yield from load_cookie(client, memcached)
+    assert 'key' in value['session']
+    assert value['session']['key'] == 'value'