# Buggy code
def wait_for_pipeline_state(
    self,
    pipeline_name: str,
    pipeline_id: str,
    instance_url: str,
    namespace: str = "default",
    success_states: Optional[List[str]] = None,
    failure_states: Optional[List[str]] = None,
    timeout: int = 5 * 60,
):
    """
    Polls pipeline state and raises an exception if the state is one of
    `failure_states` or the operation timeouted.
    """
    failure_states = failure_states or FAILURE_STATES
    success_states = success_states or SUCCESS_STATES
    start_time = monotonic()
    current_state = None
    while monotonic() - start_time < timeout:
        current_state = self._get_workflow_state(
            pipeline_name=pipeline_name,
            pipeline_id=pipeline_id,
            instance_url=instance_url,
            namespace=namespace,
        )

        if current_state in success_states:
            return
        if current_state in failure_states:
            raise AirflowException(
                f"Pipeline {pipeline_name} state {current_state} is not "
                f"one of {success_states}"
            )
        sleep(30)

    # Time is up!
    raise AirflowException(
        f"Pipeline {pipeline_name} state {current_state} is not "
        f"one of {success_states} after {timeout}s"
    )


# Traceback (airflow.exceptions.AirflowException)
[2020-07-28 08:16:46,390] {taskinstance.py:1059} ERROR - Retrieving a pipeline state failed with code 404
Traceback (most recent call last)
File "/usr/local/lib/airflow/airflow/models/taskinstance.py", line 930, in _run_raw_tas
result = task_copy.execute(context=context
File "/usr/local/lib/airflow/airflow/providers/google/cloud/operators/datafusion.py", line 700, in execut
runtime_args=self.runtime_args
File "/usr/local/lib/airflow/airflow/providers/google/cloud/hooks/datafusion.py", line 498, in start_pipelin
instance_url=instance_url
File "/usr/local/lib/airflow/airflow/providers/google/cloud/hooks/datafusion.py", line 114, in wait_for_pipeline_stat
namespace=namespace
File "/usr/local/lib/airflow/airflow/providers/google/cloud/hooks/datafusion.py", line 442, in _get_workflow_stat
f"Retrieving a pipeline state failed with code {response.status}
airflow.exceptions.AirflowException: Retrieving a pipeline state failed with code 404

# Fixed code
def wait_for_pipeline_state(
    self,
    pipeline_name: str,
    pipeline_id: str,
    instance_url: str,
    namespace: str = "default",
    success_states: Optional[List[str]] = None,
    failure_states: Optional[List[str]] = None,
    timeout: int = 5 * 60,
):
    """
    Polls pipeline state and raises an exception if the state is one of
    `failure_states` or the operation timeouted.
    """
    failure_states = failure_states or FAILURE_STATES
    success_states = success_states or SUCCESS_STATES
    start_time = monotonic()
    current_state = None
    while monotonic() - start_time < timeout:
        try:
            current_state = self._get_workflow_state(
                pipeline_name=pipeline_name,
                pipeline_id=pipeline_id,
                instance_url=instance_url,
                namespace=namespace,
            )
        except AirflowException:
            pass  # Because the pipeline may not be visible in system yet
        if current_state in success_states:
            return
        if current_state in failure_states:
            raise AirflowException(
                f"Pipeline {pipeline_name} state {current_state} is not "
                f"one of {success_states}"
            )
        sleep(30)

    # Time is up!
    raise AirflowException(
        f"Pipeline {pipeline_name} state {current_state} is not "
        f"one of {success_states} after {timeout}s"
    )