Check for same task instead of Equality to detect Duplicate Tasks (#8828

)
apache · U-Ozdemir · Apr 30, 2020 · Apr 30, 2020 · Apr 30, 2020 · May 1, 2020
commit 15273f0ea05ec579c631ce26b5d620233ebdc4d2
diff --git a/airflow/example_dags/example_complex.py b/airflow/example_dags/example_complex.py
@@ -81,7  81,7 @@
     )
 
     create_tag_template_field_result2 = BashOperator(
-        task_id="create_tag_template_field_result", bash_command="echo create_tag_template_field_result"
         task_id="create_tag_template_field_result2", bash_command="echo create_tag_template_field_result"
     )
 
     # Delete

diff --git a/airflow/models/baseoperator.py b/airflow/models/baseoperator.py
@@ -36,7  36,7 @@
 from sqlalchemy.orm import Session
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException, DuplicateTaskIdFound
 from airflow.exceptions import AirflowException
 from airflow.lineage import apply_lineage, prepare_lineage
 from airflow.models.base import Operator
 from airflow.models.pool import Pool
@@ -600,9  600,8 @@ def dag(self, dag: Any):
                 "The DAG assigned to {} can not be changed.".format(self))
         elif self.task_id not in dag.task_dict:
             dag.add_task(self)
-        elif self.task_id in dag.task_dict and dag.task_dict[self.task_id] != self:
-            raise DuplicateTaskIdFound(
-                "Task id '{}' has already been added to the DAG".format(self.task_id))
         elif self.task_id in dag.task_dict and dag.task_dict[self.task_id] is not self:
             dag.add_task(self)
 
         self._dag = dag  # pylint: disable=attribute-defined-outside-init
 

diff --git a/airflow/models/dag.py b/airflow/models/dag.py
@@ -1337,7  1337,7 @@ def add_task(self, task):
         elif task.end_date and self.end_date:
             task.end_date = min(task.end_date, self.end_date)
 
-        if task.task_id in self.task_dict and self.task_dict[task.task_id] != task:
         if task.task_id in self.task_dict and self.task_dict[task.task_id] is not task:
             raise DuplicateTaskIdFound(
                 "Task id '{}' has already been added to the DAG".format(task.task_id))
         else:

diff --git a/airflow/providers/google/cloud/example_dags/example_datacatalog.py b/airflow/providers/google/cloud/example_dags/example_datacatalog.py
@@ -181,7  181,7 @@
 
     # [START howto_operator_gcp_datacatalog_create_tag_template_field_result2]
     create_tag_template_field_result2 = BashOperator(
-        task_id="create_tag_template_field_result",
         task_id="create_tag_template_field_result2",
         bash_command="echo \"{{ task_instance.xcom_pull('create_tag_template_field') }}\"",
     )
     # [END howto_operator_gcp_datacatalog_create_tag_template_field_result2]

diff --git a/airflow/providers/google/cloud/example_dags/example_gcs_to_gcs.py b/airflow/providers/google/cloud/example_dags/example_gcs_to_gcs.py
@@ -100,7  100,7 @@
 
     # [START howto_operator_gcs_to_gcs_delimiter]
     copy_files_with_delimiter = GCSToGCSOperator(
-        task_id="copy_files_with_wildcard",
         task_id="copy_files_with_delimiter",
         source_bucket=BUCKET_1_SRC,
         source_object="data/",
         destination_bucket=BUCKET_1_DST,

diff --git a/tests/models/test_dag.py b/tests/models/test_dag.py
@@ -979,34  979,18 @@ def test_duplicate_task_ids_not_allowed_with_dag_context_manager(self):
 
         self.assertEqual(dag.task_dict, {op1.task_id: op1})
 
-        # Also verify that DAGs with duplicate task_ids don't raise errors
-        with DAG("test_dag_1", start_date=DEFAULT_DATE) as dag1:
-            op3 = DummyOperator(task_id="t3")
-            op4 = BashOperator(task_id="t4", bash_command="sleep 1")
-            op3 >> op4
-
-        self.assertEqual(dag1.task_dict, {op3.task_id: op3, op4.task_id: op4})
-
     def test_duplicate_task_ids_not_allowed_without_dag_context_manager(self):
         """Verify tasks with Duplicate task_id raises error"""
         with self.assertRaisesRegex(
             DuplicateTaskIdFound, "Task id 't1' has already been added to the DAG"
         ):
             dag = DAG("test_dag", start_date=DEFAULT_DATE)
             op1 = DummyOperator(task_id="t1", dag=dag)
-            op2 = BashOperator(task_id="t1", bash_command="sleep 1", dag=dag)
             op2 = DummyOperator(task_id="t1", dag=dag)
             op1 >> op2
 
         self.assertEqual(dag.task_dict, {op1.task_id: op1})
 
-        # Also verify that DAGs with duplicate task_ids don't raise errors
-        dag1 = DAG("test_dag_1", start_date=DEFAULT_DATE)
-        op3 = DummyOperator(task_id="t3", dag=dag1)
-        op4 = DummyOperator(task_id="t4", dag=dag1)
-        op3 >> op4
-
-        self.assertEqual(dag1.task_dict, {op3.task_id: op3, op4.task_id: op4})
-
     def test_duplicate_task_ids_for_same_task_is_allowed(self):
         """Verify that same tasks with Duplicate task_id do not raise error"""
         with DAG("test_dag", start_date=DEFAULT_DATE) as dag:

diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py
@@ -373,20  373,20 @@ def test_ti_updates_with_task(self, session=None):
         """
         test that updating the executor_config propogates to the TaskInstance DB
         """
-        dag = models.DAG(dag_id='test_run_pooling_task')
-        task = DummyOperator(task_id='test_run_pooling_task_op', dag=dag, owner='airflow',
-                             executor_config={'foo': 'bar'},
-                             start_date=timezone.datetime(2016, 2, 1, 0, 0, 0))
         with models.DAG(dag_id='test_run_pooling_task') as dag:
             task = DummyOperator(task_id='test_run_pooling_task_op', owner='airflow',
                                  executor_config={'foo': 'bar'},
                                  start_date=timezone.datetime(2016, 2, 1, 0, 0, 0))
         ti = TI(
             task=task, execution_date=timezone.utcnow())
 
         ti.run(session=session)
         tis = dag.get_task_instances()
         self.assertEqual({'foo': 'bar'}, tis[0].executor_config)
-
-        task2 = DummyOperator(task_id='test_run_pooling_task_op', dag=dag, owner='airflow',
-                              executor_config={'bar': 'baz'},
-                              start_date=timezone.datetime(2016, 2, 1, 0, 0, 0))
         with models.DAG(dag_id='test_run_pooling_task') as dag:
             task2 = DummyOperator(task_id='test_run_pooling_task_op', owner='airflow',
                                   executor_config={'bar': 'baz'},
                                   start_date=timezone.datetime(2016, 2, 1, 0, 0, 0))
 
         ti = TI(
             task=task2, execution_date=timezone.utcnow())

diff --git a/tests/providers/amazon/aws/operators/test_s3_to_sftp.py b/tests/providers/amazon/aws/operators/test_s3_to_sftp.py
@@ -133,7  133,7 @@ def test_s3_to_sftp_operation(self):
     def delete_remote_resource(self):
         # check the remote file content
         remove_file_task = SSHOperator(
-            task_id="test_check_file",
             task_id="test_rm_file",
             ssh_hook=self.hook,
             command="rm {0}".format(self.sftp_path),
             do_xcom_push=True,

diff --git a/tests/providers/google/cloud/operators/test_mlengine_utils.py b/tests/providers/google/cloud/operators/test_mlengine_utils.py
@@ -131,16  131,18 @@ def test_successful_run(self):
             self.assertEqual('err=0.9', result)
 
     def test_failures(self):
-        dag = DAG(
-            'test_dag',
-            default_args={
-                'owner': 'airflow',
-                'start_date': DEFAULT_DATE,
-                'end_date': DEFAULT_DATE,
-                'project_id': 'test-project',
-                'region': 'us-east1',
-            },
-            schedule_interval='@daily')
         def create_test_dag(dag_id):
             dag = DAG(
                 dag_id,
                 default_args={
                     'owner': 'airflow',
                     'start_date': DEFAULT_DATE,
                     'end_date': DEFAULT_DATE,
                     'project_id': 'test-project',
                     'region': 'us-east1',
                 },
                 schedule_interval='@daily')
             return dag
 
         input_with_model = self.INPUT_MISSING_ORIGIN.copy()
         other_params_but_models = {
@@ -151,26  153,30 @@ def test_failures(self):
             'prediction_path': input_with_model['outputPath'],
             'metric_fn_and_keys': (self.metric_fn, ['err']),
             'validate_fn': (lambda x: 'err=%.1f' % x['err']),
-            'dag': dag,
         }
 
         with self.assertRaisesRegex(AirflowException, 'Missing model origin'):
-            mlengine_operator_utils.create_evaluate_ops(**other_params_but_models)
             mlengine_operator_utils.create_evaluate_ops(
                 dag=create_test_dag('test_dag_1'), **other_params_but_models)
 
         with self.assertRaisesRegex(AirflowException, 'Ambiguous model origin'):
-            mlengine_operator_utils.create_evaluate_ops(model_uri='abc', model_name='cde',
-                                                        **other_params_but_models)
             mlengine_operator_utils.create_evaluate_ops(
                 dag=create_test_dag('test_dag_2'), model_uri='abc', model_name='cde',
                 **other_params_but_models)
 
         with self.assertRaisesRegex(AirflowException, 'Ambiguous model origin'):
-            mlengine_operator_utils.create_evaluate_ops(model_uri='abc', version_name='vvv',
-                                                        **other_params_but_models)
             mlengine_operator_utils.create_evaluate_ops(
                 dag=create_test_dag('test_dag_3'), model_uri='abc', version_name='vvv',
                 **other_params_but_models)
 
         with self.assertRaisesRegex(AirflowException, '`metric_fn` param must be callable'):
             params = other_params_but_models.copy()
             params['metric_fn_and_keys'] = (None, ['abc'])
-            mlengine_operator_utils.create_evaluate_ops(model_uri='gs://blah', **params)
             mlengine_operator_utils.create_evaluate_ops(
                 dag=create_test_dag('test_dag_4'), model_uri='gs://blah', **params)
 
         with self.assertRaisesRegex(AirflowException, '`validate_fn` param must be callable'):
             params = other_params_but_models.copy()
             params['validate_fn'] = None
-            mlengine_operator_utils.create_evaluate_ops(model_uri='gs://blah', **params)
             mlengine_operator_utils.create_evaluate_ops(
                 dag=create_test_dag('test_dag_5'), model_uri='gs://blah', **params)
diff --git a/tests/providers/google/cloud/sensors/test_gcs.py b/tests/providers/google/cloud/sensors/test_gcs.py
@@ -205,7  205,7 @@ def setUp(self):
         self.dag = dag
 
         self.sensor = GCSUploadSessionCompleteSensor(
-            task_id='sensor',
             task_id='sensor_1',
             bucket='test-bucket',
             prefix='test-prefix/path',
             inactivity_period=12,
@@ -227,7  227,7 @@ def test_files_deleted_between_pokes_throw_error(self):
     @mock.patch('airflow.providers.google.cloud.sensors.gcs.get_time', mock_time)
     def test_files_deleted_between_pokes_allow_delete(self):
         self.sensor = GCSUploadSessionCompleteSensor(
-            task_id='sensor',
             task_id='sensor_2',
             bucket='test-bucket',
             prefix='test-prefix/path',
             inactivity_period=12,

diff --git a/tests/providers/microsoft/azure/operators/test_file_to_wasb.py b/tests/providers/microsoft/azure/operators/test_file_to_wasb.py
@@ -45,7  45,7 @@ def setUp(self):
 
     def test_init(self):
         operator = FileToWasbOperator(
-            task_id='wasb_operator',
             task_id='wasb_operator_1',
             dag=self.dag,
             **self._config
         )
@@ -58,7  58,7 @@ def test_init(self):
         self.assertEqual(operator.retries, self._config['retries'])
 
         operator = FileToWasbOperator(
-            task_id='wasb_operator',
             task_id='wasb_operator_2',
             dag=self.dag,
             load_options={'timeout': 2},
             **self._config

diff --git a/tests/providers/microsoft/azure/operators/test_wasb_delete_blob.py b/tests/providers/microsoft/azure/operators/test_wasb_delete_blob.py
@@ -42,7  42,7 @@ def setUp(self):
 
     def test_init(self):
         operator = WasbDeleteBlobOperator(
-            task_id='wasb_operator',
             task_id='wasb_operator_1',
             dag=self.dag,
             **self._config
         )
@@ -53,7  53,7 @@ def test_init(self):
         self.assertEqual(operator.ignore_if_missing, False)
 
         operator = WasbDeleteBlobOperator(
-            task_id='wasb_operator',
             task_id='wasb_operator_2',
             dag=self.dag,
             is_prefix=True,
             ignore_if_missing=True,

diff --git a/tests/providers/microsoft/azure/sensors/test_wasb.py b/tests/providers/microsoft/azure/sensors/test_wasb.py
@@ -43,7  43,7 @@ def setUp(self):
 
     def test_init(self):
         sensor = WasbBlobSensor(
-            task_id='wasb_sensor',
             task_id='wasb_sensor_1',
             dag=self.dag,
             **self._config
         )
@@ -54,7  54,7 @@ def test_init(self):
         self.assertEqual(sensor.timeout, self._config['timeout'])
 
         sensor = WasbBlobSensor(
-            task_id='wasb_sensor',
             task_id='wasb_sensor_2',
             dag=self.dag,
             check_options={'timeout': 2},
             **self._config
@@ -94,7  94,7 @@ def setUp(self):
 
     def test_init(self):
         sensor = WasbPrefixSensor(
-            task_id='wasb_sensor',
             task_id='wasb_sensor_1',
             dag=self.dag,
             **self._config
         )
@@ -105,7  105,7 @@ def test_init(self):
         self.assertEqual(sensor.timeout, self._config['timeout'])
 
         sensor = WasbPrefixSensor(
-            task_id='wasb_sensor',
             task_id='wasb_sensor_2',
             dag=self.dag,
             check_options={'timeout': 2},
             **self._config

diff --git a/tests/providers/sftp/operators/test_sftp.py b/tests/providers/sftp/operators/test_sftp.py
@@ -75,7  75,7 @@ def test_pickle_file_transfer_put(self):
 
         # put test file to remote
         put_test_task = SFTPOperator(
-            task_id="test_sftp",
             task_id="put_test_task",
             ssh_hook=self.hook,
             local_filepath=self.test_local_filepath,
             remote_filepath=self.test_remote_filepath,
@@ -89,7  89,7 @@ def test_pickle_file_transfer_put(self):
 
         # check the remote file content
         check_file_task = SSHOperator(
-            task_id="test_check_file",
             task_id="check_file_task",
             ssh_hook=self.hook,
             command="cat {0}".format(self.test_remote_filepath),
             do_xcom_push=True,
@@ -99,7  99,7 @@ def test_pickle_file_transfer_put(self):
         ti3 = TaskInstance(task=check_file_task, execution_date=timezone.utcnow())
         ti3.run()
         self.assertEqual(
-            ti3.xcom_pull(task_ids='test_check_file', key='return_value').strip(),
             ti3.xcom_pull(task_ids=check_file_task.task_id, key='return_value').strip(),
             test_local_file_content)
 
     @conf_vars({('core', 'enable_xcom_pickling'): 'True'})
@@ -178,7  178,7 @@ def test_json_file_transfer_put(self):
 
         # put test file to remote
         put_test_task = SFTPOperator(
-            task_id="test_sftp",
             task_id="put_test_task",
             ssh_hook=self.hook,
             local_filepath=self.test_local_filepath,
             remote_filepath=self.test_remote_filepath,
@@ -191,7  191,7 @@ def test_json_file_transfer_put(self):
 
         # check the remote file content
         check_file_task = SSHOperator(
-            task_id="test_check_file",
             task_id="check_file_task",
             ssh_hook=self.hook,
             command="cat {0}".format(self.test_remote_filepath),
             do_xcom_push=True,
@@ -201,7  201,7 @@ def test_json_file_transfer_put(self):
         ti3 = TaskInstance(task=check_file_task, execution_date=timezone.utcnow())
         ti3.run()
         self.assertEqual(
-            ti3.xcom_pull(task_ids='test_check_file', key='return_value').strip(),
             ti3.xcom_pull(task_ids=check_file_task.task_id, key='return_value').strip(),
             b64encode(test_local_file_content).decode('utf-8'))
 
     @conf_vars({('core', 'enable_xcom_pickling'): 'True'})
@@ -362,7  362,7 @@ def test_arg_checking(self):
         with self.assertRaisesRegex(AirflowException,
                                     "Cannot operate without ssh_hook or ssh_conn_id."):
             task_0 = SFTPOperator(
-                task_id="test_sftp",
                 task_id="test_sftp_0",
                 local_filepath=self.test_local_filepath,
                 remote_filepath=self.test_remote_filepath,
                 operation=SFTPOperation.PUT,
@@ -372,7  372,7 @@ def test_arg_checking(self):
 
         # if ssh_hook is invalid/not provided, use ssh_conn_id to create SSHHook
         task_1 = SFTPOperator(
-            task_id="test_sftp",
             task_id="test_sftp_1",
             ssh_hook="string_rather_than_SSHHook",  # invalid ssh_hook
             ssh_conn_id=TEST_CONN_ID,
             local_filepath=self.test_local_filepath,
@@ -387,7  387,7 @@ def test_arg_checking(self):
         self.assertEqual(task_1.ssh_hook.ssh_conn_id, TEST_CONN_ID)
 
         task_2 = SFTPOperator(
-            task_id="test_sftp",
             task_id="test_sftp_2",
             ssh_conn_id=TEST_CONN_ID,  # no ssh_hook provided
             local_filepath=self.test_local_filepath,
             remote_filepath=self.test_remote_filepath,
@@ -402,7  402,7 @@ def test_arg_checking(self):
 
         # if both valid ssh_hook and ssh_conn_id are provided, ignore ssh_conn_id
         task_3 = SFTPOperator(
-            task_id="test_sftp",
             task_id="test_sftp_3",
             ssh_hook=self.hook,
             ssh_conn_id=TEST_CONN_ID,
             local_filepath=self.test_local_filepath,