diff --git a/airflow/bin/cli.py b/airflow/bin/cli.py index b2aa9191bcd44..77fddbcd3ffe3 100644 --- a/airflow/bin/cli.py +++ b/airflow/bin/cli.py @@ -1633,8 +1633,14 @@ def list_dag_runs(args, dag=None): @cli_utils.action_logging def sync_perm(args): appbuilder = cached_appbuilder() - print('Update permission, view-menu for all existing roles') + print('Updating permission, view-menu for all existing roles') appbuilder.sm.sync_roles() + print('Updating permission on all DAG views') + dags = DagBag().dags.values() + for dag in dags: + appbuilder.sm.sync_perm_for_dag( + dag.dag_id, + dag.access_control) Arg = namedtuple( @@ -2293,7 +2299,7 @@ class CLIFactory(object): }, { 'func': sync_perm, - 'help': "Update existing role's permissions.", + 'help': "Update permissions for existing roles and DAGs.", 'args': tuple(), }, { diff --git a/airflow/models/__init__.py b/airflow/models/__init__.py index c44859fd00b38..12276705eb465 100755 --- a/airflow/models/__init__.py +++ b/airflow/models/__init__.py @@ -3027,6 +3027,10 @@ class DAG(BaseDag, LoggingMixin): :param on_success_callback: Much like the ``on_failure_callback`` except that it is executed when the dag succeeds. :type on_success_callback: callable + :param access_control: Specify optional DAG-level permissions, e.g., + {'role1': {'can_dag_read'}, + 'role2': {'can_dag_read', 'can_dag_edit'}} + :type access_control: dict """ def __init__( @@ -3048,7 +3052,8 @@ def __init__( orientation=configuration.conf.get('webserver', 'dag_orientation'), catchup=configuration.conf.getboolean('scheduler', 'catchup_by_default'), on_success_callback=None, on_failure_callback=None, - params=None): + params=None, + access_control=None): self.user_defined_macros = user_defined_macros self.user_defined_filters = user_defined_filters @@ -3125,6 +3130,7 @@ def __init__( self.on_failure_callback = on_failure_callback self._old_context_manager_dags = [] + self._access_control = access_control self._comps = { 'dag_id', @@ -3346,6 +3352,14 @@ def concurrency(self): def concurrency(self, value): self._concurrency = value + @property + def access_control(self): + return self._access_control + + @access_control.setter + def access_control(self, value): + self._access_control = value + @property def description(self): return self._description diff --git a/airflow/www/security.py b/airflow/www/security.py index a81f36e1dc4e9..c87691f95de42 100644 --- a/airflow/www/security.py +++ b/airflow/www/security.py @@ -24,6 +24,7 @@ from sqlalchemy import or_ from airflow import models +from airflow.exceptions import AirflowException from airflow.www.app import appbuilder from airflow.utils.db import provide_session from airflow.utils.log.logging_mixin import LoggingMixin @@ -451,12 +452,17 @@ def sync_roles(self): self.update_admin_perm_view() self.clean_perms() - def sync_perm_for_dag(self, dag_id): + def sync_perm_for_dag(self, dag_id, access_control=None): """ Sync permissions for given dag id. The dag id surely exists in our dag bag - as only / refresh button will call this function - - :param dag_id: + as only / refresh button or cli.sync_perm will call this function + + :param dag_id: the ID of the DAG whose permissions should be updated + :type dag_id: string + :param access_control: a dict where each key is a rolename and + each value is a set() of permission names (e.g., + {'can_dag_read'} + :type access_control: dict :return: """ for dag_perm in DAG_PERMS: @@ -464,6 +470,72 @@ def sync_perm_for_dag(self, dag_id): if perm_on_dag is None: self.add_permission_view_menu(dag_perm, dag_id) + if access_control: + self._sync_dag_view_permissions(dag_id, access_control) + + def _sync_dag_view_permissions(self, dag_id, access_control): + """Set the access policy on the given DAG's ViewModel. + + :param dag_id: the ID of the DAG whose permissions should be updated + :type dag_id: string + :param access_control: a dict where each key is a rolename and + each value is a set() of permission names (e.g., + {'can_dag_read'} + :type access_control: dict + """ + def _get_or_create_dag_permission(perm_name): + dag_perm = self.find_permission_view_menu(perm_name, dag_id) + if not dag_perm: + self.log.info("Creating new permission '{}' on view '{}'".format( + perm_name, + dag_id + )) + dag_perm = self.add_permission_view_menu(perm_name, dag_id) + + return dag_perm + + def _revoke_stale_permissions(dag_view): + existing_dag_perms = self.find_permissions_view_menu(dag_view) + for perm in existing_dag_perms: + non_admin_roles = [role for role in perm.role + if role.name != 'Admin'] + for role in non_admin_roles: + target_perms_for_role = access_control.get(role.name, {}) + if perm.permission.name not in target_perms_for_role: + self.log.info("Revoking '{}' on DAG '{}' for role '{}'".format( + perm.permission, + dag_id, + role.name + )) + self.del_permission_role(role, perm) + + dag_view = self.find_view_menu(dag_id) + if dag_view: + _revoke_stale_permissions(dag_view) + + for rolename, perms in access_control.items(): + role = self.find_role(rolename) + if not role: + raise AirflowException( + "The access_control mapping for DAG '{}' includes a role " + "named '{}', but that role does not exist".format( + dag_id, + rolename)) + + perms = set(perms) + invalid_perms = perms - DAG_PERMS + if invalid_perms: + raise AirflowException( + "The access_control map for DAG '{}' includes the following " + "invalid permissions: {}; The set of valid permissions " + "is: {}".format(dag_id, + (perms - DAG_PERMS), + DAG_PERMS)) + + for perm_name in perms: + dag_perm = _get_or_create_dag_permission(perm_name) + self.add_permission_role(role, dag_perm) + def create_perm_vm_for_all_dag(self): """ Create perm-vm if not exist and insert into FAB security model for all-dags. diff --git a/airflow/www/views.py b/airflow/www/views.py index 819a5f68c991b..7bc09a14bb1bc 100644 --- a/airflow/www/views.py +++ b/airflow/www/views.py @@ -1678,10 +1678,10 @@ def refresh(self, session=None): session.merge(orm_dag) session.commit() + dag = dagbag.get_dag(dag_id) # sync dag permission - appbuilder.sm.sync_perm_for_dag(dag_id) + appbuilder.sm.sync_perm_for_dag(dag_id, dag.access_control) - dagbag.get_dag(dag_id) flash("DAG [{}] is now fresh as a daisy".format(dag_id)) return redirect(request.referrer) @@ -1690,9 +1690,9 @@ def refresh(self, session=None): @action_logging def refresh_all(self): dagbag.collect_dags(only_if_updated=False) - for dag_id in dagbag.dags: - # sync permissions for all dags - appbuilder.sm.sync_perm_for_dag(dag_id) + # sync permissions for all dags + for dag_id, dag in dagbag.dags.items(): + appbuilder.sm.sync_perm_for_dag(dag_id, dag.access_control) flash("All DAGs are now up to date") return redirect('/') diff --git a/tests/core.py b/tests/core.py index ad68f7d88a615..9cb9369caeb1c 100644 --- a/tests/core.py +++ b/tests/core.py @@ -1313,13 +1313,40 @@ def test_cli_remove_user_role(self): "User should have been removed from role 'Viewer'" ) - def test_cli_sync_perm(self): - # test whether sync_perm cli will throw exceptions or not + @mock.patch("airflow.bin.cli.DagBag") + def test_cli_sync_perm(self, dagbag_mock): + self.expect_dagbag_contains([ + DAG('has_access_control', + access_control={ + 'Public': {'can_dag_read'} + }), + DAG('no_access_control') + ], dagbag_mock) + self.appbuilder.sm = mock.Mock() + args = self.parser.parse_args([ 'sync_perm' ]) cli.sync_perm(args) + self.appbuilder.sm.sync_roles.assert_called_once() + + self.assertEqual(2, + len(self.appbuilder.sm.sync_perm_for_dag.mock_calls)) + self.appbuilder.sm.sync_perm_for_dag.assert_any_call( + 'has_access_control', + {'Public': {'can_dag_read'}} + ) + self.appbuilder.sm.sync_perm_for_dag.assert_any_call( + 'no_access_control', + None, + ) + + def expect_dagbag_contains(self, dags, dagbag_mock): + dagbag = mock.Mock() + dagbag.dags = {dag.dag_id: dag for dag in dags} + dagbag_mock.return_value = dagbag + def test_cli_list_tasks(self): for dag_id in self.dagbag.dags.keys(): args = self.parser.parse_args(['list_tasks', dag_id]) diff --git a/tests/www/test_security.py b/tests/www/test_security.py index 611effbfa8d4b..93ebb22956ff2 100644 --- a/tests/www/test_security.py +++ b/tests/www/test_security.py @@ -30,6 +30,7 @@ from sqlalchemy import Column, Integer, String, Date, Float +from airflow.exceptions import AirflowException from airflow.www.security import AirflowSecurityManager, DAG_PERMS @@ -173,9 +174,9 @@ def test_has_access(self, mock_has_view_access): mock_has_view_access.return_value = True self.assertTrue(self.security_manager.has_access('perm', 'view', user)) - def test_sync_perm_for_dag(self): + def test_sync_perm_for_dag_creates_permissions_on_view_menus(self): test_dag_id = 'TEST_DAG' - self.security_manager.sync_perm_for_dag(test_dag_id) + self.security_manager.sync_perm_for_dag(test_dag_id, access_control=None) for dag_perm in DAG_PERMS: self.assertIsNotNone(self.security_manager. find_permission_view_menu(dag_perm, test_dag_id)) @@ -192,3 +193,97 @@ def test_has_all_dag_access(self, mock_has_role, mock_has_perm): mock_has_perm.return_value = True self.assertTrue(self.security_manager.has_all_dags_access()) + + def test_access_control_with_non_existent_role(self): + with self.assertRaises(AirflowException) as context: + self.security_manager.sync_perm_for_dag( + dag_id='access-control-test', + access_control={ + 'this-role-does-not-exist': ['can_dag_edit', 'can_dag_read'] + }) + self.assertIn("role does not exist", str(context.exception)) + + def test_access_control_with_invalid_permission(self): + invalid_permissions = [ + 'can_varimport', # a real permission, but not a member of DAG_PERMS + 'can_eat_pudding', # clearly not a real permission + ] + for permission in invalid_permissions: + self.expect_user_is_in_role(self.user, rolename='team-a') + with self.assertRaises(AirflowException) as context: + self.security_manager.sync_perm_for_dag( + 'access_control_test', + access_control={ + 'team-a': {permission} + }) + self.assertIn("invalid permissions", str(context.exception)) + + def test_access_control_is_set_on_init(self): + self.expect_user_is_in_role(self.user, rolename='team-a') + self.security_manager.sync_perm_for_dag( + 'access_control_test', + access_control={ + 'team-a': ['can_dag_edit', 'can_dag_read'] + }) + self.assert_user_has_dag_perms( + perms=['can_dag_edit', 'can_dag_read'], + dag_id='access_control_test', + ) + + self.expect_user_is_in_role(self.user, rolename='NOT-team-a') + self.assert_user_does_not_have_dag_perms( + perms=['can_dag_edit', 'can_dag_read'], + dag_id='access_control_test', + ) + + def test_access_control_stale_perms_are_revoked(self): + READ_WRITE = {'can_dag_read', 'can_dag_edit'} + READ_ONLY = {'can_dag_read'} + + self.expect_user_is_in_role(self.user, rolename='team-a') + self.security_manager.sync_perm_for_dag( + 'access_control_test', + access_control={'team-a': READ_WRITE}) + self.assert_user_has_dag_perms( + perms=READ_WRITE, + dag_id='access_control_test', + ) + + self.security_manager.sync_perm_for_dag( + 'access_control_test', + access_control={'team-a': READ_ONLY}) + self.assert_user_has_dag_perms( + perms=['can_dag_read'], + dag_id='access_control_test', + ) + self.assert_user_does_not_have_dag_perms( + perms=['can_dag_edit'], + dag_id='access_control_test', + ) + + def expect_user_is_in_role(self, user, rolename): + self.security_manager.init_role(rolename, [], []) + role = self.security_manager.find_role(rolename) + if not role: + self.security_manager.add_role(rolename) + role = self.security_manager.find_role(rolename) + user.roles = [role] + self.security_manager.update_user(user) + + def assert_user_has_dag_perms(self, perms, dag_id): + for perm in perms: + self.assertTrue( + self._has_dag_perm(perm, dag_id), + "User should have '{}' on DAG '{}'".format(perm, dag_id)) + + def assert_user_does_not_have_dag_perms(self, dag_id, perms): + for perm in perms: + self.assertFalse( + self._has_dag_perm(perm, dag_id), + "User should not have '{}' on DAG '{}'".format(perm, dag_id)) + + def _has_dag_perm(self, perm, dag_id): + return self.security_manager.has_access( + perm, + dag_id, + self.user)