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:
parent
3900742d1f
commit
fa50c23a43
|
@ -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():
|
||||||
|
|
Loading…
Reference in a new issue