Initial commit
This commit is contained in:
John Wang
168
api/extensions/ext_session.py
Normal file
168
api/extensions/ext_session.py
Normal file
@@ -0,0 +1,168 @@
|
||||
import redis
|
||||
from flask import request
|
||||
from flask_session import Session, SqlAlchemySessionInterface, RedisSessionInterface
|
||||
from flask_session.sessions import total_seconds
|
||||
from itsdangerous import want_bytes
|
||||
|
||||
from extensions.ext_database import db
|
||||
|
||||
sess = Session()
|
||||
|
||||
|
||||
def init_app(app):
|
||||
sqlalchemy_session_interface = CustomSqlAlchemySessionInterface(
|
||||
app,
|
||||
db,
|
||||
app.config.get('SESSION_SQLALCHEMY_TABLE', 'sessions'),
|
||||
app.config.get('SESSION_KEY_PREFIX', 'session:'),
|
||||
app.config.get('SESSION_USE_SIGNER', False),
|
||||
app.config.get('SESSION_PERMANENT', True)
|
||||
)
|
||||
|
||||
session_type = app.config.get('SESSION_TYPE')
|
||||
if session_type == 'sqlalchemy':
|
||||
app.session_interface = sqlalchemy_session_interface
|
||||
elif session_type == 'redis':
|
||||
sess_redis_client = redis.Redis()
|
||||
sess_redis_client.connection_pool = redis.ConnectionPool(**{
|
||||
'host': app.config.get('SESSION_REDIS_HOST', 'localhost'),
|
||||
'port': app.config.get('SESSION_REDIS_PORT', 6379),
|
||||
'password': app.config.get('SESSION_REDIS_PASSWORD', None),
|
||||
'db': app.config.get('SESSION_REDIS_DB', 2),
|
||||
'encoding': 'utf-8',
|
||||
'encoding_errors': 'strict',
|
||||
'decode_responses': False
|
||||
})
|
||||
|
||||
app.extensions['session_redis'] = sess_redis_client
|
||||
|
||||
app.session_interface = CustomRedisSessionInterface(
|
||||
sess_redis_client,
|
||||
app.config.get('SESSION_KEY_PREFIX', 'session:'),
|
||||
app.config.get('SESSION_USE_SIGNER', False),
|
||||
app.config.get('SESSION_PERMANENT', True)
|
||||
)
|
||||
|
||||
|
||||
class CustomSqlAlchemySessionInterface(SqlAlchemySessionInterface):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
app,
|
||||
db,
|
||||
table,
|
||||
key_prefix,
|
||||
use_signer=False,
|
||||
permanent=True,
|
||||
sequence=None,
|
||||
autodelete=False,
|
||||
):
|
||||
if db is None:
|
||||
from flask_sqlalchemy import SQLAlchemy
|
||||
|
||||
db = SQLAlchemy(app)
|
||||
self.db = db
|
||||
self.key_prefix = key_prefix
|
||||
self.use_signer = use_signer
|
||||
self.permanent = permanent
|
||||
self.autodelete = autodelete
|
||||
self.sequence = sequence
|
||||
self.has_same_site_capability = hasattr(self, "get_cookie_samesite")
|
||||
|
||||
class Session(self.db.Model):
|
||||
__tablename__ = table
|
||||
|
||||
if sequence:
|
||||
id = self.db.Column( # noqa: A003, VNE003, A001
|
||||
self.db.Integer, self.db.Sequence(sequence), primary_key=True
|
||||
)
|
||||
else:
|
||||
id = self.db.Column( # noqa: A003, VNE003, A001
|
||||
self.db.Integer, primary_key=True
|
||||
)
|
||||
|
||||
session_id = self.db.Column(self.db.String(255), unique=True)
|
||||
data = self.db.Column(self.db.LargeBinary)
|
||||
expiry = self.db.Column(self.db.DateTime)
|
||||
|
||||
def __init__(self, session_id, data, expiry):
|
||||
self.session_id = session_id
|
||||
self.data = data
|
||||
self.expiry = expiry
|
||||
|
||||
def __repr__(self):
|
||||
return f"<Session data {self.data}>"
|
||||
|
||||
self.sql_session_model = Session
|
||||
|
||||
def save_session(self, *args, **kwargs):
|
||||
if request.blueprint == 'service_api':
|
||||
return
|
||||
elif request.method == 'OPTIONS':
|
||||
return
|
||||
elif request.endpoint and request.endpoint == 'health':
|
||||
return
|
||||
return super().save_session(*args, **kwargs)
|
||||
|
||||
|
||||
class CustomRedisSessionInterface(RedisSessionInterface):
|
||||
|
||||
def save_session(self, app, session, response):
|
||||
if request.blueprint == 'service_api':
|
||||
return
|
||||
elif request.method == 'OPTIONS':
|
||||
return
|
||||
elif request.endpoint and request.endpoint == 'health':
|
||||
return
|
||||
|
||||
if not self.should_set_cookie(app, session):
|
||||
return
|
||||
domain = self.get_cookie_domain(app)
|
||||
path = self.get_cookie_path(app)
|
||||
if not session:
|
||||
if session.modified:
|
||||
self.redis.delete(self.key_prefix + session.sid)
|
||||
response.delete_cookie(
|
||||
app.config["SESSION_COOKIE_NAME"], domain=domain, path=path
|
||||
)
|
||||
return
|
||||
|
||||
# Modification case. There are upsides and downsides to
|
||||
# emitting a set-cookie header each request. The behavior
|
||||
# is controlled by the :meth:`should_set_cookie` method
|
||||
# which performs a quick check to figure out if the cookie
|
||||
# should be set or not. This is controlled by the
|
||||
# SESSION_REFRESH_EACH_REQUEST config flag as well as
|
||||
# the permanent flag on the session itself.
|
||||
# if not self.should_set_cookie(app, session):
|
||||
# return
|
||||
conditional_cookie_kwargs = {}
|
||||
httponly = self.get_cookie_httponly(app)
|
||||
secure = self.get_cookie_secure(app)
|
||||
if self.has_same_site_capability:
|
||||
conditional_cookie_kwargs["samesite"] = self.get_cookie_samesite(app)
|
||||
expires = self.get_expiration_time(app, session)
|
||||
|
||||
if session.permanent:
|
||||
value = self.serializer.dumps(dict(session))
|
||||
if value is not None:
|
||||
self.redis.setex(
|
||||
name=self.key_prefix + session.sid,
|
||||
value=value,
|
||||
time=total_seconds(app.permanent_session_lifetime),
|
||||
)
|
||||
|
||||
if self.use_signer:
|
||||
session_id = self._get_signer(app).sign(want_bytes(session.sid)).decode("utf-8")
|
||||
else:
|
||||
session_id = session.sid
|
||||
response.set_cookie(
|
||||
app.config["SESSION_COOKIE_NAME"],
|
||||
session_id,
|
||||
expires=expires,
|
||||
httponly=httponly,
|
||||
domain=domain,
|
||||
path=path,
|
||||
secure=secure,
|
||||
**conditional_cookie_kwargs,
|
||||
)
|
||||
Reference in New Issue
Block a user