You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by as...@apache.org on 2022/04/08 21:24:39 UTC

[airflow] branch main updated: Speed up `has_access` decorator by ~200ms (#22858)

This is an automated email from the ASF dual-hosted git repository.

ash 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 d7d2b49f29 Speed up `has_access` decorator by ~200ms (#22858)
d7d2b49f29 is described below

commit d7d2b49f29af007ff7f386a405cd1ba882a1818e
Author: Ash Berlin-Taylor <as...@apache.org>
AuthorDate: Fri Apr 8 22:24:26 2022 +0100

    Speed up `has_access` decorator by ~200ms (#22858)
    
    
    Using the ORM to create all the Role, Permission, Action and Resource
    objects, only to throw them all away _on every request_ is slow. And
    since we are now using the API more and more in the UI it's starting to
    get noticeable
    
    This changes the `user.perm` property to issue a custom query that
    returns the tuple of action_name, permission_name we want, bypassing the
    ORM object inflation entirely, and since `user.roles` isn't needed in
    most requests we no longer eagerly load that.
    
    * Fix tests
    
    Caching issues that only crop up in tests (but not ever a problem in the
    request life cycle of webserver
---
 .codespellignorelines                   |  1 +
 airflow/www/auth.py                     |  8 ++++----
 airflow/www/fab_security/sqla/models.py | 30 +++++++++++++++++++++++-------
 tests/test_utils/api_connexion_utils.py |  6 ++++--
 tests/www/test_security.py              | 13 ++++++++-----
 tests/www/views/test_views_base.py      |  2 +-
 6 files changed, 41 insertions(+), 19 deletions(-)

diff --git a/.codespellignorelines b/.codespellignorelines
index 7b8a9bf5e2..d641f0aaaa 100644
--- a/.codespellignorelines
+++ b/.codespellignorelines
@@ -1,2 +1,3 @@
             f"DELETE {source_table} FROM { ', '.join(_from_name(tbl) for tbl in stmt.froms) }"
         for frm in source_query.selectable.froms:
+    roles = relationship("Role", secondary=assoc_user_role, backref="user", lazy="selectin")
diff --git a/airflow/www/auth.py b/airflow/www/auth.py
index 4c9b33b510..cbb9582526 100644
--- a/airflow/www/auth.py
+++ b/airflow/www/auth.py
@@ -35,7 +35,10 @@ def has_access(permissions: Optional[Sequence[Tuple[str, str]]] = None) -> Calla
             __tracebackhide__ = True  # Hide from pytest traceback.
 
             appbuilder = current_app.appbuilder
-            if not g.user.is_anonymous and not g.user.perms:
+
+            if appbuilder.sm.check_authorization(permissions, request.args.get('dag_id', None)):
+                return func(*args, **kwargs)
+            elif not g.user.is_anonymous and not g.user.perms:
                 return (
                     render_template(
                         'airflow/no_roles_permissions.html',
@@ -46,9 +49,6 @@ def has_access(permissions: Optional[Sequence[Tuple[str, str]]] = None) -> Calla
                     ),
                     403,
                 )
-
-            if appbuilder.sm.check_authorization(permissions, request.args.get('dag_id', None)):
-                return func(*args, **kwargs)
             else:
                 access_denied = "Access is Denied"
                 flash(access_denied, "danger")
diff --git a/airflow/www/fab_security/sqla/models.py b/airflow/www/fab_security/sqla/models.py
index d91f54fef0..c6eb65ead1 100644
--- a/airflow/www/fab_security/sqla/models.py
+++ b/airflow/www/fab_security/sqla/models.py
@@ -20,9 +20,9 @@ import datetime
 # This product contains a modified portion of 'Flask App Builder' developed by Daniel Vaz Gaspar.
 # (https://github.com/dpgaspar/Flask-AppBuilder).
 # Copyright 2013, Daniel Vaz Gaspar
-from typing import TYPE_CHECKING, Union
+from typing import TYPE_CHECKING, Set, Tuple, Union
 
-from flask import g
+from flask import current_app, g
 from flask_appbuilder.models.sqla import Model
 from sqlalchemy import (
     Boolean,
@@ -181,7 +181,7 @@ class User(Model):
     last_login = Column(DateTime)
     login_count = Column(Integer)
     fail_login_count = Column(Integer)
-    roles = relationship("Role", secondary=assoc_user_role, backref="user", lazy="joined")
+    roles = relationship("Role", secondary=assoc_user_role, backref="user", lazy="selectin")
     created_on = Column(DateTime, default=datetime.datetime.now, nullable=True)
     changed_on = Column(DateTime, default=datetime.datetime.now, nullable=True)
 
@@ -229,10 +229,24 @@ class User(Model):
 
     @property
     def perms(self):
-        perms = set()
-        for role in self.roles:
-            perms.update((perm.action.name, perm.resource.name) for perm in role.permissions)
-        return perms
+        if not self._perms:
+            # Using the ORM here is _slow_ (Creating lots of objects to then throw them away) since this is in
+            # the path for every request. Avoid it if we can!
+            if current_app:
+                sm = current_app.appbuilder.sm
+                self._perms: Set[Tuple[str, str]] = set(
+                    sm.get_session.query(sm.action_model.name, sm.resource_model.name)
+                    .join(sm.permission_model.action)
+                    .join(sm.permission_model.resource)
+                    .join(sm.permission_model.role)
+                    .filter(sm.role_model.user.contains(self))
+                    .all()
+                )
+            else:
+                self._perms = {
+                    (perm.action.name, perm.resource.name) for role in self.roles for perm in role.permissions
+                }
+        return self._perms
 
     def get_id(self):
         return self.id
@@ -243,6 +257,8 @@ class User(Model):
     def __repr__(self):
         return self.get_full_name()
 
+    _perms = None
+
 
 class RegisterUser(Model):
     """Represents a user registration."""
diff --git a/tests/test_utils/api_connexion_utils.py b/tests/test_utils/api_connexion_utils.py
index f4553c634d..88a6988534 100644
--- a/tests/test_utils/api_connexion_utils.py
+++ b/tests/test_utils/api_connexion_utils.py
@@ -56,6 +56,8 @@ def create_user(app, username, role_name=None, email=None, permissions=None):
     if role_name:
         delete_role(app, role_name)
         role = create_role(app, role_name, permissions)
+    else:
+        role = []
 
     return appbuilder.sm.add_user(
         username=username,
@@ -80,12 +82,12 @@ def create_role(app, name, permissions=None):
     return role
 
 
-def set_user_single_role(app, username, role_name):
+def set_user_single_role(app, user, role_name):
     role = create_role(app, role_name)
-    user = app.appbuilder.sm.find_user(username)
     if role not in user.roles:
         user.roles = [role]
         app.appbuilder.sm.update_user(user)
+        user._perms = None
 
 
 def delete_role(app, name):
diff --git a/tests/www/test_security.py b/tests/www/test_security.py
index 6bb663d140..e1da8c0d54 100644
--- a/tests/www/test_security.py
+++ b/tests/www/test_security.py
@@ -425,9 +425,10 @@ def test_get_current_user_permissions(app):
         ) as user:
             assert user.perms == {(action, resource)}
 
-            user._perms = None
-            user.roles = []
-
+        with create_user_scope(
+            app,
+            username='no_perms',
+        ) as user:
             assert len(user.perms) == 0
 
 
@@ -617,7 +618,7 @@ def test_access_control_is_set_on_init(
             )
 
             security_manager.bulk_sync_roles([{'role': negated_role, 'perms': []}])
-            set_user_single_role(app, username, role_name=negated_role)
+            set_user_single_role(app, user, role_name=negated_role)
             assert_user_does_not_have_dag_perms(
                 perms=[permissions.ACTION_CAN_EDIT, permissions.ACTION_CAN_READ],
                 dag_id='access_control_test',
@@ -640,7 +641,7 @@ def test_access_control_stale_perms_are_revoked(
             role_name=role_name,
             permissions=[],
         ) as user:
-            set_user_single_role(app, username, role_name='team-a')
+            set_user_single_role(app, user, role_name='team-a')
             security_manager._sync_dag_view_permissions(
                 'access_control_test', access_control={'team-a': READ_WRITE}
             )
@@ -649,6 +650,8 @@ def test_access_control_stale_perms_are_revoked(
             security_manager._sync_dag_view_permissions(
                 'access_control_test', access_control={'team-a': READ_ONLY}
             )
+            # Clear the cache, to make it pick up new rol perms
+            user._perms = None
             assert_user_has_dag_perms(
                 perms=[permissions.ACTION_CAN_READ], dag_id='access_control_test', user=user
             )
diff --git a/tests/www/views/test_views_base.py b/tests/www/views/test_views_base.py
index a23f80d59d..aeab454d66 100644
--- a/tests/www/views/test_views_base.py
+++ b/tests/www/views/test_views_base.py
@@ -30,7 +30,7 @@ from tests.test_utils.www import check_content_in_response, check_content_not_in
 
 
 def test_index(admin_client):
-    with assert_queries_count(12):
+    with assert_queries_count(14):
         resp = admin_client.get('/', follow_redirects=True)
     check_content_in_response('DAGs', resp)