From da464eddb894f3ebdecaf03986112ba1fdfd260a Mon Sep 17 00:00:00 2001 From: Phani Kumar Date: Tue, 16 May 2023 14:58:45 +0530 Subject: [PATCH 1/4] Optimize deferred execution for DataprocSubmitJobOperator --- airflow/providers/google/cloud/operators/dataproc.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/airflow/providers/google/cloud/operators/dataproc.py b/airflow/providers/google/cloud/operators/dataproc.py index 985278d4d9d06..44fe37f0a63ad 100644 --- a/airflow/providers/google/cloud/operators/dataproc.py +++ b/airflow/providers/google/cloud/operators/dataproc.py @@ -2029,6 +2029,16 @@ def execute(self, context: Context): ) self.job_id = new_job_id + + job = self.hook.get_job(project_id=self.project_id, region=self.region, job_id=self.job_id) + state = job.status.state + if state == JobStatus.State.DONE: + return self.job_id + elif state == JobStatus.State.ERROR: + raise AirflowException(f"Job failed:\n{job}") + elif state == JobStatus.State.CANCELLED: + raise AirflowException(f"Job was cancelled:\n{job}") + if self.deferrable: self.defer( trigger=DataprocSubmitTrigger( From eee7ca68754ccbd6111570920ad22743f89bf362 Mon Sep 17 00:00:00 2001 From: Phani Kumar Date: Tue, 16 May 2023 18:00:56 +0530 Subject: [PATCH 2/4] Add test --- .../google/cloud/operators/test_dataproc.py | 31 +++++++++++++++++-- 1 file changed, 29 insertions(+), 2 deletions(-) diff --git a/tests/providers/google/cloud/operators/test_dataproc.py b/tests/providers/google/cloud/operators/test_dataproc.py index 9c2e202f97c5a..33ad5bfacd8c2 100644 --- a/tests/providers/google/cloud/operators/test_dataproc.py +++ b/tests/providers/google/cloud/operators/test_dataproc.py @@ -23,7 +23,7 @@ import pytest from google.api_core.exceptions import AlreadyExists, NotFound from google.api_core.retry import Retry -from google.cloud.dataproc_v1 import Batch +from google.cloud.dataproc_v1 import Batch, JobStatus from airflow.exceptions import ( AirflowException, @@ -134,7 +134,8 @@ "gke_cluster_target": "projects/project_id/locations/region/clusters/gke_cluster_name", "node_pool_target": [ { - "node_pool": "projects/project_id/locations/region/clusters/gke_cluster_name/nodePools/dp", # noqa + "node_pool": "projects/project_id/locations/region/clusters/gke_cluster_name/nodePools/" + "dp", "roles": ["DEFAULT"], } ], @@ -1058,6 +1059,32 @@ def test_execute_deferrable(self, mock_trigger_hook, mock_hook): assert isinstance(exc.value.trigger, DataprocSubmitTrigger) assert exc.value.method_name == GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME + @mock.patch(DATAPROC_PATH.format("DataprocHook")) + @mock.patch("airflow.providers.google.cloud.operators.dataproc.DataprocSubmitJobOperator.defer") + @mock.patch("airflow.providers.google.cloud.operators.dataproc.DataprocHook.submit_job") + def test_dataproc_operator_execute_async_done_before_defer(self, mock_submit_job, mock_defer, mock_hook): + mock_submit_job.return_value.reference.job_id = TEST_JOB_ID + job_status = mock_hook.return_value.get_job.return_value.status + job_status.state = JobStatus.State.DONE + + op = DataprocSubmitJobOperator( + task_id=TASK_ID, + region=GCP_REGION, + project_id=GCP_PROJECT, + job={}, + gcp_conn_id=GCP_CONN_ID, + retry=RETRY, + asynchronous=True, + timeout=TIMEOUT, + metadata=METADATA, + request_id=REQUEST_ID, + impersonation_chain=IMPERSONATION_CHAIN, + deferrable=True, + ) + + op.execute(context=self.mock_context) + assert not mock_defer.called + @mock.patch(DATAPROC_PATH.format("DataprocHook")) def test_on_kill(self, mock_hook): job = {} From c3527ab9d42a66eb99d78edd3a38232410505165 Mon Sep 17 00:00:00 2001 From: Phani Kumar Date: Tue, 16 May 2023 18:02:27 +0530 Subject: [PATCH 3/4] Add test --- tests/providers/google/cloud/operators/test_dataproc.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/providers/google/cloud/operators/test_dataproc.py b/tests/providers/google/cloud/operators/test_dataproc.py index 33ad5bfacd8c2..5494a88ba9330 100644 --- a/tests/providers/google/cloud/operators/test_dataproc.py +++ b/tests/providers/google/cloud/operators/test_dataproc.py @@ -134,8 +134,7 @@ "gke_cluster_target": "projects/project_id/locations/region/clusters/gke_cluster_name", "node_pool_target": [ { - "node_pool": "projects/project_id/locations/region/clusters/gke_cluster_name/nodePools/" - "dp", + "node_pool": "projects/project_id/locations/region/clusters/gke_cluster_name/nodePools/dp", # noqa "roles": ["DEFAULT"], } ], From 2352117d0317dfe12669271314d9a1e149bdc072 Mon Sep 17 00:00:00 2001 From: Phani Kumar Date: Fri, 2 Jun 2023 11:54:21 +0530 Subject: [PATCH 4/4] Apply review suggestions --- .../google/cloud/operators/dataproc.py | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/airflow/providers/google/cloud/operators/dataproc.py b/airflow/providers/google/cloud/operators/dataproc.py index 44fe37f0a63ad..3019a536d91a2 100644 --- a/airflow/providers/google/cloud/operators/dataproc.py +++ b/airflow/providers/google/cloud/operators/dataproc.py @@ -2029,17 +2029,15 @@ def execute(self, context: Context): ) self.job_id = new_job_id - - job = self.hook.get_job(project_id=self.project_id, region=self.region, job_id=self.job_id) - state = job.status.state - if state == JobStatus.State.DONE: - return self.job_id - elif state == JobStatus.State.ERROR: - raise AirflowException(f"Job failed:\n{job}") - elif state == JobStatus.State.CANCELLED: - raise AirflowException(f"Job was cancelled:\n{job}") - if self.deferrable: + job = self.hook.get_job(project_id=self.project_id, region=self.region, job_id=self.job_id) + state = job.status.state + if state == JobStatus.State.DONE: + return self.job_id + elif state == JobStatus.State.ERROR: + raise AirflowException(f"Job failed:\n{job}") + elif state == JobStatus.State.CANCELLED: + raise AirflowException(f"Job was cancelled:\n{job}") self.defer( trigger=DataprocSubmitTrigger( job_id=self.job_id,