Repository: apache/airflow
File: providers/src/airflow/providers/google/cloud/operators/bigquery.py (Python)

Code before:
                self.configuration["labels"] = automatic_labels

    def _submit_job(
        self,
        hook: BigQueryHook,
        job_id: str,
    ) -> BigQueryJob:
        # Annotate the job with dag and task id labels
        self._add_job_labels()

        # Submit a new job without waiting for it to complete.
        return hook.insert_job(
            configuration=self.configuration,
            project_id=self.project_id,
            location=self.location,
            job_id=job_id,
            timeout=self.result_timeout,
            retry=self.result_retry,
            nowait=True,
        )

    def _handle_job_error(self, job: BigQueryJob | UnknownJob) -> None:
        self.log.info("Job %s is completed. Checking the job status", self.job_id)
        # I've noticed that sometimes BigQuery jobs transiently report the wrong status, causing the Airflow job to be incorrectly marked as successful.
        # To avoid this, we refresh the job properties before checking the final state and handling any errors.
        while job.state != "DONE":
            job.reload(timeout=self.result_timeout, retry=self.result_retry)
        # Log any transient errors encountered during the job execution
        for error in job.errors or []:
            self.log.error("BigQuery Job Error:", error)
        if job.error_result:
            raise AirflowException(f"BigQuery job {job.job_id} failed: {job.error_result}")

    def execute(self, context: Any):
        hook = BigQueryHook(
            gcp_conn_id=self.gcp_conn_id,
            impersonation_chain=self.impersonation_chain,
        )
        self.hook = hook
        if self.project_id is None:
            self.project_id = hook.project_id

        self.job_id = hook.generate_job_id(
            job_id=self.job_id,
            dag_id=self.dag_id,
            task_id=self.task_id,
            logical_date=context["logical_date"],
            configuration=self.configuration,
            force_rerun=self.force_rerun,
        )


Review comment (bug):
I don't follow why we should do this.
Is there a bug report to BigQuery that we can track?

I am very against such permanent workarounds. Airflow should not be the fixer for upstream library bugs. Its OK to do this a temporary workaround when upstream library is engaged and working on a fix but otherwise not.


Code after:
            elif "labels" not in self.configuration:
                self.configuration["labels"] = automatic_labels

    def _submit_job(
        self,
        hook: BigQueryHook,
        job_id: str,
    ) -> BigQueryJob:
        # Annotate the job with dag and task id labels
        self._add_job_labels()

        # Submit a new job without waiting for it to complete.
        return hook.insert_job(
            configuration=self.configuration,
            project_id=self.project_id,
            location=self.location,
            job_id=job_id,
            timeout=self.result_timeout,
            retry=self.result_retry,
            nowait=True,
        )

    def _handle_job_error(self, job: BigQueryJob | UnknownJob) -> None:
        self.log.info("Job %s is completed. Checking the job status", self.job_id)
        # Log any transient errors encountered during the job execution
        for error in job.errors or []:
            self.log.error("BigQuery Job Error: %s", error)
        if job.error_result:
            raise AirflowException(f"BigQuery job {job.job_id} failed: {job.error_result}")
        # Check the final state.
        if job.state != "DONE":
            raise AirflowException(f"Job failed with state: {job.state}")

    def execute(self, context: Any):
        hook = BigQueryHook(
            gcp_conn_id=self.gcp_conn_id,
            impersonation_chain=self.impersonation_chain,
        )
        self.hook = hook
        if self.project_id is None:
            self.project_id = hook.project_id

        self.job_id = hook.generate_job_id(
            job_id=self.job_id,
            dag_id=self.dag_id,
            task_id=self.task_id,
            logical_date=context["logical_date"],
            configuration=self.configuration,
            force_rerun=self.force_rerun,
        )