Allow RedisSessionStore to connect to sentinel (#1307)

* Allow RedisSessionStore to connect to sentinel

* Reuse flask_limiter redis storage

Co-authored-by: Adrià Casajús <adria.casajus@proton.ch>
This commit is contained in:
Carlos Quintana 2022-09-23 10:23:07 +02:00 committed by GitHub
parent 3900742d1f
commit fa50c23a43
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23

View file

@ -2,7 +2,7 @@ import uuid
from typing import Optional from typing import Optional
import flask import flask
import redis import limits.storage
from flask import current_app, session from flask import current_app, session
from flask_login import logout_user from flask_login import logout_user
@ -29,8 +29,9 @@ class ServerSession(CallbackDict, SessionMixin):
class RedisSessionStore(SessionInterface): class RedisSessionStore(SessionInterface):
def __init__(self, redis, app): def __init__(self, redis_w, redis_r, app):
self._redis = redis self._redis_w = redis_w
self._redis_r = redis_r
self._app = app self._app = app
@classmethod @classmethod
@ -59,7 +60,7 @@ class RedisSessionStore(SessionInterface):
def purge_session(self, session: ServerSession): def purge_session(self, session: ServerSession):
try: try:
self._redis.delete(self._get_key(session.session_id)) self._redis_w.delete(self._get_key(session.session_id))
session.session_id = str(uuid.uuid4()) session.session_id = str(uuid.uuid4())
except AttributeError: except AttributeError:
pass pass
@ -69,7 +70,7 @@ class RedisSessionStore(SessionInterface):
if not session_id: if not session_id:
return ServerSession(session_id=str(uuid.uuid4())) return ServerSession(session_id=str(uuid.uuid4()))
val = self._redis.get(self._get_key(session_id)) val = self._redis_r.get(self._get_key(session_id))
if val is not None: if val is not None:
try: try:
data = pickle.loads(val) data = pickle.loads(val)
@ -87,7 +88,7 @@ class RedisSessionStore(SessionInterface):
secure = self.get_cookie_secure(app) secure = self.get_cookie_secure(app)
expires = self.get_expiration_time(app, session) expires = self.get_expiration_time(app, session)
val = pickle.dumps(dict(session)) val = pickle.dumps(dict(session))
self._redis.setex( self._redis_w.setex(
name=self._get_key(session.session_id), name=self._get_key(session.session_id),
value=val, value=val,
time=int(app.permanent_session_lifetime.total_seconds()), time=int(app.permanent_session_lifetime.total_seconds()),
@ -107,7 +108,18 @@ class RedisSessionStore(SessionInterface):
def set_redis_session(app: flask.Flask, redis_url: str): def set_redis_session(app: flask.Flask, redis_url: str):
app.session_interface = RedisSessionStore(redis.from_url(redis_url), app) if redis_url.startswith("redis://"):
storage = limits.storage.RedisStorage(redis_url)
app.session_interface = RedisSessionStore(storage.storage, storage.storage, app)
elif redis_url.startswith("redis+sentinel://"):
storage = limits.storage.RedisSentinelStorage(redis_url)
app.session_interface = RedisSessionStore(
storage.storage, storage.storage_slave, app
)
else:
raise RuntimeError(
f"Tried to set_redis_session with an invalid redis url: ${redis_url}"
)
def logout_session(): def logout_session():