[AIRFLOW-2694] Declare permissions in DAG definition

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(),
         },
         {

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

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,19 +452,90 @@ 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:
             perm_on_dag = self.find_permission_view_menu(dag_perm, 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.

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('/')
 

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])

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)