You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by ur...@apache.org on 2023/07/28 06:16:21 UTC
[airflow] branch main updated: Add get_user method to auth manager (#32838)
This is an automated email from the ASF dual-hosted git repository.
uranusjr pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 319045492d Add get_user method to auth manager (#32838)
319045492d is described below
commit 319045492d2559bd856a43a1fa810adf59358d7d
Author: Vincent <97...@users.noreply.github.com>
AuthorDate: Fri Jul 28 02:16:13 2023 -0400
Add get_user method to auth manager (#32838)
---
airflow/api/auth/backend/basic_auth.py | 2 +-
.../api_connexion/endpoints/dag_run_endpoint.py | 6 ++--
.../endpoints/role_and_permission_endpoint.py | 2 +-
.../endpoints/task_instance_endpoint.py | 4 +--
airflow/api_connexion/endpoints/user_endpoint.py | 2 +-
.../schemas/role_and_permission_schema.py | 2 +-
airflow/api_connexion/schemas/user_schema.py | 2 +-
airflow/auth/managers/base_auth_manager.py | 11 +++++++
airflow/auth/managers/fab/auth/anonymous_user.py | 4 ++-
airflow/auth/managers/fab/fab_auth_manager.py | 16 +++++++--
.../sqla => auth/managers/fab}/models.py | 5 +--
airflow/auth/managers/model/__init__.py | 17 ++++++++++
.../auth/anonymous_user.py => model/base_user.py} | 38 +++++++++-------------
airflow/cli/commands/role_command.py | 2 +-
.../versions/0073_2_0_0_prefix_dag_permissions.py | 2 +-
airflow/models/__init__.py | 2 --
airflow/utils/db.py | 6 ++--
airflow/www/decorators.py | 4 +--
airflow/www/fab_security/manager.py | 5 ++-
airflow/www/fab_security/sqla/manager.py | 4 +--
airflow/www/security.py | 2 +-
docs/apache-airflow/img/airflow_erd.sha256 | 2 +-
.../endpoints/test_role_and_permission_endpoint.py | 2 +-
.../api_connexion/endpoints/test_user_endpoint.py | 2 +-
tests/api_connexion/schemas/test_user_schema.py | 2 +-
tests/auth/managers/fab/test_fab_auth_manager.py | 27 ++++++++++++---
tests/cli/commands/test_role_command.py | 2 +-
tests/test_utils/db.py | 2 +-
tests/www/test_security.py | 2 +-
29 files changed, 114 insertions(+), 65 deletions(-)
diff --git a/airflow/api/auth/backend/basic_auth.py b/airflow/api/auth/backend/basic_auth.py
index 3f802fde63..dbfeaa5f80 100644
--- a/airflow/api/auth/backend/basic_auth.py
+++ b/airflow/api/auth/backend/basic_auth.py
@@ -24,8 +24,8 @@ from flask import Response, request
from flask_appbuilder.const import AUTH_LDAP
from flask_login import login_user
+from airflow.auth.managers.fab.models import User
from airflow.utils.airflow_flask_app import get_airflow_app
-from airflow.www.fab_security.sqla.models import User
CLIENT_AUTH: tuple[str, str] | Any | None = None
diff --git a/airflow/api_connexion/endpoints/dag_run_endpoint.py b/airflow/api_connexion/endpoints/dag_run_endpoint.py
index 490923c6c3..377b79aa11 100644
--- a/airflow/api_connexion/endpoints/dag_run_endpoint.py
+++ b/airflow/api_connexion/endpoints/dag_run_endpoint.py
@@ -21,7 +21,6 @@ from http import HTTPStatus
import pendulum
from connexion import NoContent
from flask import g
-from flask_login import current_user
from marshmallow import ValidationError
from sqlalchemy import delete, or_, select
from sqlalchemy.orm import Session
@@ -68,6 +67,7 @@ from airflow.utils.session import NEW_SESSION, provide_session
from airflow.utils.state import DagRunState
from airflow.utils.types import DagRunType
from airflow.www.decorators import action_logging
+from airflow.www.extensions.init_auth_manager import get_auth_manager
RESOURCE_EVENT_PREFIX = "dag_run"
@@ -357,7 +357,7 @@ def post_dag_run(*, dag_id: str, session: Session = NEW_SESSION) -> APIResponse:
)
dag_run_note = post_body.get("note")
if dag_run_note:
- current_user_id = getattr(current_user, "id", None)
+ current_user_id = get_auth_manager().get_user_id()
dag_run.note = (dag_run_note, current_user_id)
return dagrun_schema.dump(dag_run)
except ValueError as ve:
@@ -478,7 +478,7 @@ def set_dag_run_note(*, dag_id: str, dag_run_id: str, session: Session = NEW_SES
except ValidationError as err:
raise BadRequest(detail=str(err))
- current_user_id = getattr(current_user, "id", None)
+ current_user_id = get_auth_manager().get_user_id()
if dag_run.dag_run_note is None:
dag_run.note = (new_note, current_user_id)
else:
diff --git a/airflow/api_connexion/endpoints/role_and_permission_endpoint.py b/airflow/api_connexion/endpoints/role_and_permission_endpoint.py
index 609c45893f..34b0f5478a 100644
--- a/airflow/api_connexion/endpoints/role_and_permission_endpoint.py
+++ b/airflow/api_connexion/endpoints/role_and_permission_endpoint.py
@@ -34,9 +34,9 @@ from airflow.api_connexion.schemas.role_and_permission_schema import (
role_schema,
)
from airflow.api_connexion.types import APIResponse, UpdateMask
+from airflow.auth.managers.fab.models import Action, Role
from airflow.security import permissions
from airflow.utils.airflow_flask_app import get_airflow_app
-from airflow.www.fab_security.sqla.models import Action, Role
from airflow.www.security import AirflowSecurityManager
diff --git a/airflow/api_connexion/endpoints/task_instance_endpoint.py b/airflow/api_connexion/endpoints/task_instance_endpoint.py
index 55db7ef8b9..c0f1f7079e 100644
--- a/airflow/api_connexion/endpoints/task_instance_endpoint.py
+++ b/airflow/api_connexion/endpoints/task_instance_endpoint.py
@@ -51,6 +51,7 @@ from airflow.utils.airflow_flask_app import get_airflow_app
from airflow.utils.db import get_query_count
from airflow.utils.session import NEW_SESSION, provide_session
from airflow.utils.state import DagRunState, TaskInstanceState
+from airflow.www.extensions.init_auth_manager import get_auth_manager
T = TypeVar("T")
@@ -692,9 +693,8 @@ def set_task_instance_note(
raise NotFound(error_message)
ti, sla_miss = result
- from flask_login import current_user
- current_user_id = getattr(current_user, "id", None)
+ current_user_id = get_auth_manager().get_user_id()
if ti.task_instance_note is None:
ti.note = (new_note, current_user_id)
else:
diff --git a/airflow/api_connexion/endpoints/user_endpoint.py b/airflow/api_connexion/endpoints/user_endpoint.py
index 2a88fb1b24..c8ba1a7cc7 100644
--- a/airflow/api_connexion/endpoints/user_endpoint.py
+++ b/airflow/api_connexion/endpoints/user_endpoint.py
@@ -34,9 +34,9 @@ from airflow.api_connexion.schemas.user_schema import (
user_schema,
)
from airflow.api_connexion.types import APIResponse, UpdateMask
+from airflow.auth.managers.fab.models import Role, User
from airflow.security import permissions
from airflow.utils.airflow_flask_app import get_airflow_app
-from airflow.www.fab_security.sqla.models import Role, User
@security.requires_access([(permissions.ACTION_CAN_READ, permissions.RESOURCE_USER)])
diff --git a/airflow/api_connexion/schemas/role_and_permission_schema.py b/airflow/api_connexion/schemas/role_and_permission_schema.py
index 324336c288..30f23e2f35 100644
--- a/airflow/api_connexion/schemas/role_and_permission_schema.py
+++ b/airflow/api_connexion/schemas/role_and_permission_schema.py
@@ -21,7 +21,7 @@ from typing import NamedTuple
from marshmallow import Schema, fields
from marshmallow_sqlalchemy import SQLAlchemySchema, auto_field
-from airflow.www.fab_security.sqla.models import Action, Permission, Resource, Role
+from airflow.auth.managers.fab.models import Action, Permission, Resource, Role
class ActionSchema(SQLAlchemySchema):
diff --git a/airflow/api_connexion/schemas/user_schema.py b/airflow/api_connexion/schemas/user_schema.py
index 843ad32f02..fb30da9a02 100644
--- a/airflow/api_connexion/schemas/user_schema.py
+++ b/airflow/api_connexion/schemas/user_schema.py
@@ -23,7 +23,7 @@ from marshmallow_sqlalchemy import SQLAlchemySchema, auto_field
from airflow.api_connexion.parameters import validate_istimezone
from airflow.api_connexion.schemas.role_and_permission_schema import RoleSchema
-from airflow.www.fab_security.sqla.models import User
+from airflow.auth.managers.fab.models import User
class UserCollectionItemSchema(SQLAlchemySchema):
diff --git a/airflow/auth/managers/base_auth_manager.py b/airflow/auth/managers/base_auth_manager.py
index 17b7b61aeb..a234efcaca 100644
--- a/airflow/auth/managers/base_auth_manager.py
+++ b/airflow/auth/managers/base_auth_manager.py
@@ -20,6 +20,7 @@ from __future__ import annotations
from abc import abstractmethod
from typing import TYPE_CHECKING
+from airflow.auth.managers.model.base_user import BaseUser
from airflow.exceptions import AirflowException
from airflow.utils.log.logging_mixin import LoggingMixin
@@ -42,6 +43,16 @@ class BaseAuthManager(LoggingMixin):
"""Return the username associated to the user in session."""
...
+ @abstractmethod
+ def get_user(self) -> BaseUser:
+ """Return the user associated to the user in session."""
+ ...
+
+ @abstractmethod
+ def get_user_id(self) -> str:
+ """Return the user ID associated to the user in session."""
+ ...
+
@abstractmethod
def is_logged_in(self) -> bool:
"""Return whether the user is logged in."""
diff --git a/airflow/auth/managers/fab/auth/anonymous_user.py b/airflow/auth/managers/fab/auth/anonymous_user.py
index b97ec69e3e..c11669f0c6 100644
--- a/airflow/auth/managers/fab/auth/anonymous_user.py
+++ b/airflow/auth/managers/fab/auth/anonymous_user.py
@@ -20,8 +20,10 @@ from __future__ import annotations
from flask import current_app
from flask_login import AnonymousUserMixin
+from airflow.auth.managers.model.base_user import BaseUser
-class AnonymousUser(AnonymousUserMixin):
+
+class AnonymousUser(AnonymousUserMixin, BaseUser):
"""User object used when no active user is logged in."""
_roles: set[tuple[str, str]] = set()
diff --git a/airflow/auth/managers/fab/fab_auth_manager.py b/airflow/auth/managers/fab/fab_auth_manager.py
index ac5ae45c43..085b6e997b 100644
--- a/airflow/auth/managers/fab/fab_auth_manager.py
+++ b/airflow/auth/managers/fab/fab_auth_manager.py
@@ -22,6 +22,7 @@ from flask_login import current_user
from airflow import AirflowException
from airflow.auth.managers.base_auth_manager import BaseAuthManager
+from airflow.auth.managers.fab.models import User
from airflow.auth.managers.fab.security_manager_override import FabAirflowSecurityManagerOverride
@@ -39,13 +40,22 @@ class FabAuthManager(BaseAuthManager):
For backward compatibility reasons, the username in FAB auth manager is the concatenation of the
first name and the last name.
"""
- first_name = current_user.first_name or ""
- last_name = current_user.last_name or ""
+ user = self.get_user()
+ first_name = user.first_name or ""
+ last_name = user.last_name or ""
return f"{first_name} {last_name}".strip()
+ def get_user(self) -> User:
+ """Return the user associated to the user in session."""
+ return current_user
+
+ def get_user_id(self) -> str:
+ """Return the user ID associated to the user in session."""
+ return str(self.get_user().get_id())
+
def is_logged_in(self) -> bool:
"""Return whether the user is logged in."""
- return current_user and not current_user.is_anonymous
+ return not self.get_user().is_anonymous
def get_security_manager_override_class(self) -> type:
"""Return the security manager override."""
diff --git a/airflow/www/fab_security/sqla/models.py b/airflow/auth/managers/fab/models.py
similarity index 98%
rename from airflow/www/fab_security/sqla/models.py
rename to airflow/auth/managers/fab/models.py
index 5c83588801..475499707e 100644
--- a/airflow/www/fab_security/sqla/models.py
+++ b/airflow/auth/managers/fab/models.py
@@ -40,6 +40,7 @@ from sqlalchemy import (
)
from sqlalchemy.orm import backref, declared_attr, relationship
+from airflow.auth.managers.model.base_user import BaseUser
from airflow.models.base import Base
"""
@@ -139,7 +140,7 @@ assoc_user_role = Table(
)
-class User(Model):
+class User(Model, BaseUser):
"""Represents an Airflow user which has roles assigned to it."""
__tablename__ = "ab_user"
@@ -185,7 +186,7 @@ class User(Model):
@classmethod
def get_user_id(cls):
try:
- return g.user.id
+ return g.user.get_id()
except Exception:
return None
diff --git a/airflow/auth/managers/model/__init__.py b/airflow/auth/managers/model/__init__.py
new file mode 100644
index 0000000000..217e5db960
--- /dev/null
+++ b/airflow/auth/managers/model/__init__.py
@@ -0,0 +1,17 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
diff --git a/airflow/auth/managers/fab/auth/anonymous_user.py b/airflow/auth/managers/model/base_user.py
similarity index 50%
copy from airflow/auth/managers/fab/auth/anonymous_user.py
copy to airflow/auth/managers/model/base_user.py
index b97ec69e3e..329846881a 100644
--- a/airflow/auth/managers/fab/auth/anonymous_user.py
+++ b/airflow/auth/managers/model/base_user.py
@@ -17,32 +17,26 @@
# under the License.
from __future__ import annotations
-from flask import current_app
-from flask_login import AnonymousUserMixin
+from abc import abstractmethod
-class AnonymousUser(AnonymousUserMixin):
- """User object used when no active user is logged in."""
-
- _roles: set[tuple[str, str]] = set()
- _perms: set[tuple[str, str]] = set()
+class BaseUser:
+ """User model interface."""
@property
- def roles(self):
- if not self._roles:
- public_role = current_app.appbuilder.get_app.config["AUTH_ROLE_PUBLIC"]
- self._roles = {current_app.appbuilder.sm.find_role(public_role)} if public_role else set()
- return self._roles
+ def is_authenticated(self) -> bool:
+ return not self.is_anonymous
- @roles.setter
- def roles(self, roles):
- self._roles = roles
- self._perms = set()
+ @property
+ @abstractmethod
+ def is_active(self) -> bool:
+ ...
@property
- def perms(self):
- if not self._perms:
- self._perms = {
- (perm.action.name, perm.resource.name) for role in self.roles for perm in role.permissions
- }
- return self._perms
+ @abstractmethod
+ def is_anonymous(self) -> bool:
+ ...
+
+ @abstractmethod
+ def get_id(self) -> str:
+ ...
diff --git a/airflow/cli/commands/role_command.py b/airflow/cli/commands/role_command.py
index db11d69dd2..4e439fa2fc 100644
--- a/airflow/cli/commands/role_command.py
+++ b/airflow/cli/commands/role_command.py
@@ -23,11 +23,11 @@ import itertools
import json
import os
+from airflow.auth.managers.fab.models import Action, Permission, Resource, Role
from airflow.cli.simple_table import AirflowConsole
from airflow.utils import cli as cli_utils
from airflow.utils.cli import suppress_logs_and_warning
from airflow.utils.providers_configuration_loader import providers_configuration_loaded
-from airflow.www.fab_security.sqla.models import Action, Permission, Resource, Role
from airflow.www.security import EXISTING_ROLES
diff --git a/airflow/migrations/versions/0073_2_0_0_prefix_dag_permissions.py b/airflow/migrations/versions/0073_2_0_0_prefix_dag_permissions.py
index 660da4ac2c..44f930a04b 100644
--- a/airflow/migrations/versions/0073_2_0_0_prefix_dag_permissions.py
+++ b/airflow/migrations/versions/0073_2_0_0_prefix_dag_permissions.py
@@ -28,7 +28,7 @@ from flask_appbuilder import SQLA
from airflow import settings
from airflow.security import permissions
-from airflow.www.fab_security.sqla.models import Action, Permission, Resource
+from airflow.auth.managers.fab.models import Action, Permission, Resource
# revision identifiers, used by Alembic.
revision = "849da589634d"
diff --git a/airflow/models/__init__.py b/airflow/models/__init__.py
index 45690b5ae0..4d15dfe5c7 100644
--- a/airflow/models/__init__.py
+++ b/airflow/models/__init__.py
@@ -59,12 +59,10 @@ def import_all_models():
for name in __lazy_imports:
__getattr__(name)
- import airflow.jobs.job
import airflow.models.dagwarning
import airflow.models.dataset
import airflow.models.serialized_dag
import airflow.models.tasklog
- import airflow.www.fab_security.sqla.models
def __getattr__(name):
diff --git a/airflow/utils/db.py b/airflow/utils/db.py
index 9ef504e739..f0f60d1dc0 100644
--- a/airflow/utils/db.py
+++ b/airflow/utils/db.py
@@ -720,8 +720,8 @@ def _get_flask_db(sql_database_uri):
def _create_db_from_orm(session):
from alembic import command
+ from airflow.auth.managers.fab.models import Model
from airflow.models.base import Base
- from airflow.www.fab_security.sqla.models import Model
def _create_flask_session_tbl(sql_database_uri):
db = _get_flask_db(sql_database_uri)
@@ -998,7 +998,7 @@ def check_username_duplicates(session: Session) -> Iterable[str]:
:param session: session of the sqlalchemy
:rtype: str
"""
- from airflow.www.fab_security.sqla.models import RegisterUser, User
+ from airflow.auth.managers.fab.models import RegisterUser, User
for model in [User, RegisterUser]:
dups = []
@@ -1730,8 +1730,8 @@ def drop_airflow_models(connection):
:param connection: SQLAlchemy Connection
:return: None
"""
+ from airflow.auth.managers.fab.models import Model
from airflow.models.base import Base
- from airflow.www.fab_security.sqla.models import Model
Base.metadata.drop_all(connection)
Model.metadata.drop_all(connection)
diff --git a/airflow/www/decorators.py b/airflow/www/decorators.py
index af316e3ed0..c74be2635e 100644
--- a/airflow/www/decorators.py
+++ b/airflow/www/decorators.py
@@ -26,7 +26,7 @@ from itertools import chain
from typing import Callable, TypeVar, cast
import pendulum
-from flask import after_this_request, g, request
+from flask import after_this_request, request
from pendulum.parsing.exceptions import ParserError
from airflow.models import Log
@@ -88,7 +88,7 @@ def action_logging(func: Callable | None = None, event: str | None = None) -> Ca
if not get_auth_manager().is_logged_in():
user = "anonymous"
else:
- user = f"{g.user.username} ({g.user.get_full_name()})"
+ user = get_auth_manager().get_user_name()
fields_skip_logging = {"csrf_token", "_csrf_token"}
extra_fields = [
diff --git a/airflow/www/fab_security/manager.py b/airflow/www/fab_security/manager.py
index dc31a35522..fca6a3083c 100644
--- a/airflow/www/fab_security/manager.py
+++ b/airflow/www/fab_security/manager.py
@@ -67,12 +67,11 @@ from flask_appbuilder.security.views import (
from flask_jwt_extended import current_user as current_user_jwt
from flask_limiter import Limiter
from flask_limiter.util import get_remote_address
-from flask_login import current_user
from werkzeug.security import check_password_hash
+from airflow.auth.managers.fab.models import Action, Permission, RegisterUser, Resource, Role, User
from airflow.configuration import conf
from airflow.www.extensions.init_auth_manager import get_auth_manager
-from airflow.www.fab_security.sqla.models import Action, Permission, RegisterUser, Resource, Role, User
# This product contains a modified portion of 'Flask App Builder' developed by Daniel Vaz Gaspar.
# (https://github.com/dpgaspar/Flask-AppBuilder).
@@ -1539,4 +1538,4 @@ class BaseSecurityManager:
@staticmethod
def before_request():
"""Hook runs before request."""
- g.user = current_user
+ g.user = get_auth_manager().get_user()
diff --git a/airflow/www/fab_security/sqla/manager.py b/airflow/www/fab_security/sqla/manager.py
index c90ca52654..6ce9580c29 100644
--- a/airflow/www/fab_security/sqla/manager.py
+++ b/airflow/www/fab_security/sqla/manager.py
@@ -26,8 +26,7 @@ from sqlalchemy import and_, func, inspect, literal
from sqlalchemy.orm.exc import MultipleResultsFound
from werkzeug.security import generate_password_hash
-from airflow.www.fab_security.manager import BaseSecurityManager
-from airflow.www.fab_security.sqla.models import (
+from airflow.auth.managers.fab.models import (
Action,
Permission,
RegisterUser,
@@ -36,6 +35,7 @@ from airflow.www.fab_security.sqla.models import (
User,
assoc_permission_role,
)
+from airflow.www.fab_security.manager import BaseSecurityManager
log = logging.getLogger(__name__)
diff --git a/airflow/www/security.py b/airflow/www/security.py
index 7f534e78b8..98384839ab 100644
--- a/airflow/www/security.py
+++ b/airflow/www/security.py
@@ -24,6 +24,7 @@ from flask import g
from sqlalchemy import or_
from sqlalchemy.orm import Session, joinedload
+from airflow.auth.managers.fab.models import Permission, Resource, Role, User
from airflow.auth.managers.fab.views.user_details import CustomUserDBModelView
from airflow.exceptions import AirflowException, RemovedInAirflow3Warning
from airflow.models import DagBag, DagModel
@@ -32,7 +33,6 @@ from airflow.utils.log.logging_mixin import LoggingMixin
from airflow.utils.session import NEW_SESSION, provide_session
from airflow.www.extensions.init_auth_manager import get_auth_manager
from airflow.www.fab_security.sqla.manager import SecurityManager
-from airflow.www.fab_security.sqla.models import Permission, Resource, Role, User
from airflow.www.fab_security.views import (
ActionModelView,
CustomResetMyPasswordView,
diff --git a/docs/apache-airflow/img/airflow_erd.sha256 b/docs/apache-airflow/img/airflow_erd.sha256
index dc75b8aee6..f10af90f5b 100644
--- a/docs/apache-airflow/img/airflow_erd.sha256
+++ b/docs/apache-airflow/img/airflow_erd.sha256
@@ -1 +1 @@
-902b20991528af2fceec9c967f4a7a59b921e0fc2e39b6bb2c4727b3f61a56b7
\ No newline at end of file
+7ff18b1eafa528dbdfb62d75151a526de280f73e1edf7fea18da7c64644f0da9
\ No newline at end of file
diff --git a/tests/api_connexion/endpoints/test_role_and_permission_endpoint.py b/tests/api_connexion/endpoints/test_role_and_permission_endpoint.py
index 4ba36597e6..0a62de1043 100644
--- a/tests/api_connexion/endpoints/test_role_and_permission_endpoint.py
+++ b/tests/api_connexion/endpoints/test_role_and_permission_endpoint.py
@@ -19,8 +19,8 @@ from __future__ import annotations
import pytest
from airflow.api_connexion.exceptions import EXCEPTIONS_LINK_MAP
+from airflow.auth.managers.fab.models import Role
from airflow.security import permissions
-from airflow.www.fab_security.sqla.models import Role
from airflow.www.security import EXISTING_ROLES
from tests.test_utils.api_connexion_utils import (
assert_401,
diff --git a/tests/api_connexion/endpoints/test_user_endpoint.py b/tests/api_connexion/endpoints/test_user_endpoint.py
index a84b7f1087..f29f998a84 100644
--- a/tests/api_connexion/endpoints/test_user_endpoint.py
+++ b/tests/api_connexion/endpoints/test_user_endpoint.py
@@ -22,10 +22,10 @@ import pytest
from sqlalchemy.sql.functions import count
from airflow.api_connexion.exceptions import EXCEPTIONS_LINK_MAP
+from airflow.auth.managers.fab.models import User
from airflow.security import permissions
from airflow.utils import timezone
from airflow.utils.session import create_session
-from airflow.www.fab_security.sqla.models import User
from tests.test_utils.api_connexion_utils import assert_401, create_user, delete_user
from tests.test_utils.config import conf_vars
diff --git a/tests/api_connexion/schemas/test_user_schema.py b/tests/api_connexion/schemas/test_user_schema.py
index 042948429a..cf52e61b85 100644
--- a/tests/api_connexion/schemas/test_user_schema.py
+++ b/tests/api_connexion/schemas/test_user_schema.py
@@ -19,8 +19,8 @@ from __future__ import annotations
import pytest
from airflow.api_connexion.schemas.user_schema import user_collection_item_schema, user_schema
+from airflow.auth.managers.fab.models import User
from airflow.utils import timezone
-from airflow.www.fab_security.sqla.models import User
from tests.test_utils.api_connexion_utils import create_role, delete_role
TEST_EMAIL = "test@example.org"
diff --git a/tests/auth/managers/fab/test_fab_auth_manager.py b/tests/auth/managers/fab/test_fab_auth_manager.py
index 313165b56b..96ddbc7f3e 100644
--- a/tests/auth/managers/fab/test_fab_auth_manager.py
+++ b/tests/auth/managers/fab/test_fab_auth_manager.py
@@ -23,8 +23,8 @@ import pytest
from airflow import AirflowException
from airflow.auth.managers.fab.fab_auth_manager import FabAuthManager
+from airflow.auth.managers.fab.models import User
from airflow.auth.managers.fab.security_manager_override import FabAirflowSecurityManagerOverride
-from airflow.www.fab_security.sqla.models import User
from airflow.www.security import ApplessAirflowSecurityManager
@@ -44,21 +44,38 @@ class TestFabAuthManager:
(None, "Last", "Last"),
],
)
- @mock.patch("flask_login.utils._get_user")
- def test_get_user_name(self, mock_current_user, first_name, last_name, expected, auth_manager):
+ @mock.patch.object(FabAuthManager, "get_user")
+ def test_get_user_name(self, mock_get_user, first_name, last_name, expected, auth_manager):
user = User()
user.first_name = first_name
user.last_name = last_name
- mock_current_user.return_value = user
+ mock_get_user.return_value = user
assert auth_manager.get_user_name() == expected
@mock.patch("flask_login.utils._get_user")
- def test_is_logged_in(self, mock_current_user, auth_manager):
+ def test_get_user(self, mock_current_user, auth_manager):
user = Mock()
user.is_anonymous.return_value = True
mock_current_user.return_value = user
+ assert auth_manager.get_user() == user
+
+ @mock.patch.object(FabAuthManager, "get_user")
+ def test_get_user_id(self, mock_get_user, auth_manager):
+ user_id = "test"
+ user = Mock()
+ user.get_id.return_value = user_id
+ mock_get_user.return_value = user
+
+ assert auth_manager.get_user_id() == user_id
+
+ @mock.patch.object(FabAuthManager, "get_user")
+ def test_is_logged_in(self, mock_get_user, auth_manager):
+ user = Mock()
+ user.is_anonymous.return_value = True
+ mock_get_user.return_value = user
+
assert auth_manager.is_logged_in() is False
def test_get_security_manager_override_class_return_fab_security_manager_override(self, auth_manager):
diff --git a/tests/cli/commands/test_role_command.py b/tests/cli/commands/test_role_command.py
index 8d8772dcf9..544d8e9560 100644
--- a/tests/cli/commands/test_role_command.py
+++ b/tests/cli/commands/test_role_command.py
@@ -23,10 +23,10 @@ from contextlib import redirect_stdout
import pytest
+from airflow.auth.managers.fab.models import Role
from airflow.cli.commands import role_command
from airflow.security import permissions
from airflow.utils.cli_app_builder import get_application_builder
-from airflow.www.fab_security.sqla.models import Role
TEST_USER1_EMAIL = "test-user1@example.com"
TEST_USER2_EMAIL = "test-user2@example.com"
diff --git a/tests/test_utils/db.py b/tests/test_utils/db.py
index 79074a1ddb..43b6185eea 100644
--- a/tests/test_utils/db.py
+++ b/tests/test_utils/db.py
@@ -17,6 +17,7 @@
# under the License.
from __future__ import annotations
+from airflow.auth.managers.fab.models import Permission, Resource, assoc_permission_role
from airflow.jobs.job import Job
from airflow.models import (
Connection,
@@ -50,7 +51,6 @@ from airflow.models.serialized_dag import SerializedDagModel
from airflow.security.permissions import RESOURCE_DAG_PREFIX
from airflow.utils.db import add_default_pool_if_not_exists, create_default_connections, reflect_tables
from airflow.utils.session import create_session
-from airflow.www.fab_security.sqla.models import Permission, Resource, assoc_permission_role
def clear_db_runs():
diff --git a/tests/www/test_security.py b/tests/www/test_security.py
index 58ccc778e1..99659de2c7 100644
--- a/tests/www/test_security.py
+++ b/tests/www/test_security.py
@@ -31,13 +31,13 @@ from sqlalchemy import Column, Date, Float, Integer, String
from airflow.auth.managers.fab.auth.anonymous_user import AnonymousUser
from airflow.auth.managers.fab.fab_auth_manager import FabAuthManager
+from airflow.auth.managers.fab.models import User, assoc_permission_role
from airflow.exceptions import AirflowException
from airflow.models import DagModel
from airflow.models.base import Base
from airflow.models.dag import DAG
from airflow.security import permissions
from airflow.www import app as application
-from airflow.www.fab_security.sqla.models import User, assoc_permission_role
from airflow.www.utils import CustomSQLAInterface
from tests.test_utils.api_connexion_utils import (
create_user,