Correctly restore upstream_task_ids when deserializing Operators (apache#8775)

ashb · web-flow · commit 280f1f0c4cc4 · 2020-05-10T11:41:47.000+01:00
This test exposed a bug in one of the example dags, that wasn't caught by apache#6549. That will be a fixed in a separate issue, but it caused the round-trip tests to fail here Fixes apache#8720
diff --git a/airflow/providers/google/cloud/example_dags/example_gcs.py b/airflow/providers/google/cloud/example_dags/example_gcs.py
@@ -126,8 +126,8 @@
     )
 
     # [START howto_operator_gcs_delete_bucket]
-    delete_bucket_1 = GCSDeleteBucketOperator(task_id="delete_bucket", bucket_name=BUCKET_1)
-    delete_bucket_2 = GCSDeleteBucketOperator(task_id="delete_bucket", bucket_name=BUCKET_2)
+    delete_bucket_1 = GCSDeleteBucketOperator(task_id="delete_bucket_1", bucket_name=BUCKET_1)
+    delete_bucket_2 = GCSDeleteBucketOperator(task_id="delete_bucket_2", bucket_name=BUCKET_2)
     # [END howto_operator_gcs_delete_bucket]
 
     [create_bucket1, create_bucket2] >> list_buckets >> list_buckets_result
diff --git a/airflow/serialization/serialized_objects.py b/airflow/serialization/serialized_objects.py
@@ -598,7 +598,7 @@ def deserialize_dag(cls, encoded_dag: Dict[str, Any]) -> 'SerializedDAG':
             for task_id in serializable_task.downstream_task_ids:
                 # Bypass set_upstream etc here - it does more than we want
                 # noinspection PyProtectedMember
-                dag.task_dict[task_id]._upstream_task_ids.add(task_id)  # pylint: disable=protected-access
+                dag.task_dict[task_id]._upstream_task_ids.add(serializable_task.task_id)  # noqa: E501 # pylint: disable=protected-access
 
         return dag
 
diff --git a/tests/serialization/test_dag_serialization.py b/tests/serialization/test_dag_serialization.py
@@ -355,6 +355,9 @@ def validate_deserialized_task(self, serialized_task, task,):
         assert serialized_task.task_type == task.task_type
         assert set(serialized_task.template_fields) == set(task.template_fields)
 
+        assert serialized_task.upstream_task_ids == task.upstream_task_ids
+        assert serialized_task.downstream_task_ids == task.downstream_task_ids
+
         for field in fields_to_check:
             assert getattr(serialized_task, field) == getattr(task, field), \
                 f'{task.dag.dag_id}.{task.task_id}.{field} does not match'

Original file line number	Diff line number	Diff line change
`@@ -126,8 +126,8 @@`
`126`	`126`	`)`
`127`	`127`
`128`	`128`	`# [START howto_operator_gcs_delete_bucket]`
`129`		`- delete_bucket_1 = GCSDeleteBucketOperator(task_id="delete_bucket", bucket_name=BUCKET_1)`
`130`		`- delete_bucket_2 = GCSDeleteBucketOperator(task_id="delete_bucket", bucket_name=BUCKET_2)`
	`129`	`+ delete_bucket_1 = GCSDeleteBucketOperator(task_id="delete_bucket_1", bucket_name=BUCKET_1)`
	`130`	`+ delete_bucket_2 = GCSDeleteBucketOperator(task_id="delete_bucket_2", bucket_name=BUCKET_2)`
`131`	`131`	`# [END howto_operator_gcs_delete_bucket]`
`132`	`132`
`133`	`133`	`[create_bucket1, create_bucket2] >> list_buckets >> list_buckets_result`