From b59b0fdceddfd74722316455734a74522a640f5a Mon Sep 17 00:00:00 2001 From: Louis Simoneau Date: Fri, 27 Mar 2020 07:27:29 +1100 Subject: [PATCH] [AIRFLOW-6399] Serialization: DAG access_control field should be decorated field in DAG serialization (#7879) Co-authored by: Louis Simoneau Co-authored-by: Long Le Xich (cherry picked from commit 2d0db948f0bffa7077b4dd581d32cae6ad1aee0e) --- airflow/serialization/serialized_objects.py | 2 +- tests/serialization/test_dag_serialization.py | 25 ++++++++++++++++--- 2 files changed, 23 insertions(+), 4 deletions(-) diff --git a/airflow/serialization/serialized_objects.py b/airflow/serialization/serialized_objects.py index 593e985af43f09..3d2029ac866b0e 100644 --- a/airflow/serialization/serialized_objects.py +++ b/airflow/serialization/serialized_objects.py @@ -500,7 +500,7 @@ class SerializedDAG(DAG, BaseSerialization): not pickle-able. SerializedDAG works for all DAGs. """ - _decorated_fields = {'schedule_interval', 'default_args'} + _decorated_fields = {'schedule_interval', 'default_args', '_access_control'} @staticmethod def __get_constructor_defaults(): # pylint: disable=no-method-argument diff --git a/tests/serialization/test_dag_serialization.py b/tests/serialization/test_dag_serialization.py index b5c150daeceb1c..0d074481ba71ad 100644 --- a/tests/serialization/test_dag_serialization.py +++ b/tests/serialization/test_dag_serialization.py @@ -95,6 +95,18 @@ }, ], "timezone": "UTC", + "_access_control": { + "__type": "dict", + "__var": { + "test_role": { + "__type": "set", + "__var": [ + "can_dag_read", + "can_dag_edit" + ] + } + } + } }, } @@ -116,6 +128,9 @@ def make_simple_dag(): }, start_date=datetime(2019, 8, 1), is_paused_upon_creation=False, + access_control={ + "test_role": {"can_dag_read", "can_dag_edit"} + } ) as dag: CustomOperator(task_id='custom_task') BashOperator(task_id='bash_task', bash_command='echo {{ task.task_id }}', owner='airflow') @@ -215,12 +230,16 @@ def validate_serialized_dag(self, json_dag, ground_truth_dag): def sorted_serialized_dag(dag_dict): """ - Sorts the "tasks" list in the serialised dag python dictionary - This is needed as the order of tasks should not matter but assertEqual - would fail if the order of tasks list changes in dag dictionary + Sorts the "tasks" list and "access_control" permissions in the + serialised dag python dictionary. This is needed as the order of + items should not matter but assertEqual would fail if the order of + items changes in the dag dictionary """ dag_dict["dag"]["tasks"] = sorted(dag_dict["dag"]["tasks"], key=lambda x: sorted(x.keys())) + dag_dict["dag"]["_access_control"]["__var"]["test_role"]["__var"] = sorted( + dag_dict["dag"]["_access_control"]["__var"]["test_role"]["__var"] + ) return dag_dict self.assertEqual(sorted_serialized_dag(ground_truth_dag),