From b9f2dc0067c9ea9c803956ca3ccbbba0a9215a81 Mon Sep 17 00:00:00 2001 From: Kyle Smith Date: Thu, 17 Mar 2022 15:51:45 -0400 Subject: [PATCH] add session handler for mongodb: MongoDBSessionInterface --- README.md | 13 ++++- quart_session/__init__.py | 20 ++++++- quart_session/sessions.py | 108 ++++++++++++++++++++++++++++++++++---- setup.py | 3 +- 4 files changed, 130 insertions(+), 14 deletions(-) diff --git a/README.md b/README.md index 58adfdd..5e5ec77 100644 --- a/README.md +++ b/README.md @@ -101,6 +101,18 @@ app.config['SESSION_TYPE'] = 'memcached' Session(app) ``` +### MongoDB + +via `motor`. + +```python3 +app = Quart(__name__) +app.config['SESSION_TYPE'] = 'mongodb' +app.config['SESSION_MONGODB_URI'] = 'mongodb://localhost:27017/my_database' +app.config['SESSION_MONGODB_COLLECTION'] = 'sessions' +Session(app) +``` + ### JSON serializer [flask-session](https://pypi.org/project/Flask-Session/) uses `pickle` @@ -174,7 +186,6 @@ by explicitly setting `SESSION_REVERSE_PROXY` to `True`. ## Future development -- `MongoDBSessionInterface` - `FileSystemSessionInterface` - `GoogleCloudDatastoreSessionInterface` - Pytest diff --git a/quart_session/__init__.py b/quart_session/__init__.py index 8c665de..41bc5fd 100644 --- a/quart_session/__init__.py +++ b/quart_session/__init__.py @@ -10,13 +10,19 @@ :license: BSD, see LICENSE for more details. """ -__version__ = '1.0.4' +__version__ = '1.0.5-dev' import os from quart import Quart -from .sessions import RedisSessionInterface, RedisTrioSessionInterface, MemcachedSessionInterface, NullSessionInterface +from .sessions import ( + RedisSessionInterface, + RedisTrioSessionInterface, + MemcachedSessionInterface, + MongoDBSessionInterface, + NullSessionInterface +) class Session(object): @@ -133,6 +139,16 @@ class Session(object): use_signer=config['SESSION_USE_SIGNER'], permanent=config['SESSION_PERMANENT'], **config) + elif config['SESSION_TYPE'] == 'mongodb': + session_interface = MongoDBSessionInterface( + mongodb_uri=config['SESSION_MONGODB_URI'], + collection=config['SESSION_MONGODB_COLLECTION'], + client_kwargs=config.get('SESSION_MONGODB_CLIENT_KWARGS', {}), + set_callback=config.get('SESSION_MONGODB_SET_CALLBACK'), + key_prefix=config['SESSION_KEY_PREFIX'], + use_signer=config['SESSION_USE_SIGNER'], + permanent=config['SESSION_PERMANENT'], + **config) elif config['SESSION_TYPE'] == 'null': app.logger.warning(f"{backend_warning}. Currently using: null") session_interface = NullSessionInterface( diff --git a/quart_session/sessions.py b/quart_session/sessions.py index fa52faf..3dea95e 100644 --- a/quart_session/sessions.py +++ b/quart_session/sessions.py @@ -11,8 +11,9 @@ """ import time from typing import Optional -from uuid import uuid4 +from uuid import uuid4, UUID import asyncio +import functools from quart import Quart, current_app from quart.wrappers import BaseRequestWebsocket, Response @@ -60,6 +61,10 @@ class MemcachedSession(ServerSideSession): pass +class MongoDBSession(ServerSideSession): + pass + + class NullSession(ServerSideSession): pass @@ -120,14 +125,17 @@ class SessionInterface(QuartSessionInterface): options['sid'] = self._generate_sid() return self.session_class(**options) - try: - data = self.serializer.loads(val) - except: - app.logger.warning(f"Failed to deserialize session " - f"data for sid: {sid}. Generating new sid.") - app.logger.debug(f"data: {val}") - options['sid'] = self._generate_sid() - return self.session_class(**options) + if self.serializer is None: + data = val + else: + try: + data = self.serializer.loads(val) + except: + app.logger.warning(f"Failed to deserialize session " + f"data for sid: {sid}. Generating new sid.") + app.logger.debug(f"data: {val}") + options['sid'] = self._generate_sid() + return self.session_class(**options) protection = self._config['SESSION_PROTECTION'] if protection is True and addr is not None and \ @@ -169,7 +177,11 @@ class SessionInterface(QuartSessionInterface): secure = self.get_cookie_secure(app) expires = self.get_expiration_time(app, session) - val = self.serializer.dumps(dict(session)) + if self.serializer is None: + val = dict(session) + else: + val = self.serializer.dumps(dict(session)) + await self.set(key=session_key, value=val, app=app) if self.use_signer: session_id = self._get_signer(app).sign(want_bytes(session.sid)) @@ -360,6 +372,82 @@ class MemcachedSessionInterface(SessionInterface): return await self.backend.delete(key) +def _convert_key_to_uuid(func): + """ + convert the session UUID to a UUID object for mongodb + + example: + "session:b8ebbf02-cc7a-4b0b-824f-22a984c8c0b8" -> + UUID("b8ebbf02-cc7a-4b0b-824f-22a984c8c0b8") + + """ + + @functools.wraps(func) + async def wrapper(*args, **kwargs): + if 'key' in kwargs: + key = kwargs['key'] + try: + if key.startswith('session:'): + _, _uuid = tuple(key.split(':')) + kwargs['key'] = UUID(_uuid) + except Exception as e: + current_app.logger.warning( + f"session could not be converted to a uuid object: {key}" + ) + return await func(*args, **kwargs) + return wrapper + + +class MongoDBSessionInterface(SessionInterface): + # mongodb does not a serializer as many object types are properly handled by the connector + serializer = None + session_class = MongoDBSession + + def __init__(self, mongodb_uri, collection, client_kwargs={}, set_callback=None, **kwargs): + from motor.motor_asyncio import AsyncIOMotorClient + + super().__init__(**kwargs) + self.mongodb_uri = mongodb_uri + self.client_kwargs = client_kwargs + self.set_callback = set_callback + self._collection = collection + self._client = AsyncIOMotorClient(self.mongodb_uri, uuidRepresentation='standard', **self.client_kwargs) + self._database = self._client.get_database() + + async def create(self, app: Quart) -> None: + pass + + @_convert_key_to_uuid + async def get(self, key, app): + value = await self.collection.find_one({'_id': key}, {'data': True}) + if value: + return value.get('data', {}) + else: + return None + + @_convert_key_to_uuid + async def set(self, key, value, expiry=None, app=None): + doc = { + 'data': value, + } + + # allows the document to be modified prior upsert + if callable(self.set_callback): + self.set_callback(doc) + + await self.collection.update_one({ + '_id': key + }, { + '$set': doc + }, + upsert=True + ) + + @property + def collection(self): + return self._database.get_collection(self._collection) + + class NullSessionInterface(SessionInterface): """This class does absolutely nothing""" session_class = NullSession diff --git a/setup.py b/setup.py index 169f9e4..98d187b 100644 --- a/setup.py +++ b/setup.py @@ -24,7 +24,7 @@ INSTALL_REQUIRES = [ setup( name='Quart-Session', - version='1.0.4', + version='1.0.5-dev', url='https://github.com/sferdi0/quart-session', license='BSD', author='Sander', @@ -40,6 +40,7 @@ setup( tests_require=INSTALL_REQUIRES + ["asynctest", "hypothesis", "pytest", "pytest-asyncio"], extras_require={ "dotenv": ["python-dotenv"], + "mongodb": ["motor>=2.5.1"], "redis": ["aioredis>=2.0.0"] }, classifiers=[