Change Airflow BigQueryInsertJobOperator and BigQueryGetDataOperator Priority to Batch

Question

I've read the Apache Airflow documentation for operators for BigQuery jobs here (https://airflow.apache.org/docs/apache-airflow-providers-google/stable/_modules/airflow/providers/google/cloud/operators/bigquery.html#BigQueryGetDataOperator) and I can't find how to change the job priority to batch. How is it can be done?

Elad Elad · Accepted Answer · 2021-06-23T07:19:51

BigQueryExecuteQueryOperator has priority param that can be set with INTERACTIVE/BATCH the default is INTERACTIVE:

execute_insert_query = BigQueryExecuteQueryOperator(
    task_id="execute_insert_query",
    sql=INSERT_ROWS_QUERY,
    use_legacy_sql=False,
    location=location,
    priority='BATCH',
)

The BigQueryInsertJobOperator doesn't have it. I think you can create a custom operator that inherits from BigQueryInsertJobOperator and adds it by overwriting the _submit_job function:

class MyBigQueryInsertJobOperator(BigQueryInsertJobOperator):
    def __init__(
        self,
        priority: str = 'INTERACTIVE',
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)
        self.priority = priority

    def _submit_job(
        self,
        hook: BigQueryHook,
        job_id: str,
    ) -> BigQueryJob:
        # Submit a new job
        job = hook.insert_job(
            configuration=self.configuration,
            project_id=self.project_id,
            location=self.location,
            job_id=job_id,
            priority=self.priority,
        )
        # Start the job and wait for it to complete and get the result.
        job.result()
        return job

I didn't test though but it should work.

Change Airflow BigQueryInsertJobOperator and BigQueryGetDataOperator Priority to Batch

1 Answers