BigQueryInsertJobOperator is broken on any type of job except `query`

[vaaalik]

Apache Airflow Provider(s)

google

Versions of Apache Airflow Providers

apache-airflow-providers-google==7.0.0

Apache Airflow version

2.2.5

Operating System

MacOS 12.2.1

Deployment

Official Apache Airflow Helm Chart

Deployment details

No response

What happened

We are using BigQueryInsertJobOperator to load data from parquet files in Google Cloud Storage with this kind of configuration:

BigQueryInsertJobOperator(
        task_id="load_to_bq",
        configuration={
            "load": {
                "writeDisposition": "WRITE_APPEND",
                "createDisposition": "CREATE_IF_NEEDED",
                "destinationTable": destination_table,
                "sourceUris": source_files
                "sourceFormat": "PARQUET"
            }
        }

After upgrade to apache-airflow-providers-google==7.0.0 all load jobs are now broken. I believe that problem lies in this line:

airflow/airflow/providers/google/cloud/operators/bigquery.py

Line 2170 in 5bfacf8

table = job.to_api_repr()["configuration"]["query"]["destinationTable"]

So it's trying to get the destination table from query job config and makes it impossible to use any other type of job.

What you think should happen instead

No response

How to reproduce

Use BigQueryInsertJobOperator to submit any type of job except query

Anything else

Traceback (most recent call last):
  File "/home/airflow/.local/lib/python3.9/site-packages/airflow/providers/google/cloud/operators/bigquery.py", line 2170, in execute
    table = job.to_api_repr()["configuration"]["query"]["destinationTable"]
KeyError: 'query'

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

[boring-cyborg]

Thanks for opening your first issue here! Be sure to follow the issue template!

[raphaelauv]

#23165 @wojsamjan

[potiuk]

Do you think #23165 fixes this issue ? Can we close it @raphaelauv ?

[raphaelauv]

No It's the reason of the issue

[potiuk]

Ah OK. That was not clear @wojsamjan -> can you please take a look and fix it ?

[kardoo]

I have exactly the issue.
I moved all my queries from operator BigQueryExecuteQueryOperator to BigQueryInsertJobOperator because it will be deprecated and now all queries no longer work.
What could be the best approach? Wait for a fix? Make a custom operator? I cannot find a workarround.
Thanks

[raphaelauv]

try this

from airflow.providers.google.cloud.operators.bigquery import BigQueryInsertJobOperator
from typing import Any
from google.api_core.exceptions import Conflict
from airflow.exceptions import AirflowException
from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook
from airflow.providers.google.cloud.links.bigquery import BigQueryTableLink


class FixBigQueryInsertJobOperator(BigQueryInsertJobOperator):

    def execute(self, context: Any):
        hook = BigQueryHook(
            gcp_conn_id=self.gcp_conn_id,
            delegate_to=self.delegate_to,
            impersonation_chain=self.impersonation_chain,
        )
        self.hook = hook
        job_id = self._job_id(context)
        try:
            job = self._submit_job(hook, job_id)
            self._handle_job_error(job)
        except Conflict:
            # If the job already exists retrieve it
            job = hook.get_job(
                project_id=self.project_id,
                location=self.location,
                job_id=job_id,
            )
            if job.state in self.reattach_states:
                # We are reattaching to a job
                job.result(timeout=self.result_timeout, retry=self.result_retry)
                self._handle_job_error(job)
            else:
                # Same job configuration so we need force_rerun
                raise AirflowException(
                    f"Job with id: {job_id} already exists and is in {job.state} state. If you "
                    f"want to force rerun it consider setting `force_rerun=True`."
                    f"Or, if you want to reattach in this scenario add {job.state} to `reattach_states`"
                )

        if "query" in job.to_api_repr()["configuration"]:
            if "destinationTable" in job.to_api_repr()["configuration"]["query"]:
                table = job.to_api_repr()["configuration"]["query"]["destinationTable"]
                BigQueryTableLink.persist(
                    context=context,
                    task_instance=self,
                    dataset_id=table["datasetId"],
                    project_id=table["projectId"],
                    table_id=table["tableId"],
                )
        self.job_id = job.job_id
        return job.job_id

[kardoo]

Although I was avoiding this solution because I would have to refactor hundreds of queries, I guess its the best solution.
I wonder, if a fix arrives quickly enough so I could avoid this work. :)
Thank you so much for sharing.

[eladkal]

This is open source project - the fix will arrive when someone decides to spend the time on fixing it.
If you volunteer I'd be happy assign this issue to you and assist with code review.

[kardoo]

My apologies if I was misunderstood. If I have time in the next couple of weeks, I will let you know.

[takuma11248250]

hi.
I have exactly the issue.
When do you expect this issue to be resolved?