Update reset_user_sessions to work from either CLI or web

airflow/auth/managers/fab/security_manager/override.py

@@ -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,
@@ -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:
@@ -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",
             )
 
@@ -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>", "*"))

tests/www/views/test_views_custom_user_views.py

@@ -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()
@@ -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]
+        )