Add handling state of existing Dataproc batch (#24924)

This change avoids Airflow marking tasks as 'Success' even if the existing Batch is in a 'Failed' state. We check the various states, and ensure that the Airflow task state reflects the actual state of the Dataproc Batch. Co-authored-by: Daniel van der Ende <[email protected]>
apache · Jul 12, 2022 · b777514 · b777514
1 parent 1c7a4ac
commit b777514
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 0 deletions.
diff --git a/airflow/providers/google/cloud/operators/dataproc.py b/airflow/providers/google/cloud/operators/dataproc.py
@@ -2049,6 +2049,22 @@ def execute(self, context: 'Context'):
  timeout=self.timeout,
  metadata=self.metadata,
  )
+
+ # The existing batch may be a in a number of states other than 'SUCCEEDED'
+ if result.state != Batch.State.SUCCEEDED:
+ if result.state == Batch.State.FAILED or result.state == Batch.State.CANCELLED:
+ raise AirflowException(
+ f"Existing Batch {self.batch_id} failed or cancelled. "
+ f"Error: {result.state_message}"
+ )
+ else:
+ # Batch state is either: RUNNING, PENDING, CANCELLING, or UNSPECIFIED
+ self.log.info(
+ f"Batch {self.batch_id} is in state {result.state.name}."
+ "Waiting for state change..."
+ )
+ result = hook.wait_for_operation(timeout=self.timeout, operation=result)
+
  batch_id = self.batch_id or result.name.split('/')[-1]
  DataprocLink.persist(context=context, task_instance=self, url=DATAPROC_BATCH_LINK, resource=batch_id)
  return Batch.to_dict(result)

diff --git a/tests/providers/google/cloud/operators/test_dataproc.py b/tests/providers/google/cloud/operators/test_dataproc.py
@@ -23,6 +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 airflow import AirflowException
 from airflow.exceptions import AirflowTaskTimeout
@@ -1658,6 +1659,35 @@ def test_execute(self, mock_hook, to_dict_mock):
  metadata=METADATA,
  )
 
+ @mock.patch(DATAPROC_PATH.format("Batch.to_dict"))
+ @mock.patch(DATAPROC_PATH.format("DataprocHook"))
+ def test_execute_batch_failed(self, mock_hook, to_dict_mock):
+ op = DataprocCreateBatchOperator(
+ task_id=TASK_ID,
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ region=GCP_LOCATION,
+ project_id=GCP_PROJECT,
+ batch=BATCH,
+ batch_id=BATCH_ID,
+ request_id=REQUEST_ID,
+ retry=RETRY,
+ timeout=TIMEOUT,
+ metadata=METADATA,
+ )
+ mock_hook.return_value.create_batch.side_effect = AlreadyExists("")
+ mock_hook.return_value.get_batch.return_value.state = Batch.State.FAILED
+ with pytest.raises(AirflowException):
+ op.execute(context=MagicMock())
+ mock_hook.return_value.get_batch.assert_called_once_with(
+ batch_id=BATCH_ID,
+ region=GCP_LOCATION,
+ project_id=GCP_PROJECT,
+ retry=RETRY,
+ timeout=TIMEOUT,
+ metadata=METADATA,
+ )
+
 
 class TestDataprocDeleteBatchOperator:
  @mock.patch(DATAPROC_PATH.format("DataprocHook"))