This is an automated email from the ASF dual-hosted git repository.
vincbeck 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 72e27157974 Add declarative configuration for FAB custom roles (#72820)
72e27157974 is described below
commit 72e27157974a1531ad4b3cba09c56389987ad27f
Author: ucaeon <[email protected]>
AuthorDate: Sat Sep 12 00:23:50 2026 +0900
Add declarative configuration for FAB custom roles (#72820)
Allow deployments to provision custom roles from configuration while
preserving permissions maintained through the UI and CLI. Role creation must
not leave partial permissions or alter a role created concurrently.
---
providers/fab/docs/auth-manager/access-control.rst | 58 +++
providers/fab/provider.yaml | 13 +
.../fab/auth_manager/security_manager/override.py | 101 ++++-
.../src/airflow/providers/fab/get_provider_info.py | 7 +
.../security_manager/test_custom_roles.py | 439 +++++++++++++++++++++
.../tests/unit/fab/auth_manager/test_security.py | 54 +++
6 files changed, 671 insertions(+), 1 deletion(-)
diff --git a/providers/fab/docs/auth-manager/access-control.rst
b/providers/fab/docs/auth-manager/access-control.rst
index 761becef418..1e0a6051b3f 100644
--- a/providers/fab/docs/auth-manager/access-control.rst
+++ b/providers/fab/docs/auth-manager/access-control.rst
@@ -82,6 +82,64 @@ other users. ``Admin`` users have ``Op`` permission plus
additional permissions:
Custom Roles
'''''''''''''
+Declarative custom roles
+^^^^^^^^^^^^^^^^^^^^^^^^
+
+Use ``[fab] custom_roles`` to create custom roles and their permissions when
the API server starts:
+
+.. code-block:: ini
+
+ [fab]
+ custom_roles = {"PythonTester": [{"action": "can_read", "resource":
"DAGs"}], "Analyst": []}
+
+The equivalent environment variable is:
+
+.. code-block:: bash
+
+ export AIRFLOW__FAB__CUSTOM_ROLES='{"PythonTester": [{"action":
"can_read", "resource": "DAGs"}], "Analyst": []}'
+
+The default is ``{}``, which creates no roles. Each role maps to a list of
permission objects
+containing exactly ``action`` and ``resource``. Use actual FAB names such as
``can_edit`` and
+``DAGs``, not Python constant names such as ``ACTION_CAN_EDIT`` or a
``permission`` key.
+Names must be non-empty strings: role names support up to 64 characters,
actions up to 100,
+and resources up to 250. Repeated permission pairs are applied once.
+
+Action and resource names must already be registered in the FAB database when
this
+configuration is applied. They are validated independently: a new permission
pairing an
+existing action with an existing resource is allowed. This configuration does
not create
+actions or resources. Unknown names raise a configuration error before any
configured
+role is created.
+
+Default-role permissions are synchronized before this validation. Permissions
supplied by
+plugins or individual Dags must have their action and resource names
registered beforehand;
+names registered later during startup are not available to this initialization
step.
+
+All entries are validated before any configured role is created, including
entries for existing
+and built-in roles. Invalid JSON or an invalid structure raises a
configuration error.
+After validation, built-in roles (``Admin``, ``Viewer``, ``User``, ``Op``, and
``Public``)
+are skipped with a warning.
+
+Only missing roles are created. Like ``airflow roles import``, an existing
role name is skipped.
+Changing the configuration does not update existing permissions, so UI and CLI
edits are preserved.
+Removing a role from the configuration does not delete it from the database.
Deleting a configured
+role from the database allows it to be created again at the next
initialization.
+
+Each new role and its declared permissions are committed together. If creation
fails, that role's
+transaction is rolled back; roles successfully created earlier are retained.
If another process
+creates the same role first, its role is left unchanged by this configuration.
+
+An empty list, such as ``"Analyst": []``, declares no permissions. Normal FAB
initialization still
+adds ``can_read`` on ``Website`` to custom roles. The existing default-role
and permission
+maintenance behavior is unchanged.
+
+Startup initialization follows ``[fab] update_fab_perms``. When it is
disabled, configuration is
+not applied on startup. Running ``airflow sync-perm`` explicitly applies it
regardless of that flag.
+This configuration does not create users or assign roles to them.
+
+Registered names do not guarantee that every action-resource combination
grants access to
+an existing feature. For access to individual Dags, prefer the Dag's
``DAG(access_control=...)``
+configuration rather than managing those permissions here.
+
Dag Level Role
^^^^^^^^^^^^^^
``Admin`` can create a set of roles which are only allowed to view a certain
set of Dags. This is called Dag level access. Each Dag defined in the Dag model
table
diff --git a/providers/fab/provider.yaml b/providers/fab/provider.yaml
index 31586edde8e..4ee49559ea5 100644
--- a/providers/fab/provider.yaml
+++ b/providers/fab/provider.yaml
@@ -173,6 +173,19 @@ config:
type: string
example: ~
default: "True"
+ custom_roles:
+ description: |
+ JSON object mapping custom role names to lists of objects with
``action`` and
+ ``resource`` keys, using FAB permission names such as ``can_read``
and ``DAGs``.
+ Missing roles are created with their permissions during role
initialization.
+ Existing roles are skipped, preserving changes made through the UI
or CLI.
+ Built-in roles are ignored. Startup initialization requires
``update_fab_perms``.
+ An empty list declares no permissions; normal FAB initialization
still grants
+ custom roles ``can_read`` on ``Website``.
+ version_added: ~
+ type: string
+ example: '{"PythonTester": [{"action": "can_read", "resource":
"DAGs"}], "Analyst": []}'
+ default: "{{}}"
auth_backends:
description: |
Comma separated list of auth backends to authenticate users of the
API.
diff --git
a/providers/fab/src/airflow/providers/fab/auth_manager/security_manager/override.py
b/providers/fab/src/airflow/providers/fab/auth_manager/security_manager/override.py
index 5823fa19f08..0ac69dab596 100644
---
a/providers/fab/src/airflow/providers/fab/auth_manager/security_manager/override.py
+++
b/providers/fab/src/airflow/providers/fab/auth_manager/security_manager/override.py
@@ -72,7 +72,7 @@ from sqlalchemy.exc import IntegrityError,
MultipleResultsFound
from sqlalchemy.orm import joinedload
from werkzeug.security import check_password_hash, generate_password_hash
-from airflow.providers.common.compat.sdk import conf
+from airflow.providers.common.compat.sdk import AirflowConfigException, conf
from airflow.providers.fab.auth_manager.models import (
Action,
Group,
@@ -1241,11 +1241,110 @@ class
FabAirflowSecurityManagerOverride(AirflowSecurityManagerV2):
# Sync the default roles (Admin, Viewer, User, Op, public) with
related permissions
self.bulk_sync_roles(self.ROLE_CONFIGS)
+ self.create_roles_from_config()
+
self.add_homepage_access_to_custom_roles()
# init existing roles, the rest role could be created through UI.
self.update_admin_permission()
self.clean_perms()
+ def _get_custom_roles_config(self) -> dict[str, list[tuple[str, str]]]:
+ config = conf.getjson("fab", "custom_roles", fallback={})
+ if not isinstance(config, dict):
+ raise AirflowConfigException("[fab] custom_roles must be a JSON
object")
+
+ roles: dict[str, list[tuple[str, str]]] = {}
+ for name, items in config.items():
+ if not isinstance(name, str) or not name.strip() or len(name) > 64:
+ raise AirflowConfigException("[fab] custom_roles role names
must contain 1 to 64 characters")
+ if not isinstance(items, list):
+ raise AirflowConfigException(f"[fab] custom_roles[{name!r}]
must be a list")
+ perms: set[tuple[str, str]] = set()
+ for index, item in enumerate(items):
+ location = f"[fab] custom_roles[{name!r}][{index}]"
+ if not isinstance(item, dict) or item.keys() != {"action",
"resource"}:
+ raise AirflowConfigException(f"{location} must contain
only 'action' and 'resource'")
+ for key, limit in (("action", 100), ("resource", 250)):
+ value = item[key]
+ if not isinstance(value, str) or not value.strip() or
len(value) > limit:
+ raise AirflowConfigException(
+ f"{location}.{key} must be a non-empty string of
at most {limit} characters"
+ )
+ perms.add((item["action"], item["resource"]))
+ roles[name] = sorted(perms)
+ return roles
+
+ def create_roles_from_config(self) -> None:
+ """Create missing configured roles, preserving permissions on existing
roles."""
+ roles = self._get_custom_roles_config()
+ for name, perms in roles.items():
+ for action_name, resource_name in perms:
+ self._get_configured_action_and_resource(name, action_name,
resource_name)
+ for name, perms in roles.items():
+ if name in EXISTING_ROLES:
+ log.warning("Skipping built-in role '%s' in [fab]
custom_roles", name)
+ continue
+ self._create_role_from_config(name, perms)
+
+ def _get_configured_action_and_resource(
+ self, role_name: str, action_name: str, resource_name: str
+ ) -> tuple[Action, Resource]:
+ action = self.get_action(action_name)
+ if action is None:
+ raise AirflowConfigException(
+ f"Unknown action {action_name!r} in [fab]
custom_roles[{role_name!r}]"
+ )
+ resource = self.get_resource(resource_name)
+ if resource is None:
+ raise AirflowConfigException(
+ f"Unknown resource {resource_name!r} in [fab]
custom_roles[{role_name!r}]"
+ )
+ return action, resource
+
+ def _create_role_from_config(self, name: str, perms: list[tuple[str,
str]]) -> None:
+ # FAB's public creation helpers commit individually; a configured role
must
+ # become visible only after all its declared permissions have been
attached.
+ for attempt in range(3):
+ conflict: tuple[str, str] | None = None
+ try:
+ if self.find_role(name) is not None:
+ return
+ role = self.role_model()
+ role.name = name
+ role.permissions = []
+ self.session.add(role)
+ self.session.flush()
+
+ for action_name, resource_name in perms:
+ perm = self.get_permission(action_name, resource_name)
+ if perm is None:
+ action, resource =
self._get_configured_action_and_resource(
+ name, action_name, resource_name
+ )
+ conflict = (action_name, resource_name)
+ perm = self.permission_model()
+ perm.action = action
+ perm.resource = resource
+ self.session.add(perm)
+ self.session.flush()
+ conflict = None
+ if perm not in role.permissions:
+ role.permissions.append(perm)
+
+ self.session.commit()
+ return
+ except IntegrityError:
+ self.session.rollback()
+ if self.find_role(name) is not None:
+ return
+ if conflict is None or attempt == 2:
+ raise
+ if self.get_permission(*conflict) is None:
+ raise
+ except Exception:
+ self.session.rollback()
+ raise
+
def create_perm_vm_for_all_dag(self) -> None:
"""Create perm-vm if not exist and insert into FAB security model for
all-dags."""
# create perm for global logical dag
diff --git a/providers/fab/src/airflow/providers/fab/get_provider_info.py
b/providers/fab/src/airflow/providers/fab/get_provider_info.py
index 332774bca93..ae6b9fe73c9 100644
--- a/providers/fab/src/airflow/providers/fab/get_provider_info.py
+++ b/providers/fab/src/airflow/providers/fab/get_provider_info.py
@@ -107,6 +107,13 @@ def get_provider_info():
"example": None,
"default": "True",
},
+ "custom_roles": {
+ "description": "JSON object mapping custom role names
to lists of objects with ``action`` and\n``resource`` keys, using FAB
permission names such as ``can_read`` and ``DAGs``.\nMissing roles are created
with their permissions during role initialization.\nExisting roles are skipped,
preserving changes made through the UI or CLI.\nBuilt-in roles are ignored.
Startup initialization requires ``update_fab_perms``.\nAn empty list declares
no permissions; normal FAB init [...]
+ "version_added": None,
+ "type": "string",
+ "example": '{"PythonTester": [{"action": "can_read",
"resource": "DAGs"}], "Analyst": []}',
+ "default": "{{}}",
+ },
"auth_backends": {
"description": "Comma separated list of auth backends
to authenticate users of the API.\n",
"version_added": "2.0.0",
diff --git
a/providers/fab/tests/unit/fab/auth_manager/security_manager/test_custom_roles.py
b/providers/fab/tests/unit/fab/auth_manager/security_manager/test_custom_roles.py
new file mode 100644
index 00000000000..c74c507bd70
--- /dev/null
+++
b/providers/fab/tests/unit/fab/auth_manager/security_manager/test_custom_roles.py
@@ -0,0 +1,439 @@
+# 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.
+from __future__ import annotations
+
+import json
+from concurrent.futures import ThreadPoolExecutor
+from threading import Barrier
+from types import SimpleNamespace
+from unittest import mock
+from uuid import uuid4
+
+import pytest
+from sqlalchemy import create_engine, event, select
+from sqlalchemy.exc import IntegrityError
+from sqlalchemy.orm import Session
+
+from airflow.providers.common.compat.sdk import AirflowConfigException, conf
+from airflow.providers.fab.auth_manager.models import (
+ Action,
+ Permission,
+ Resource,
+ Role,
+ assoc_permission_role,
+)
+from airflow.providers.fab.auth_manager.security_manager.override import
FabAirflowSecurityManagerOverride
+
+from tests_common.test_utils.config import conf_vars
+
+
+def make_security_manager(session):
+ manager = object.__new__(FabAirflowSecurityManagerOverride)
+ manager.appbuilder = SimpleNamespace(session=session)
+ return manager
+
+
[email protected]
+def role_engine(tmp_path):
+ configured_engine = create_engine(conf.get("database", "sql_alchemy_conn"))
+ database = f"custom_roles_{uuid4().hex}"
+ if configured_engine.dialect.name == "sqlite":
+ engine = create_engine(f"sqlite:///{tmp_path / 'roles.db'}")
+ else:
+ with
configured_engine.connect().execution_options(isolation_level="AUTOCOMMIT") as
connection:
+ connection.exec_driver_sql(f"CREATE DATABASE {database}")
+ engine = create_engine(configured_engine.url.set(database=database))
+ try:
+ Role.metadata.create_all(
+ engine,
+ tables=[
+ Action.__table__,
+ Resource.__table__,
+ Permission.__table__,
+ Role.__table__,
+ assoc_permission_role,
+ ],
+ )
+ yield engine
+ finally:
+ engine.dispose()
+ if configured_engine.dialect.name != "sqlite":
+ with
configured_engine.connect().execution_options(isolation_level="AUTOCOMMIT") as
connection:
+ connection.exec_driver_sql(f"DROP DATABASE {database}")
+ configured_engine.dispose()
+
+
[email protected]
+def manager(role_engine):
+ with Session(role_engine) as session:
+ yield make_security_manager(session)
+
+
+def get_role_permissions(role):
+ return {(permission.action.name, permission.resource.name) for permission
in role.permissions}
+
+
+class TestCustomRoleValidation:
+ @pytest.mark.parametrize(
+ "config",
+ [
+ [],
+ None,
+ False,
+ 1,
+ "roles",
+ {"": []},
+ {" ": []},
+ {"x" * 65: []},
+ {"Analyst": {}},
+ {"Admin": "oops"},
+ {"Analyst": [None]},
+ {"Analyst": [{"permission": "can_read", "resource": "DAGs"}]},
+ {"Analyst": [{"action": "can_read"}]},
+ {"Analyst": [{"action": "can_read", "resource": "DAGs", "extra":
True}]},
+ {"Analyst": [{"action": None, "resource": "DAGs"}]},
+ {"Analyst": [{"action": " ", "resource": "DAGs"}]},
+ {"Analyst": [{"action": "x" * 101, "resource": "DAGs"}]},
+ {"Analyst": [{"action": "can_read", "resource": 42}]},
+ {"Analyst": [{"action": "can_read", "resource": ""}]},
+ {"Analyst": [{"action": "can_read", "resource": "x" * 251}]},
+ ],
+ )
+ def test_invalid_config_does_not_start_creation(self, config):
+ manager = make_security_manager(mock.Mock(spec=Session))
+ with conf_vars({("fab", "custom_roles"): json.dumps(config)}):
+ with pytest.raises(AirflowConfigException, match="custom_roles"):
+ manager.create_roles_from_config()
+ assert manager.session.mock_calls == []
+
+ @conf_vars({("fab", "custom_roles"): "{"})
+ def test_invalid_json_does_not_start_creation(self):
+ manager = make_security_manager(mock.Mock(spec=Session))
+ with pytest.raises(AirflowConfigException, match="custom_roles"):
+ manager.create_roles_from_config()
+ assert manager.session.mock_calls == []
+
+ @conf_vars({("fab", "custom_roles"): '{"Valid": [], "Invalid": "oops"}'})
+ def test_validates_all_roles_before_creation(self):
+ manager = make_security_manager(mock.Mock(spec=Session))
+ with pytest.raises(AirflowConfigException, match="Invalid"):
+ manager.create_roles_from_config()
+ assert manager.session.mock_calls == []
+
+ @conf_vars({("fab", "custom_roles"): "{}"})
+ def test_empty_config_does_not_access_database(self):
+ manager = make_security_manager(mock.Mock(spec=Session))
+ manager.create_roles_from_config()
+ assert manager.session.mock_calls == []
+
+
@mock.patch("airflow.providers.fab.auth_manager.security_manager.override.log",
autospec=True)
+ @conf_vars({("fab", "custom_roles"): '{"Admin": [], "Viewer": [], "User":
[], "Op": [], "Public": []}'})
+ def test_builtin_roles_are_skipped(self, mock_log):
+ manager = make_security_manager(mock.Mock(spec=Session))
+ manager.create_roles_from_config()
+ assert manager.session.mock_calls == []
+ assert mock_log.warning.call_count == 5
+
+
[email protected]_test
+class TestCustomRolePersistence:
+ @pytest.mark.parametrize("role_name", ["Analyst", "Admin", "Existing"])
+ @pytest.mark.parametrize("unknown", ["action", "resource"])
+ def test_rejects_unknown_names_before_creating_any_role(self, manager,
role_name, unknown):
+ manager.create_permission("can_read", "DAGs")
+ manager.add_role("Existing")
+ item = {"action": "can_read", "resource": "DAGs"}
+ item[unknown] = "typo"
+ config = {"First": [], role_name: [item]}
+ with conf_vars({("fab", "custom_roles"): json.dumps(config)}):
+ with pytest.raises(AirflowConfigException, match=f"Unknown
{unknown} 'typo'"):
+ manager.create_roles_from_config()
+ assert {role.name for role in manager.get_all_roles()} == {"Existing"}
+ assert manager.get_action("typo") is None
+ assert manager.get_resource("typo") is None
+ assert len(manager.session.scalars(select(Permission)).all()) == 1
+
+ @mock.patch.object(FabAirflowSecurityManagerOverride, "get_resource",
autospec=True)
+ @mock.patch.object(FabAirflowSecurityManagerOverride, "get_permission",
autospec=True)
+ @conf_vars(
+ {
+ (
+ "fab",
+ "custom_roles",
+ ): '{"Analyst": [{"action": "can_read", "resource": "DAGs"},
{"action": "can_read", "resource": "dags"}]}'
+ }
+ )
+ def test_deduplicates_names_resolving_to_same_permission(
+ self, mock_get_permission, mock_get_resource, manager
+ ):
+ permission = Permission(action=Action(name="can_read"),
resource=Resource(name="DAGs"))
+ manager.session.add(permission)
+ manager.session.commit()
+ mock_get_permission.return_value = permission
+ mock_get_resource.return_value = permission.resource
+ manager.create_roles_from_config()
+ assert get_role_permissions(manager.find_role("Analyst")) ==
{("can_read", "DAGs")}
+ assert
len(manager.session.execute(select(assoc_permission_role)).all()) == 1
+
+ @pytest.mark.parametrize("preexisting_permission", [False, True])
+ @conf_vars(
+ {
+ (
+ "fab",
+ "custom_roles",
+ ): '{"Analyst": [{"action": "can_read", "resource": "DAGs"},
{"action": "can_read", "resource": "DAGs"}], "Empty": []}'
+ }
+ )
+ def test_creates_roles_and_reuses_permissions(self, manager,
preexisting_permission):
+ manager.session.add_all([Action(name="can_read"),
Resource(name="DAGs")])
+ manager.session.commit()
+ if preexisting_permission:
+ manager.create_permission("can_read", "DAGs")
+ else:
+ assert manager.get_permission("can_read", "DAGs") is None
+ manager.create_roles_from_config()
+ manager.create_roles_from_config()
+ assert get_role_permissions(manager.find_role("Analyst")) ==
{("can_read", "DAGs")}
+ assert manager.find_role("Empty").permissions == []
+ assert len(manager.session.scalars(select(Permission)).all()) == 1
+ assert manager.session.scalars(select(Action.name)).all() ==
["can_read"]
+ assert manager.session.scalars(select(Resource.name)).all() == ["DAGs"]
+ assert len(manager.get_all_roles()) == 2
+
+ @conf_vars({("fab", "custom_roles"): '{"Analyst": [{"action": "can_read",
"resource": "DAGs"}]}'})
+ def test_preserves_manual_permission_changes(self, manager):
+ manager.create_permission("can_read", "DAGs")
+ manager.create_roles_from_config()
+ role = manager.find_role("Analyst")
+ replacement = manager.create_permission("can_edit", "Connections")
+ role.permissions = [replacement]
+ manager.session.commit()
+ manager.create_roles_from_config()
+ assert get_role_permissions(manager.find_role("Analyst")) ==
{("can_edit", "Connections")}
+
+ @conf_vars({("fab", "custom_roles"): '{"Analyst": "oops"}'})
+ def test_validates_existing_role(self, manager):
+ manager.add_role("Analyst")
+ with pytest.raises(AirflowConfigException, match="Analyst"):
+ manager.create_roles_from_config()
+ assert manager.find_role("Analyst").permissions == []
+
+ @pytest.mark.parametrize("failure_stage", ["permission", "commit"])
+ @conf_vars(
+ {
+ (
+ "fab",
+ "custom_roles",
+ ): '{"Completed": [], "Analyst": [{"action": "custom_action",
"resource": "Custom resource"}]}'
+ }
+ )
+ def test_rolls_back_failed_role_and_new_permissions(self, manager,
role_engine, failure_stage):
+ manager.session.add_all([Action(name="custom_action"),
Resource(name="Custom resource")])
+ manager.session.commit()
+
+ def fail_permission(session, flush_context, instances):
+ if any(isinstance(item, Permission) for item in session.new):
+ raise RuntimeError("Permission storage failed")
+
+ def fail_commit(session):
+ if manager.find_role("Analyst") is not None:
+ raise RuntimeError("Permission storage failed")
+
+ event_name, listener = (
+ ("before_flush", fail_permission)
+ if failure_stage == "permission"
+ else ("before_commit", fail_commit)
+ )
+ event.listen(manager.session, event_name, listener)
+ try:
+ with pytest.raises(RuntimeError, match="Permission storage
failed"):
+ manager.create_roles_from_config()
+ finally:
+ event.remove(manager.session, event_name, listener)
+ with Session(role_engine) as observer:
+ assert observer.scalars(select(Role.name)).all() == ["Completed"]
+ assert observer.scalars(select(Action.name)).all() ==
["custom_action"]
+ assert observer.scalars(select(Resource.name)).all() == ["Custom
resource"]
+ assert observer.scalars(select(Permission)).all() == []
+ assert observer.execute(select(assoc_permission_role)).all() == []
+ manager.create_roles_from_config()
+ assert get_role_permissions(manager.find_role("Analyst")) ==
{("custom_action", "Custom resource")}
+
+ @conf_vars({("fab", "custom_roles"): '{"Analyst": [{"action": "can_read",
"resource": "DAGs"}]}'})
+ def test_rollback_preserves_preexisting_permission(self, manager,
role_engine):
+ manager.create_permission("can_read", "DAGs")
+
+ def fail_commit(session):
+ raise RuntimeError("Connection storage failed")
+
+ event.listen(manager.session, "before_commit", fail_commit)
+ try:
+ with pytest.raises(RuntimeError, match="Connection storage
failed"):
+ manager.create_roles_from_config()
+ finally:
+ event.remove(manager.session, "before_commit", fail_commit)
+ with Session(role_engine) as observer:
+ assert observer.scalars(select(Role)).all() == []
+ assert len(observer.scalars(select(Permission)).all()) == 1
+
+ @pytest.mark.parametrize("iteration", range(5))
+ def test_concurrent_role_creation_preserves_winner_permissions(self,
role_engine, iteration):
+ with Session(role_engine) as session:
+ session.add_all([Action(name="can_read"), Action(name="can_edit"),
Resource(name="DAGs")])
+ session.commit()
+ barrier = Barrier(2, timeout=10)
+ original_find_role = FabAirflowSecurityManagerOverride.find_role
+
+ def create_role(action):
+ with Session(role_engine) as session:
+ manager = make_security_manager(session)
+ first_lookup = True
+
+ def find_role(name):
+ nonlocal first_lookup
+ role = original_find_role(manager, name)
+ if first_lookup:
+ first_lookup = False
+ assert role is None
+ barrier.wait()
+ return role
+
+ with mock.patch.object(manager, "find_role", autospec=True,
side_effect=find_role):
+ manager._create_role_from_config("Analyst", [(action,
"DAGs")])
+
+ with ThreadPoolExecutor(max_workers=2) as executor:
+ futures = [executor.submit(create_role, action) for action in
("can_read", "can_edit")]
+ for future in futures:
+ future.result(timeout=20)
+ with Session(role_engine) as observer:
+ roles = observer.scalars(select(Role)).unique().all()
+ assert len(roles) == 1
+ assert get_role_permissions(roles[0]) in ({("can_read", "DAGs")},
{("can_edit", "DAGs")})
+
+ def test_concurrent_roles_reuse_shared_permission(self, role_engine):
+ if role_engine.dialect.name == "sqlite":
+ pytest.skip("SQLite serializes writers before they can race on
shared permissions")
+ with Session(role_engine) as session:
+ session.add_all([Action(name="can_read"), Resource(name="DAGs")])
+ session.commit()
+ barrier = Barrier(2, timeout=10)
+ method = "get_permission"
+ original_lookup = getattr(FabAirflowSecurityManagerOverride, method)
+
+ def create_role(name):
+ with Session(role_engine) as session:
+ manager = make_security_manager(session)
+ first_lookup = True
+
+ def lookup(*args):
+ nonlocal first_lookup
+ result = original_lookup(manager, *args)
+ if first_lookup:
+ first_lookup = False
+ assert result is None
+ barrier.wait()
+ return result
+
+ with mock.patch.object(manager, method, autospec=True,
side_effect=lookup):
+ manager._create_role_from_config(name, [("can_read",
"DAGs")])
+
+ with ThreadPoolExecutor(max_workers=2) as executor:
+ futures = [executor.submit(create_role, name) for name in
("Analyst", "Tester")]
+ for future in futures:
+ future.result(timeout=20)
+ with Session(role_engine) as observer:
+ roles = observer.scalars(select(Role)).unique().all()
+ assert {role.name for role in roles} == {"Analyst", "Tester"}
+ assert all(get_role_permissions(role) == {("can_read", "DAGs")}
for role in roles)
+ assert len(observer.scalars(select(Permission)).all()) == 1
+
+
+class TestCustomRoleRetries:
+ def test_retries_shared_permission_conflict(self):
+ session = mock.Mock(spec=Session)
+ manager = make_security_manager(session)
+ error = IntegrityError("insert", {}, Exception("duplicate"))
+ session.flush.side_effect = [None, error, None]
+ action = Action(id=1, name="can_read")
+ resource = Resource(id=2, name="DAGs")
+ permission = Permission(id=3, action=action, resource=resource)
+ with (
+ mock.patch.object(manager, "find_role", autospec=True,
return_value=None),
+ mock.patch.object(
+ manager,
+ "get_action",
+ autospec=True,
+ return_value=action,
+ ),
+ mock.patch.object(
+ manager,
+ "get_resource",
+ autospec=True,
+ return_value=resource,
+ ),
+ mock.patch.object(
+ manager,
+ "get_permission",
+ autospec=True,
+ side_effect=[None, permission, permission],
+ ),
+ ):
+ manager._create_role_from_config("Analyst", [("can_read", "DAGs")])
+ session.rollback.assert_called_once()
+ session.commit.assert_called_once()
+ roles = [call.args[0] for call in session.add.call_args_list if
isinstance(call.args[0], Role)]
+ assert len(roles) == 2
+ assert roles[0] is not roles[1]
+ assert roles[1].permissions == [permission]
+
+ @pytest.mark.parametrize("shared_object_appeared", [False, True])
+ def test_propagates_unexplained_or_exhausted_conflicts(self,
shared_object_appeared):
+ session = mock.Mock(spec=Session)
+ manager = make_security_manager(session)
+ error = IntegrityError("insert", {}, Exception("failure"))
+ session.flush.side_effect = [None, error] * 3
+ with (
+ mock.patch.object(manager, "find_role", autospec=True,
return_value=None),
+ mock.patch.object(
+ manager,
+ "get_permission",
+ autospec=True,
+ side_effect=[None, Permission() if shared_object_appeared else
None] * 3,
+ ),
+ mock.patch.object(
+ manager,
+ "get_action",
+ autospec=True,
+ return_value=Action(name="can_read"),
+ ),
+ mock.patch.object(manager, "get_resource", autospec=True,
return_value=Resource(name="DAGs")),
+ ):
+ with pytest.raises(IntegrityError) as raised:
+ manager._create_role_from_config("Analyst", [("can_read",
"DAGs")])
+ assert raised.value is error
+ assert session.rollback.call_count == (3 if shared_object_appeared
else 1)
+ session.commit.assert_not_called()
+
+ def test_does_not_retry_unrelated_role_insert_error(self):
+ session = mock.Mock(spec=Session)
+ session.flush.side_effect = IntegrityError("insert", {},
Exception("failure"))
+ manager = make_security_manager(session)
+ with mock.patch.object(manager, "find_role", autospec=True,
return_value=None):
+ with pytest.raises(IntegrityError):
+ manager._create_role_from_config("Analyst", [])
+ session.rollback.assert_called_once()
+ session.commit.assert_not_called()
diff --git a/providers/fab/tests/unit/fab/auth_manager/test_security.py
b/providers/fab/tests/unit/fab/auth_manager/test_security.py
index f2520fbd884..18d855ce61a 100644
--- a/providers/fab/tests/unit/fab/auth_manager/test_security.py
+++ b/providers/fab/tests/unit/fab/auth_manager/test_security.py
@@ -19,6 +19,7 @@ from __future__ import annotations
import contextlib
import datetime
+import json
import logging
from typing import TYPE_CHECKING
from unittest import mock
@@ -264,6 +265,59 @@ def session(app_builder):
return app_builder.session
+class TestDeclarativeCustomRoles:
+ @pytest.fixture(autouse=True)
+ def cleanup_configured_role(self, security_manager):
+ yield
+ security_manager.session.rollback()
+ if security_manager.find_role("ConfigAnalyst") is not None:
+ security_manager.delete_role("ConfigAnalyst")
+
+ @pytest.mark.parametrize("update_permissions", [True, False])
+ @mock.patch.object(FabAuthManager, "security_manager",
new_callable=mock.PropertyMock)
+ def test_startup_creates_role_only_when_permission_updates_enabled(
+ self, mock_security_manager, security_manager, update_permissions
+ ):
+ mock_security_manager.return_value = security_manager
+ with conf_vars(
+ {
+ ("fab", "update_fab_perms"): str(update_permissions),
+ ("fab", "custom_roles"): json.dumps(
+ {"ConfigAnalyst": [{"action": "can_read", "resource":
"DAGs"}]}
+ ),
+ }
+ ):
+ FabAuthManager()._sync_appbuilder_roles()
+ role = security_manager.find_role("ConfigAnalyst")
+ if update_permissions:
+ assert {(perm.action.name, perm.resource.name) for perm in
role.permissions} == {
+ ("can_read", "DAGs"),
+ ("can_read", "Website"),
+ }
+ else:
+ assert role is None
+
+ @conf_vars({("fab", "update_fab_perms"): "False", ("fab", "custom_roles"):
'{"ConfigAnalyst": []}'})
+ def test_explicit_sync_creates_role_and_adds_homepage_permission(self,
security_manager):
+ security_manager.sync_roles()
+ role = security_manager.find_role("ConfigAnalyst")
+ assert {(perm.action.name, perm.resource.name) for perm in
role.permissions} == {
+ ("can_read", "Website")
+ }
+
+ @conf_vars({("fab", "custom_roles"): '{"ConfigAnalyst": [{"action":
"can_read", "resource": "DAGs"}]}'})
+ def
test_sync_preserves_existing_role_without_adding_configured_permissions(self,
security_manager):
+ role = security_manager.add_role("ConfigAnalyst")
+ permission = security_manager.create_permission("can_edit", "DAGs")
+ security_manager.add_permission_to_role(role, permission)
+ security_manager.sync_roles()
+ role = security_manager.find_role("ConfigAnalyst")
+ assert {(perm.action.name, perm.resource.name) for perm in
role.permissions} == {
+ ("can_edit", "DAGs"),
+ ("can_read", "Website"),
+ }
+
+
@pytest.fixture
def role(request, app, security_manager):
params = request.param
