Implement BigQueryStreamingBufferEmptySensor to handle DML operations on streaming tables

[radhwene]

Implement BigQueryStreamingBufferEmptySensor to handle DML operations on streaming tables

Fixes #59408

Problem

When using BigQuery DML operators (UPDATE, DELETE, MERGE) on tables with active streaming buffers, tasks fail with:
This is a documented BigQuery limitation. Currently, Airflow has no built-in mechanism to wait for the buffer to flush before executing DML operations, causing repeated failures until it eventually clears (within 90 minutes per Google Cloud documentation).

Solution

This PR implements BigQueryStreamingBufferEmptySensor - a composable sensor that allows users to explicitly wait for a BigQuery table's streaming buffer to empty before proceeding with DML operations.

This aligns with Airflow's design philosophy by providing:

• Explicit control over pipeline dependencies
• Composable operations that can be combined with other tasks
• Non-blocking execution via deferrable mode

Changes

1. New Sensor: BigQueryStreamingBufferEmptySensor

• Polls BigQuery table metadata to check if streaming buffer is empty
• Supports both synchronous polling and deferrable (async) execution
• 90-minute timeout aligned with Google Cloud's streaming buffer flush guarantee
• Full support for GCP service account impersonation
• Consistent with existing BigQueryTableExistenceSensor implementation

2. New Trigger: BigQueryStreamingBufferEmptyTrigger

• Async trigger for deferrable sensor execution
• Continuous polling of table metadata via BigQueryTableAsyncHook
• Event-driven callback integration
• Comprehensive error handling

3. Documentation & Examples

• Added complete documentation section in bigquery.rst
• 3 usage examples (sync, deferrable, async modes)
• Real-world workflow examples with streaming INSERT/UPDATE/DELETE operations
• 90-minute timeout configuration with explanations

Usage Example

from airflow import DAG
from airflow.providers.google.cloud.sensors.bigquery import BigQueryStreamingBufferEmptySensor
from airflow.providers.google.cloud.operators.bigquery import BigQueryInsertJobOperator

with DAG('bigquery_dml_pipeline'):
    # Wait for streaming buffer to be empty
    wait_buffer = BigQueryStreamingBufferEmptySensor(
        task_id='wait_buffer_empty',
        project_id='my-project',
        dataset_id='my_dataset',
        table_id='my_table',
        deferrable=True,  # Non-blocking execution
        timeout=5400,  # 90 minutes
    )
    
    # Then safely execute DML operation
    update_table = BigQueryInsertJobOperator(
        task_id='update_table',
        configuration={
            'query': {
                'query': 'UPDATE my_dataset.my_table SET status="processed"',
                'useLegacySql': False,
            }
        },
    )
    
    wait_buffer >> update_table

[boring-cyborg]

Congratulations on your first Pull Request and welcome to the Apache Airflow community! If you have any issues or are unsure about any anything please check our Contributors' Guide (https://github.com/apache/airflow/blob/main/contributing-docs/README.rst)
Here are some useful points:

• Pay attention to the quality of your code (ruff, mypy and type annotations). Our prek-hooks will help you with that.
• In case of a new feature add useful documentation (in docstrings or in docs/ directory). Adding a new operator? Check this short guide Consider adding an example DAG that shows how users should use it.
• Consider using Breeze environment for testing locally, it's a heavy docker but it ships with a working Airflow and a lot of integrations.
• Be patient and persistent. It might take some time to get a review or get the final approval from Committers.
• Please follow ASF Code of Conduct for all communication including (but not limited to) comments on Pull Requests, Mailing list and Slack.
• Be sure to read the Airflow Coding style.
• Always keep your Pull Requests rebased, otherwise your build might fail due to changes not related to your commits.
  Apache Airflow is a community-driven project and together we are making it better 🚀.
  In case of doubts contact the developers at:
  Mailing List: dev@airflow.apache.org
  Slack: https://s.apache.org/airflow-slack

[shahar1]

Implement BigQueryStreamingBufferEmptySensor to handle DML operations on streaming tables

Fixes #59408

Problem

When using BigQuery DML operators (UPDATE, DELETE, MERGE) on tables with active streaming buffers, tasks fail with: This is a documented BigQuery limitation. Currently, Airflow has no built-in mechanism to wait for the buffer to flush before executing DML operations, causing repeated failures until it eventually clears (within 90 minutes per Google Cloud documentation).

Solution

This PR implements BigQueryStreamingBufferEmptySensor - a composable sensor that allows users to explicitly wait for a BigQuery table's streaming buffer to empty before proceeding with DML operations.

This aligns with Airflow's design philosophy by providing:

• Explicit control over pipeline dependencies
• Composable operations that can be combined with other tasks
• Non-blocking execution via deferrable mode

Changes

1. New Sensor: BigQueryStreamingBufferEmptySensor

• Polls BigQuery table metadata to check if streaming buffer is empty
• Supports both synchronous polling and deferrable (async) execution
• 90-minute timeout aligned with Google Cloud's streaming buffer flush guarantee
• Full support for GCP service account impersonation
• Consistent with existing BigQueryTableExistenceSensor implementation

2. New Trigger: BigQueryStreamingBufferEmptyTrigger

• Async trigger for deferrable sensor execution
• Continuous polling of table metadata via BigQueryTableAsyncHook
• Event-driven callback integration
• Comprehensive error handling

3. Documentation & Examples

• Added complete documentation section in bigquery.rst
• 3 usage examples (sync, deferrable, async modes)
• Real-world workflow examples with streaming INSERT/UPDATE/DELETE operations
• 90-minute timeout configuration with explanations

Usage Example

from airflow import DAG
from airflow.providers.google.cloud.sensors.bigquery import BigQueryStreamingBufferEmptySensor
from airflow.providers.google.cloud.operators.bigquery import BigQueryInsertJobOperator

with DAG('bigquery_dml_pipeline'):
    # Wait for streaming buffer to be empty
    wait_buffer = BigQueryStreamingBufferEmptySensor(
        task_id='wait_buffer_empty',
        project_id='my-project',
        dataset_id='my_dataset',
        table_id='my_table',
        deferrable=True,  # Non-blocking execution
        timeout=5400,  # 90 minutes
    )
    
    # Then safely execute DML operation
    update_table = BigQueryInsertJobOperator(
        task_id='update_table',
        configuration={
            'query': {
                'query': 'UPDATE my_dataset.my_table SET status="processed"',
                'useLegacySql': False,
            }
        },
    )
    
    wait_buffer >> update_table

Thanks for your contribution!
Could you please attach a screenshot of a the system tests after running?

CC: @VladaZakharova @MaksYermak

[radhwene]

System Test Results

I ran a standalone test script against a live GCP project to validate the BigQueryStreamingBufferEmptySensor behavior end-to-end.

Test Flow

1. Create dataset & table in BigQuery
2. Insert rows via Streaming API (insert_rows_json)
3. Confirm streaming buffer is active (~3 rows, ~400 bytes)
4. Sensor poke() polls every 30s until buffer is flushed
5. Run DML UPDATE to confirm buffer is clear
6. Cleanup (delete dataset)

Screenshot

Tested on:

• Airflow 2.10.3 / Python 3.12
• google-cloud-bigquery SDK
• GCP project with BigQuery API enabled

[shahar1]

System Test Results

I ran a standalone test script against a live GCP project to validate the BigQueryStreamingBufferEmptySensor behavior end-to-end.

Test Flow

1. Create dataset & table in BigQuery
2. Insert rows via Streaming API (insert_rows_json)
3. Confirm streaming buffer is active (~3 rows, ~400 bytes)
4. Sensor poke() polls every 30s until buffer is flushed
5. Run DML UPDATE to confirm buffer is clear
6. Cleanup (delete dataset)

Screenshot

### Tested on: * Airflow 2.10.3 / Python 3.12 * google-cloud-bigquery SDK * GCP project with BigQuery API enabled

Perfect!

System Test Results

I ran a standalone test script against a live GCP project to validate the BigQueryStreamingBufferEmptySensor behavior end-to-end.

Test Flow

1. Create dataset & table in BigQuery
2. Insert rows via Streaming API (insert_rows_json)
3. Confirm streaming buffer is active (~3 rows, ~400 bytes)
4. Sensor poke() polls every 30s until buffer is flushed
5. Run DML UPDATE to confirm buffer is clear
6. Cleanup (delete dataset)

Screenshot

### Tested on: * Airflow 2.10.3 / Python 3.12 * google-cloud-bigquery SDK * GCP project with BigQuery API enabled

Looks great! Failing CI doesn't seem related, I'll merge from main and rerun.

[shahar1]

Almost there! I wanted to merge - but I've just realized that:

1. There are no unit tests for the new operators - could you please add for each, including checking exceptions? It's essential as a fast indicator for regressions. We currently cannot run the system tests in the automated CI, and even if we could - it would be better to detect regressions in earlier stage.
2. We try to reduce usage of AirflowException overall - could you please change the exceptions to Python built-in exceptions? (e.g., ValueError, etc.)

If you could please take care of the above, I'll happily approve and merge.

Asked to add unit tests + changing AirflowException to native Python exceptions

[radhwene]

Hi @shahar1 ,
Thanks for the review! I've added the unit tests for the exception handling cases.

Sensor tests (test_bigquery.py)

Test	Description
test_poke_table_not_found	Raises AirflowException when the table doesn't exist
test_poke_raises_on_unexpected_error	Re-raises unexpected exceptions from the BigQuery client
test_execute_complete_no_event	Raises AirflowException when no event is received in the trigger callback
test_execute_complete_error_event	Raises AirflowException when trigger returns an error status

Trigger tests (test_bigquery.py)

Test	Description
test_run_raises_on_table_not_found	Yields error event when table returns 404
test_run_raises_on_exception	Yields error event on unexpected exceptions
test_is_streaming_buffer_empty_table_not_exists	Raises AirflowException when table response is empty