Skip to content

Commit

Permalink
Update reset_user_sessions to work from either CLI or web
Browse files Browse the repository at this point in the history
This updates the `reset_user_sessions` method to not class `flash` if
it is not being run within the web context but to log the warning
message instead. Without this change, `flask fab reset-password` would
fail if a warning message needed to be printed.
  • Loading branch information
pdebelak committed Dec 4, 2023
1 parent 35a1b7a commit 1f88ad1
Show file tree
Hide file tree
Showing 2 changed files with 59 additions and 23 deletions.
46 changes: 25 additions & 21 deletions airflow/auth/managers/fab/security_manager/override.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@

import jwt
import re2
from flask import flash, g, session
from flask import flash, g, has_request_context, session
from flask_appbuilder import const
from flask_appbuilder.const import (
AUTH_DB,
Expand Down Expand Up @@ -476,17 +476,15 @@ def reset_user_sessions(self, user: User) -> None:
user_session_model = interface.sql_session_model
num_sessions = session.query(user_session_model).count()
if num_sessions > MAX_NUM_DATABASE_USER_SESSIONS:
flash(
Markup(
f"The old sessions for user {user.username} have <b>NOT</b> been deleted!<br>"
f"You have a lot ({num_sessions}) of user sessions in the 'SESSIONS' table in "
f"your database.<br> "
"This indicates that this deployment might have an automated API calls that create "
"and not reuse sessions.<br>You should consider reusing sessions or cleaning them "
"periodically using db clean.<br>"
"Make sure to reset password for the user again after cleaning the session table "
"to remove old sessions of the user."
),
_cli_safe_flash(
f"The old sessions for user {user.username} have <b>NOT</b> been deleted!<br>"
f"You have a lot ({num_sessions}) of user sessions in the 'SESSIONS' table in "
f"your database.<br> "
"This indicates that this deployment might have an automated API calls that create "
"and not reuse sessions.<br>You should consider reusing sessions or cleaning them "
"periodically using db clean.<br>"
"Make sure to reset password for the user again after cleaning the session table "
"to remove old sessions of the user.",
"warning",
)
else:
Expand All @@ -495,15 +493,13 @@ def reset_user_sessions(self, user: User) -> None:
if session_details.get("_user_id") == user.id:
session.delete(s)
else:
flash(
Markup(
"Since you are using `securecookie` session backend mechanism, we cannot prevent "
f"some old sessions for user {user.username} to be reused.<br> If you want to make sure "
"that the user is logged out from all sessions, you should consider using "
"`database` session backend mechanism.<br> You can also change the 'secret_key` "
"webserver configuration for all your webserver instances and restart the webserver. "
"This however will logout all users from all sessions."
),
_cli_safe_flash(
"Since you are using `securecookie` session backend mechanism, we cannot prevent "
f"some old sessions for user {user.username} to be reused.<br> If you want to make sure "
"that the user is logged out from all sessions, you should consider using "
"`database` session backend mechanism.<br> You can also change the 'secret_key` "
"webserver configuration for all your webserver instances and restart the webserver. "
"This however will logout all users from all sessions.",
"warning",
)

Expand Down Expand Up @@ -2666,3 +2662,11 @@ def _get_root_dag_id(self, dag_id: str) -> str:
).one()
return dm.root_dag_id or dm.dag_id
return dag_id


def _cli_safe_flash(text: str, level: str) -> None:
"""Shows a flash in a web context or prints a message if not."""
if has_request_context():
flash(Markup(text), level)
else:
getattr(log, level)(text.replace("<br>", "\n").replace("<b>", "*").replace("</b>", "*"))
36 changes: 34 additions & 2 deletions tests/www/views/test_views_custom_user_views.py
Original file line number Diff line number Diff line change
Expand Up @@ -248,8 +248,9 @@ def get_session_by_id(self, session_id: str):
return self.db.session.query(self.model).filter(self.model.session_id == session_id).scalar()

@mock.patch("airflow.auth.managers.fab.security_manager.override.flash")
@mock.patch("airflow.auth.managers.fab.security_manager.override.has_request_context", return_value=True)
@mock.patch("airflow.auth.managers.fab.security_manager.override.MAX_NUM_DATABASE_USER_SESSIONS", 1)
def test_refuse_delete(self, flash_mock):
def test_refuse_delete(self, _mock_has_context, flash_mock):
self.create_user_db_session("session_id_1", timedelta(days=1), self.user_1.id)
self.create_user_db_session("session_id_2", timedelta(days=1), self.user_2.id)
self.db.session.commit()
Expand All @@ -268,11 +269,42 @@ def test_refuse_delete(self, flash_mock):
assert self.get_session_by_id("session_id_2") is not None

@mock.patch("airflow.auth.managers.fab.security_manager.override.flash")
def test_warn_securecookie(self, flash_mock):
@mock.patch("airflow.auth.managers.fab.security_manager.override.has_request_context", return_value=True)
def test_warn_securecookie(self, _mock_has_context, flash_mock):
self.app.session_interface = SecureCookieSessionInterface()
self.security_manager.reset_password(self.user_1.id, "new_password")
assert flash_mock.called
assert (
"Since you are using `securecookie` session backend mechanism, we cannot"
in flash_mock.call_args[0][0]
)

@mock.patch("airflow.auth.managers.fab.security_manager.override.log")
@mock.patch("airflow.auth.managers.fab.security_manager.override.MAX_NUM_DATABASE_USER_SESSIONS", 1)
def test_refuse_delete_cli(self, log_mock):
self.create_user_db_session("session_id_1", timedelta(days=1), self.user_1.id)
self.create_user_db_session("session_id_2", timedelta(days=1), self.user_2.id)
self.db.session.commit()
self.db.session.flush()
assert self.db.session.query(self.model).count() == 2
assert self.get_session_by_id("session_id_1") is not None
assert self.get_session_by_id("session_id_2") is not None
self.security_manager.reset_password(self.user_1.id, "new_password")
assert log_mock.warning.called
assert (
"The old sessions for user user_to_delete_1 have *NOT* been deleted!\n"
in log_mock.warning.call_args[0][0]
)
assert self.db.session.query(self.model).count() == 2
assert self.get_session_by_id("session_id_1") is not None
assert self.get_session_by_id("session_id_2") is not None

@mock.patch("airflow.auth.managers.fab.security_manager.override.log")
def test_warn_securecookie_cli(self, log_mock):
self.app.session_interface = SecureCookieSessionInterface()
self.security_manager.reset_password(self.user_1.id, "new_password")
assert log_mock.warning.called
assert (
"Since you are using `securecookie` session backend mechanism, we cannot"
in log_mock.warning.call_args[0][0]
)

0 comments on commit 1f88ad1

Please sign in to comment.