Bugfix: Sync Access Control defined in DAGs when running sync-perm

fixes #13376
apache · Dec 29, 2020 · 6bdf0c2 · 6bdf0c2
1 parent bafd258
commit 6bdf0c2
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 13 deletions.
diff --git a/airflow/cli/commands/sync_perm_command.py b/airflow/cli/commands/sync_perm_command.py
@@ -30,6 +30,8 @@ def sync_perm(args):
     # Add missing permissions for all the Base Views
     appbuilder.add_permissions(update_perms=True)
     print('Updating permission on all DAG views')
-    dags = DagBag(read_dags_from_db=True).dags.values()
+    dagbag = DagBag(read_dags_from_db=True)
+    dagbag.collect_dags_from_db()
+    dags = dagbag.dags.values()
     for dag in dags:
         appbuilder.sm.sync_perm_for_dag(dag.dag_id, dag.access_control)
diff --git a/tests/cli/commands/test_sync_perm_command.py b/tests/cli/commands/test_sync_perm_command.py
@@ -35,13 +35,18 @@ def setUpClass(cls):
     @mock.patch("airflow.cli.commands.sync_perm_command.cached_app")
     @mock.patch("airflow.cli.commands.sync_perm_command.DagBag")
     def test_cli_sync_perm(self, dagbag_mock, mock_cached_app):
-        self.expect_dagbag_contains(
-            [
-                DAG('has_access_control', access_control={'Public': {permissions.ACTION_CAN_READ}}),
-                DAG('no_access_control'),
-            ],
-            dagbag_mock,
-        )
+        dags = [
+            DAG('has_access_control', access_control={'Public': {permissions.ACTION_CAN_READ}}),
+            DAG('no_access_control'),
+        ]
+
+        collect_dags_from_db_mock = mock.Mock()
+        dagbag = mock.Mock()
+
+        dagbag.dags = {dag.dag_id: dag for dag in dags}
+        dagbag.collect_dags_from_db = collect_dags_from_db_mock
+        dagbag_mock.return_value = dagbag
+
         appbuilder = mock_cached_app.return_value.appbuilder
         appbuilder.sm = mock.Mock()
 
@@ -51,6 +56,7 @@ def test_cli_sync_perm(self, dagbag_mock, mock_cached_app):
         assert appbuilder.sm.sync_roles.call_count == 1
 
         dagbag_mock.assert_called_once_with(read_dags_from_db=True)
+        collect_dags_from_db_mock.assert_called_once_with()
         self.assertEqual(2, len(appbuilder.sm.sync_perm_for_dag.mock_calls))
         appbuilder.sm.sync_perm_for_dag.assert_any_call(
             'has_access_control', {'Public': {permissions.ACTION_CAN_READ}}
@@ -60,8 +66,3 @@ def test_cli_sync_perm(self, dagbag_mock, mock_cached_app):
             None,
         )
         appbuilder.add_permissions.assert_called_once_with(update_perms=True)
-
-    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