Repository: apache/airflow

Issue:
[Google Cloud] DataprocCreateBatchOperator returns incorrect results and does not reattach
### Apache Airflow version

main (development)

### What happened

The provider operator for Google Cloud Dataproc Batches has two bugs: 

1. The running [operator](https://github.com/apache/airflow/blob/main/airflow/providers/google/cloud/operators/dataproc.py#L2123-L2124) returns successful even if the job transitions to State.CANCELLED or State.CANCELLING 
2. It [attempts](https://github.com/apache/airflow/blob/main/airflow/providers/google/cloud/operators/dataproc.py#L2154) to 'reattach' to a potentially running job if it AlreadyExists, but it sends the wrong type since 'result' is a Batch and needs Operation

### What you think should happen instead

A new hook that polls for batch job completion.  There is precedent for it in traditional dataproc with 'wait_for_job'.

### How to reproduce

Use the Breeze environment and a DAG that runs DataprocCreateBatchOperator.  Allow the first instance to start. 

Use the gcloud CLI to cancel the job. 

`gcloud dataproc batches cancel <batch_id> --project <project_id> --region <region>`

Observe that the task completes successfully after a 3-5 minute timeout, even though the job was cancelled. 

Run the task again with the same batch_id.  Observe the ValueError where it expects Operation but receives Batch



### Operating System

Darwin 5806 21.6.0 Darwin Kernel Version 21.6.0: Mon Aug 22 20:17:10 PDT 2022; root:xnu-8020.140.49~2/RELEASE_X86_64 x86_64

### Versions of Apache Airflow Providers

Same as dev (main) version. 

### Deployment

Other Docker-based deployment

### Deployment details

Observable in the Breeze environment, when running against real Google Infrastructure. 

### Anything else

Every time. 

### Are you willing to submit PR?

- [X] Yes I am willing to submit a PR!

### Code of Conduct

- [X] I agree to follow this project's [Code of Conduct](https://github.com/apache/airflow/blob/main/CODE_OF_CONDUCT.md)



Discussion:

Resolving patch:
diff --git a/airflow/providers/google/cloud/hooks/dataproc.py b/airflow/providers/google/cloud/hooks/dataproc.py
--- a/airflow/providers/google/cloud/hooks/dataproc.py
+++ b/airflow/providers/google/cloud/hooks/dataproc.py
@@ -985,6 +985,68 @@ def list_batches(
         )
         return result
 
+    @GoogleBaseHook.fallback_to_default_project_id
+    def wait_for_batch(
+        self,
+        batch_id: str,
+        region: str,
+        project_id: str,
+        wait_check_interval: int = 10,
+        retry: Retry | _MethodDefault = DEFAULT,
+        timeout: float | None = None,
+        metadata: Sequence[tuple[str, str]] = (),
+    ) -> Batch:
+        """
+        Wait for a Batch job to complete.
+
+        After Batch job submission, the operator will wait for the job to complete, however, this is useful
+        in the case where Airflow is restarted or the task pid is killed for any reason. In this case, the
+        Batch create will happen again, AlreadyExists will be raised and caught, then should fall to this
+        function for waiting on completion.
+
+        :param batch_id: Required. The ID to use for the batch, which will become the final component
+            of the batch's resource name.
+            This value must be 4-63 characters. Valid characters are /[a-z][0-9]-/.
+        :param region: Required. The Cloud Dataproc region in which to handle the request.
+        :param project_id: Required. The ID of the Google Cloud project that the cluster belongs to.
+        :param wait_check_interval: The amount of time to pause between checks for job completion
+        :param retry: A retry object used to retry requests to get_batch.
+            If ``None`` is specified, requests will not be retried.
+        :param timeout: The amount of time, in seconds, to wait for the create_batch request to complete.
+            Note that if ``retry`` is specified, the timeout applies to each individual attempt.
+        :param metadata: Additional metadata that is provided to the method.
+        """
+        state = None
+        first_loop: bool = True
+        while state not in [
+            Batch.State.CANCELLED,
+            Batch.State.FAILED,
+            Batch.State.SUCCEEDED,
+            Batch.State.STATE_UNSPECIFIED,
+        ]:
+            try:
+                if not first_loop:
+                    time.sleep(wait_check_interval)
+                first_loop = False
+                self.log.debug("Waiting for batch %s", batch_id)
+                result = self.get_batch(
+                    batch_id=batch_id,
+                    region=region,
+                    project_id=project_id,
+                    retry=retry,
+                    timeout=timeout,
+                    metadata=metadata,
+                )
+                state = result.state
+            except ServerError as err:
+                self.log.info(
+                    "Retrying. Dataproc API returned server error when waiting for batch id %s: %s",
+                    batch_id,
+                    err,
+                )
+
+        return result
+
 
 class DataprocAsyncHook(GoogleBaseHook):
     """
diff --git a/airflow/providers/google/cloud/operators/dataproc.py b/airflow/providers/google/cloud/operators/dataproc.py
--- a/airflow/providers/google/cloud/operators/dataproc.py
+++ b/airflow/providers/google/cloud/operators/dataproc.py
@@ -2233,7 +2233,15 @@ def __init__(
 
     def execute(self, context: Context):
         hook = DataprocHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
-        self.log.info("Creating batch")
+        # batch_id might not be set and will be generated
+        if self.batch_id:
+            link = DATAPROC_BATCH_LINK.format(
+                region=self.region, project_id=self.project_id, resource=self.batch_id
+            )
+            self.log.info("Creating batch %s", self.batch_id)
+            self.log.info("Once started, the batch job will be available at %s", link)
+        else:
+            self.log.info("Starting batch job. The batch ID will be generated since it was not provided.")
         if self.region is None:
             raise AirflowException("Region should be set here")
         try:
@@ -2275,32 +2283,37 @@ def execute(self, context: Context):
 
         except AlreadyExists:
             self.log.info("Batch with given id already exists")
-            if self.batch_id is None:
-                raise AirflowException("Batch Id should be set here")
-            result = hook.get_batch(
-                batch_id=self.batch_id,
-                region=self.region,
-                project_id=self.project_id,
-                retry=self.retry,
-                timeout=self.timeout,
-                metadata=self.metadata,
-            )
-            # The existing batch may be 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)
-
+            # This is only likely to happen if batch_id was provided
+            # Could be running if Airflow was restarted after task started
+            # poll until a final state is reached
+            if self.batch_id:
+                self.log.info("Attaching to the job (%s) if it is still running.", self.batch_id)
+                result = hook.wait_for_batch(
+                    batch_id=self.batch_id,
+                    region=self.region,
+                    project_id=self.project_id,
+                    retry=self.retry,
+                    timeout=self.timeout,
+                    metadata=self.metadata,
+                    wait_check_interval=self.polling_interval_seconds,
+                )
+        # It is possible we don't have a result in the case where batch_id was not provide, one was generated
+        # by chance, AlreadyExists was caught, but we can't reattach because we don't have the generated id
+        if result is None:
+            raise AirflowException("The job could not be reattached because the id was generated.")
+
+        # The existing batch may be a number of states other than 'SUCCEEDED'\
+        # wait_for_operation doesn't fail if the job is cancelled, so we will check for it here which also
+        # finds a cancelling|canceled|unspecified job from wait_for_batch
         batch_id = self.batch_id or result.name.split("/")[-1]
+        link = DATAPROC_BATCH_LINK.format(region=self.region, project_id=self.project_id, resource=batch_id)
+        if result.state == Batch.State.FAILED:
+            raise AirflowException(f"Batch job {batch_id} failed.  Driver Logs: {link}")
+        if result.state in (Batch.State.CANCELLED, Batch.State.CANCELLING):
+            raise AirflowException(f"Batch job {batch_id} was cancelled. Driver logs: {link}")
+        if result.state == Batch.State.STATE_UNSPECIFIED:
+            raise AirflowException(f"Batch job {batch_id} unspecified. Driver logs: {link}")
+        self.log.info("Batch job %s completed. Driver logs: %s", batch_id, link)
         DataprocLink.persist(context=context, task_instance=self, url=DATAPROC_BATCH_LINK, resource=batch_id)
         return Batch.to_dict(result)
 


Test patch:
diff --git a/tests/providers/google/cloud/hooks/test_dataproc.py b/tests/providers/google/cloud/hooks/test_dataproc.py
--- a/tests/providers/google/cloud/hooks/test_dataproc.py
+++ b/tests/providers/google/cloud/hooks/test_dataproc.py
@@ -23,6 +23,7 @@
 import pytest
 from google.api_core.gapic_v1.method import DEFAULT
 from google.cloud.dataproc_v1 import (
+    Batch,
     BatchControllerAsyncClient,
     ClusterControllerAsyncClient,
     JobControllerAsyncClient,
@@ -419,6 +420,28 @@ def test_create_batch(self, mock_client):
             timeout=None,
         )
 
+    @mock.patch(DATAPROC_STRING.format("DataprocHook.get_batch"))
+    def test_wait_for_batch(self, mock_batch):
+        mock_batch.return_value = Batch(state=Batch.State.SUCCEEDED)
+        result: Batch = self.hook.wait_for_batch(
+            batch_id=BATCH_ID,
+            region=GCP_LOCATION,
+            project_id=GCP_PROJECT,
+            wait_check_interval=1,
+            retry=DEFAULT,
+            timeout=None,
+            metadata=(),
+        )
+        mock_batch.assert_called_once_with(
+            batch_id=BATCH_ID,
+            region=GCP_LOCATION,
+            project_id=GCP_PROJECT,
+            retry=DEFAULT,
+            timeout=None,
+            metadata=(),
+        )
+        assert result.state == Batch.State.SUCCEEDED
+
     @mock.patch(DATAPROC_STRING.format("DataprocHook.get_batch_client"))
     def test_delete_batch(self, mock_client):
         self.hook.delete_batch(
diff --git a/tests/providers/google/cloud/operators/test_dataproc.py b/tests/providers/google/cloud/operators/test_dataproc.py
--- a/tests/providers/google/cloud/operators/test_dataproc.py
+++ b/tests/providers/google/cloud/operators/test_dataproc.py
@@ -1923,6 +1923,7 @@ def test_execute(self, mock_hook, to_dict_mock):
             timeout=TIMEOUT,
             metadata=METADATA,
         )
+        mock_hook.return_value.wait_for_operation.return_value = Batch(state=Batch.State.SUCCEEDED)
         op.execute(context=MagicMock())
         mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN)
         mock_hook.return_value.create_batch.assert_called_once_with(
@@ -1953,6 +1954,7 @@ def test_execute_with_result_retry(self, mock_hook, to_dict_mock):
             timeout=TIMEOUT,
             metadata=METADATA,
         )
+        mock_hook.return_value.wait_for_operation.return_value = Batch(state=Batch.State.SUCCEEDED)
         op.execute(context=MagicMock())
         mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN)
         mock_hook.return_value.create_batch.assert_called_once_with(
@@ -1982,14 +1984,91 @@ def test_execute_batch_failed(self, mock_hook, to_dict_mock):
             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
+        mock_hook.return_value.wait_for_operation.return_value = Batch(state=Batch.State.FAILED)
         with pytest.raises(AirflowException):
             op.execute(context=MagicMock())
-            mock_hook.return_value.get_batch.assert_called_once_with(
+
+    @mock.patch(DATAPROC_PATH.format("DataprocHook"))
+    def test_execute_batch_already_exists_succeeds(self, mock_hook):
+        op = DataprocCreateBatchOperator(
+            task_id=TASK_ID,
+            gcp_conn_id=GCP_CONN_ID,
+            impersonation_chain=IMPERSONATION_CHAIN,
+            region=GCP_REGION,
+            project_id=GCP_PROJECT,
+            batch=BATCH,
+            batch_id=BATCH_ID,
+            request_id=REQUEST_ID,
+            retry=RETRY,
+            timeout=TIMEOUT,
+            metadata=METADATA,
+        )
+        mock_hook.return_value.wait_for_operation.side_effect = AlreadyExists("")
+        mock_hook.return_value.wait_for_batch.return_value = Batch(state=Batch.State.SUCCEEDED)
+        op.execute(context=MagicMock())
+        mock_hook.return_value.wait_for_batch.assert_called_once_with(
+            batch_id=BATCH_ID,
+            region=GCP_REGION,
+            project_id=GCP_PROJECT,
+            wait_check_interval=5,
+            retry=RETRY,
+            timeout=TIMEOUT,
+            metadata=METADATA,
+        )
+
+    @mock.patch(DATAPROC_PATH.format("DataprocHook"))
+    def test_execute_batch_already_exists_fails(self, mock_hook):
+        op = DataprocCreateBatchOperator(
+            task_id=TASK_ID,
+            gcp_conn_id=GCP_CONN_ID,
+            impersonation_chain=IMPERSONATION_CHAIN,
+            region=GCP_REGION,
+            project_id=GCP_PROJECT,
+            batch=BATCH,
+            batch_id=BATCH_ID,
+            request_id=REQUEST_ID,
+            retry=RETRY,
+            timeout=TIMEOUT,
+            metadata=METADATA,
+        )
+        mock_hook.return_value.wait_for_operation.side_effect = AlreadyExists("")
+        mock_hook.return_value.wait_for_batch.return_value = Batch(state=Batch.State.FAILED)
+        with pytest.raises(AirflowException):
+            op.execute(context=MagicMock())
+            mock_hook.return_value.wait_for_batch.assert_called_once_with(
+                batch_id=BATCH_ID,
+                region=GCP_REGION,
+                project_id=GCP_PROJECT,
+                wait_check_interval=10,
+                retry=RETRY,
+                timeout=TIMEOUT,
+                metadata=METADATA,
+            )
+
+    @mock.patch(DATAPROC_PATH.format("DataprocHook"))
+    def test_execute_batch_already_exists_cancelled(self, mock_hook):
+        op = DataprocCreateBatchOperator(
+            task_id=TASK_ID,
+            gcp_conn_id=GCP_CONN_ID,
+            impersonation_chain=IMPERSONATION_CHAIN,
+            region=GCP_REGION,
+            project_id=GCP_PROJECT,
+            batch=BATCH,
+            batch_id=BATCH_ID,
+            request_id=REQUEST_ID,
+            retry=RETRY,
+            timeout=TIMEOUT,
+            metadata=METADATA,
+        )
+        mock_hook.return_value.wait_for_operation.side_effect = AlreadyExists("")
+        mock_hook.return_value.wait_for_batch.return_value = Batch(state=Batch.State.CANCELLED)
+        with pytest.raises(AirflowException):
+            op.execute(context=MagicMock())
+            mock_hook.return_value.wait_for_batch.assert_called_once_with(
                 batch_id=BATCH_ID,
                 region=GCP_REGION,
                 project_id=GCP_PROJECT,
+                wait_check_interval=10,
                 retry=RETRY,
                 timeout=TIMEOUT,
                 metadata=METADATA,