`BigQueryInsertJobOperator` doesn't handle job errors caused by running tasks in deferrable mode

[nathadfield]

Apache Airflow version

2.7.0

What happened

When a BigQuery job running in deferrable ends before completion - i.e. cancelled via the UI or terminated due a lack of resources - the precise error is not returned to the logs. Instead a key error from within the execute_complete function is returned.

[2023-09-07, 09:17:52 UTC] {taskinstance.py:1824} ERROR - Task failed with exception
Traceback (most recent call last):
  File "/usr/local/lib/python3.10/site-packages/airflow/providers/google/cloud/operators/bigquery.py", line 2874, in execute_complete
    raise AirflowException(event["message"])
KeyError: 'message'

https://github.com/apache/airflow/blob/main/airflow/providers/google/cloud/operators/bigquery.py#L2874-L2887

What you think should happen instead

When not running in deferrable mode, when the job is cancelled/fails during execution, the _handle_job_error is called resulting in the following output.

[2023-09-07, 09:16:35 UTC] {taskinstance.py:1824} ERROR - Task failed with exception
Traceback (most recent call last):
  File "/usr/local/lib/python3.10/site-packages/airflow/providers/google/cloud/operators/bigquery.py", line 2848, in execute
    job.result(timeout=self.result_timeout, retry=self.result_retry)
  File "/usr/local/lib/python3.10/site-packages/google/cloud/bigquery/job/query.py", line 1499, in result
    do_get_result()
  File "/usr/local/lib/python3.10/site-packages/google/cloud/bigquery/job/query.py", line 1489, in do_get_result
    super(QueryJob, self).result(retry=retry, timeout=timeout)
  File "/usr/local/lib/python3.10/site-packages/google/cloud/bigquery/job/base.py", line 728, in result
    return super(_AsyncJob, self).result(timeout=timeout, **kwargs)
  File "/usr/local/lib/python3.10/site-packages/google/api_core/future/polling.py", line 261, in result
    raise self._exception
google.api_core.exceptions.GoogleAPICallError: 200 Job execution was cancelled: User requested cancellation

It would appear that, even though _handle_job_error should get called when in deferrable mode, this doesn't happen if the job was running when it terminates.

https://github.com/apache/airflow/blob/main/airflow/providers/google/cloud/operators/bigquery.py#L2860-L2872

How to reproduce

The code below will replicate the issue, however you need to provide a query that runs for enough time that enables you to cancel the job in UI.

Queries that have syntax errors never get started and are not affected by this.

from datetime import datetime

from airflow import models
from airflow.providers.google.cloud.operators.bigquery import BigQueryInsertJobOperator

with models.DAG(
    dag_id='bq_insert_job',
    start_date=datetime(2023, 8, 31),
    catchup=False,
    schedule='0 0 * * *',
) as dag:

    test1 = BigQueryInsertJobOperator(
        task_id=f'test1',
        sql='<INSERT LONG RUNNING QUERY>',
        deferrable=False,
    )

    test2 = BigQueryInsertJobOperator(
        task_id=f'test2',
        sql=f'<INSERT LONG RUNNING QUERY>',
        deferrable=True,
    )

Operating System

n/a

Versions of Apache Airflow Providers

apache-airflow-providers-google==10.6.0

Deployment

Astronomer

Deployment details

No response

Anything else

No response

Are you willing to submit PR?

• Yes I am willing to submit a PR!

Code of Conduct

• I agree to follow this project's Code of Conduct

[nathadfield]

The issue is that the job trigger only returns TriggerEvent({"status": "error"}) and execute_complete
https://github.com/apache/airflow/blob/main/airflow/providers/google/cloud/triggers/bigquery.py#L91

I was thinking that I could add another function in BigQueryAsyncHook to return the exact error message from the job status, similar to get_job_status.

https://github.com/apache/airflow/blob/main/airflow/providers/google/cloud/hooks/bigquery.py#L3102C15-L3111

e.g.

    async def get_job_message(self, job_id: str | None, project_id: str | None = None) -> str:
        async with ClientSession() as s:
            job_client = await self.get_job_instance(project_id, job_id, s)
            job = await job_client.get_job()
            self.log.info(job)
            status = job.get("status", {})
            if status["state"] == "DONE":
                if "errorResult" in status:
                    return status["errorResult"]['message']

And also modify the trigger function call this when job_status is error such that we can yield TriggerEvent({"status": "error", "message": message}).

[nathadfield]

Actually, I think I can just modify what get_job_status returns and handle that in the trigger.