From cc1e838ab7223e31f1b264aad9766d198bb7f1ce Mon Sep 17 00:00:00 2001 From: Hossein Torabi Date: Sun, 8 Aug 2021 22:51:17 +0200 Subject: [PATCH] deduplicate running jobs on BigQueryInsertJobOperator --- airflow/providers/google/cloud/operators/bigquery.py | 7 ++----- tests/providers/google/cloud/operators/test_bigquery.py | 2 +- 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/airflow/providers/google/cloud/operators/bigquery.py b/airflow/providers/google/cloud/operators/bigquery.py index fe853c22aca35..e89d6a861343e 100644 --- a/airflow/providers/google/cloud/operators/bigquery.py +++ b/airflow/providers/google/cloud/operators/bigquery.py @@ -2247,16 +2247,13 @@ def _submit_job( hook: BigQueryHook, job_id: str, ) -> BigQueryJob: - # Submit a new job - job = hook.insert_job( + # Submit a new job and wait for it to complete and get the result. + return hook.insert_job( configuration=self.configuration, project_id=self.project_id, location=self.location, job_id=job_id, ) - # Start the job and wait for it to complete and get the result. - job.result() - return job @staticmethod def _handle_job_error(job: BigQueryJob) -> None: diff --git a/tests/providers/google/cloud/operators/test_bigquery.py b/tests/providers/google/cloud/operators/test_bigquery.py index 801034c06ce66..1576b24055478 100644 --- a/tests/providers/google/cloud/operators/test_bigquery.py +++ b/tests/providers/google/cloud/operators/test_bigquery.py @@ -930,7 +930,7 @@ def test_execute_reattach(self, mock_hook, mock_md5): } } - mock_hook.return_value.insert_job.return_value.result.side_effect = Conflict("any") + mock_hook.return_value.insert_job.side_effect = Conflict("any") job = MagicMock( job_id=real_job_id, error_result=False,