You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "radhwene (via GitHub)" <gi...@apache.org> on 2026/01/27 20:18:02 UTC
[PR] Implement BigQueryStreamingBufferEmptySensor to handle DML operations on streaming tables [airflow]
radhwene opened a new pull request, #61148:
URL: https://github.com/apache/airflow/pull/61148
# 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
```python
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
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement BigQueryStreamingBufferEmptySensor to handle DML operations on streaming tables [airflow]
Posted by "shahar1 (via GitHub)" <gi...@apache.org>.
shahar1 commented on PR #61148:
URL: https://github.com/apache/airflow/pull/61148#issuecomment-3876199675
> ## 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
> <img alt="image" width="704" height="300" src="https://private-user-images.githubusercontent.com/20400181/547578269-88360e20-2c21-4398-93db-692eab42d159.png?jwt=eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3NzA3MTMyMTMsIm5iZiI6MTc3MDcxMjkxMywicGF0aCI6Ii8yMDQwMDE4MS81NDc1NzgyNjktODgzNjBlMjAtMmMyMS00Mzk4LTkzZGItNjkyZWFiNDJkMTU5LnBuZz9YLUFtei1BbGdvcml0aG09QVdTNC1ITUFDLVNIQTI1NiZYLUFtei1DcmVkZW50aWFsPUFLSUFWQ09EWUxTQTUzUFFLNFpBJTJGMjAyNjAyMTAlMkZ1cy1lYXN0LTElMkZzMyUyRmF3czRfcmVxdWVzdCZYLUFtei1EYXRlPTIwMjYwMjEwVDA4NDE1M1omWC1BbXotRXhwaXJlcz0zMDAmWC1BbXotU2lnbmF0dXJlPWFmYmI5MWIyYWVkMmE1YjFmNDM2OTIyZTYwMzM0ODY4MWEzYWQ0YjliODdjN2Q4Y2EzMTdhOGU3N2JjYmEzZjgmWC1BbXotU2lnbmVkSGVhZGVycz1ob3N0In0.ff9PTo_vZc4N5EoHNmmfvWcJu58budud4FxzmP19T4o">
> ### 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
> <img alt="image" width="704" height="300" src="https://private-user-images.githubusercontent.com/20400181/547578269-88360e20-2c21-4398-93db-692eab42d159.png?jwt=eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3NzA3MTMyMzQsIm5iZiI6MTc3MDcxMjkzNCwicGF0aCI6Ii8yMDQwMDE4MS81NDc1NzgyNjktODgzNjBlMjAtMmMyMS00Mzk4LTkzZGItNjkyZWFiNDJkMTU5LnBuZz9YLUFtei1BbGdvcml0aG09QVdTNC1ITUFDLVNIQTI1NiZYLUFtei1DcmVkZW50aWFsPUFLSUFWQ09EWUxTQTUzUFFLNFpBJTJGMjAyNjAyMTAlMkZ1cy1lYXN0LTElMkZzMyUyRmF3czRfcmVxdWVzdCZYLUFtei1EYXRlPTIwMjYwMjEwVDA4NDIxNFomWC1BbXotRXhwaXJlcz0zMDAmWC1BbXotU2lnbmF0dXJlPTM4NDNjZmE5MGM5MzI0MzMxNmEyODc4MDNmMDYzZGZiMTY5NGU1YWJiNzFjMzBiYWQ4ZGEwNDg1ODQxZTljMTAmWC1BbXotU2lnbmVkSGVhZGVycz1ob3N0In0.lr7C7apPW1CBJCCD5qkkBEEaXndsO82jvw9-pFVi7_M">
> ### 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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement BigQueryStreamingBufferEmptySensor to handle DML operations on streaming tables [airflow]
Posted by "radhwene (via GitHub)" <gi...@apache.org>.
radhwene commented on code in PR #61148:
URL: https://github.com/apache/airflow/pull/61148#discussion_r2868865141
##########
providers/google/src/airflow/providers/google/cloud/sensors/bigquery.py:
##########
@@ -77,17 +76,7 @@ def __init__(
**kwargs,
) -> None:
if deferrable and "poke_interval" not in kwargs:
- # TODO: Remove once deprecated
- if "polling_interval" in kwargs:
- kwargs["poke_interval"] = kwargs["polling_interval"]
- warnings.warn(
- "Argument `poll_interval` is deprecated and will be removed "
- "in a future release. Please use `poke_interval` instead.",
- AirflowProviderDeprecationWarning,
- stacklevel=2,
- )
- else:
- kwargs["poke_interval"] = 5
Review Comment:
You’re right, thanks for pointing this out. I restored the deprecated code path for the current operator to keep it aligned with the Google provider deprecation policy.
Could you please revalidate the PR when you have a moment?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement BigQueryStreamingBufferEmptySensor to handle DML operations on streaming tables [airflow]
Posted by "MaksYermak (via GitHub)" <gi...@apache.org>.
MaksYermak commented on code in PR #61148:
URL: https://github.com/apache/airflow/pull/61148#discussion_r2815604019
##########
providers/google/tests/system/google/cloud/bigquery/example_bigquery_sensors.py:
##########
@@ -152,22 +162,105 @@
)
# [END howto_sensor_bigquery_table_partition_async]
+ # [START howto_sensor_bigquery_streaming_buffer_empty]
+ check_streaming_buffer_empty = BigQueryStreamingBufferEmptySensor(
+ task_id="check_streaming_buffer_empty",
+ project_id=PROJECT_ID,
+ dataset_id=DATASET_NAME,
+ table_id=TABLE_NAME,
+ poke_interval=30,
+ timeout=5400, # 90 minutes - Google Cloud flushes streaming buffer within 90 minutes
+ )
+ # [END howto_sensor_bigquery_streaming_buffer_empty]
+
+ # Streaming operations: INSERT, UPDATE, DELETE
+ # These operations write data to the streaming buffer before being flushed to persistent storage
+ stream_insert = BigQueryInsertJobOperator(
+ task_id="stream_insert",
+ configuration={
+ "query": {
+ "query": STREAMING_INSERT_QUERY,
+ "useLegacySql": False,
+ }
+ },
+ )
+
+ stream_update = BigQueryInsertJobOperator(
+ task_id="stream_update",
+ configuration={
+ "query": {
+ "query": STREAMING_UPDATE_QUERY,
+ "useLegacySql": False,
+ }
+ },
+ )
+
+ stream_delete = BigQueryInsertJobOperator(
+ task_id="stream_delete",
+ configuration={
+ "query": {
+ "query": STREAMING_DELETE_QUERY,
+ "useLegacySql": False,
+ }
+ },
+ )
+
+ # [START howto_sensor_bigquery_streaming_buffer_empty_defered]
+ check_streaming_buffer_empty_def = BigQueryStreamingBufferEmptySensor(
+ task_id="check_streaming_buffer_empty_def",
+ project_id=PROJECT_ID,
+ dataset_id=DATASET_NAME,
+ table_id=TABLE_NAME,
+ deferrable=True,
+ poke_interval=30,
+ timeout=5400, # 90 minutes - Google Cloud flushes streaming buffer within 90 minutes
+ )
+ # [END howto_sensor_bigquery_streaming_buffer_empty_defered]
+
+ # [START howto_sensor_bigquery_streaming_buffer_empty_async]
+ check_streaming_buffer_empty_async = BigQueryStreamingBufferEmptySensor(
+ task_id="check_streaming_buffer_empty_async",
+ project_id=PROJECT_ID,
+ dataset_id=DATASET_NAME,
+ table_id=TABLE_NAME,
+ poke_interval=30,
+ timeout=5400, # 90 minutes - Google Cloud flushes streaming buffer within 90 minutes
+ )
+ # [END howto_sensor_bigquery_streaming_buffer_empty_async]
Review Comment:
@radhwene could you explain the reason of this task? I do not see any `async` operation here it is just the same task as for sync operation, could you explain please?
##########
providers/google/src/airflow/providers/google/cloud/sensors/bigquery.py:
##########
@@ -256,3 +257,133 @@ def execute_complete(self, context: dict[str, Any], event: dict[str, str] | None
message = "No event received in trigger callback"
raise AirflowException(message)
+
+
+class BigQueryStreamingBufferEmptySensor(BaseSensorOperator):
+ """
+ Sensor for checking whether the streaming buffer in a BigQuery table is empty.
+
+ The BigQueryStreamingBufferEmptySensor waits for the streaming buffer in a specified
+ BigQuery table to be empty before proceeding. It can be used in ETL pipelines to ensure
+ that recent streamed data has been processed before continuing downstream tasks.
+
+ :param project_id: The Google Cloud project ID where the BigQuery table resides.
+ :param dataset_id: The ID of the dataset containing the BigQuery table.
+ :param table_id: The ID of the BigQuery table to monitor.
+ :param gcp_conn_id: The Airflow connection ID for GCP. Defaults to "google_cloud_default".
+ :param impersonation_chain: Optional service account to impersonate using short-term
+ credentials, or chained list of accounts required to get the access_token
+ of the last account in the list, which will be impersonated in the request.
+ If set as a string, the account must grant the originating account
+ the Service Account Token Creator IAM role.
+ If set as a sequence, the identities from the list must grant
+ Service Account Token Creator IAM role to the directly preceding identity, with first
+ account from the list granting this role to the originating account (templated).
+ :param deferrable: Run sensor in deferrable mode. Defaults to False.
+ """
+
+ template_fields: Sequence[str] = (
+ "project_id",
+ "dataset_id",
+ "table_id",
+ "impersonation_chain",
+ )
+
+ ui_color = "#f0eee4"
+
+ def __init__(
+ self,
+ *,
+ project_id: str,
+ dataset_id: str,
+ table_id: str,
+ gcp_conn_id: str = "google_cloud_default",
+ impersonation_chain: str | Sequence[str] | None = None,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ) -> None:
+ if deferrable and "poke_interval" not in kwargs:
+ # TODO: Remove once deprecated
+ if "polling_interval" in kwargs:
+ kwargs["poke_interval"] = kwargs["polling_interval"]
+ warnings.warn(
+ "Argument `polling_interval` is deprecated and will be removed "
+ "in a future release. Please use `poke_interval` instead.",
+ AirflowProviderDeprecationWarning,
+ stacklevel=2,
+ )
Review Comment:
@radhwene could you please remove this code, because it is deprecated we should not add a deprecate code to the new operator?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement BigQueryStreamingBufferEmptySensor to handle DML operations on streaming tables [airflow]
Posted by "shahar1 (via GitHub)" <gi...@apache.org>.
shahar1 commented on PR #61148:
URL: https://github.com/apache/airflow/pull/61148#issuecomment-3919204555
There's one failing test in the CI :(
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement BigQueryStreamingBufferEmptySensor to handle DML operations on streaming tables [airflow]
Posted by "boring-cyborg[bot] (via GitHub)" <gi...@apache.org>.
boring-cyborg[bot] commented on PR #61148:
URL: https://github.com/apache/airflow/pull/61148#issuecomment-3807341598
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]( https://github.com/apache/airflow/blob/main/contributing-docs/08_static_code_checks.rst#prerequisites-for-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](https://github.com/apache/airflow/blob/main/airflow-core/docs/howto/custom-operator.rst) Consider adding an example DAG that shows how users should use it.
- Consider using [Breeze environment](https://github.com/apache/airflow/blob/main/dev/breeze/doc/README.rst) 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](https://www.apache.org/foundation/policies/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]( https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#coding-style-and-best-practices).
- 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
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement BigQueryStreamingBufferEmptySensor to handle DML operations on streaming tables [airflow]
Posted by "radhwene (via GitHub)" <gi...@apache.org>.
radhwene commented on PR #61148:
URL: https://github.com/apache/airflow/pull/61148#issuecomment-4103396154
@potiuk Thanks for the report! The CI failures were caused by the branch being significantly behind main (360 commits). I've rebased the branch onto the latest main and squashed all commits into a single clean commit. The branch is now up to date.
The failing checks (Postgres, MySQL, Sqlite, MinSQLAlchemy, LatestBoto serialization tests) are unrelated to my changes — they were caused by drift from the base branch.
I'll wait for the new CI run to complete and mark the PR as "Ready for review" once all checks pass.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement BigQueryStreamingBufferEmptySensor to handle DML operations on streaming tables [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #61148:
URL: https://github.com/apache/airflow/pull/61148#issuecomment-4387606480
@radhwene This draft PR has been inactive for 9 days since the last triage comment and no response from the author. Closing to keep the queue clean.
You are welcome to reopen this PR when you resume work, or to open a new one addressing the issues previously raised. There is no rush — take your time.
---
_Note: This comment was drafted by an AI-assisted triage tool and may contain mistakes. Once you have addressed the points above, an Apache Airflow maintainer — a real person — will take the next look at your PR. We use this [two-stage triage process](https://github.com/apache/airflow/blob/main/contributing-docs/25_maintainer_pr_triage.md#why-the-first-pass-is-automated) so that our maintainers' limited time is spent where it matters most: the conversation with you._
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement BigQueryStreamingBufferEmptySensor to handle DML operations on streaming tables [airflow]
Posted by "radhwene (via GitHub)" <gi...@apache.org>.
radhwene commented on PR #61148:
URL: https://github.com/apache/airflow/pull/61148#issuecomment-3907874121
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 |
<img width="1536" height="383" alt="image" src="https://github.com/user-attachments/assets/8225ff7e-09b2-4db8-b6d5-2a8b992fa0cc" />
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement BigQueryStreamingBufferEmptySensor to handle DML operations on streaming tables [airflow]
Posted by "radhwene (via GitHub)" <gi...@apache.org>.
radhwene commented on code in PR #61148:
URL: https://github.com/apache/airflow/pull/61148#discussion_r2836200628
##########
providers/google/tests/system/google/cloud/bigquery/example_bigquery_sensors.py:
##########
@@ -152,22 +162,105 @@
)
# [END howto_sensor_bigquery_table_partition_async]
+ # [START howto_sensor_bigquery_streaming_buffer_empty]
+ check_streaming_buffer_empty = BigQueryStreamingBufferEmptySensor(
+ task_id="check_streaming_buffer_empty",
+ project_id=PROJECT_ID,
+ dataset_id=DATASET_NAME,
+ table_id=TABLE_NAME,
+ poke_interval=30,
+ timeout=5400, # 90 minutes - Google Cloud flushes streaming buffer within 90 minutes
+ )
+ # [END howto_sensor_bigquery_streaming_buffer_empty]
+
+ # Streaming operations: INSERT, UPDATE, DELETE
+ # These operations write data to the streaming buffer before being flushed to persistent storage
+ stream_insert = BigQueryInsertJobOperator(
+ task_id="stream_insert",
+ configuration={
+ "query": {
+ "query": STREAMING_INSERT_QUERY,
+ "useLegacySql": False,
+ }
+ },
+ )
+
+ stream_update = BigQueryInsertJobOperator(
+ task_id="stream_update",
+ configuration={
+ "query": {
+ "query": STREAMING_UPDATE_QUERY,
+ "useLegacySql": False,
+ }
+ },
+ )
+
+ stream_delete = BigQueryInsertJobOperator(
+ task_id="stream_delete",
+ configuration={
+ "query": {
+ "query": STREAMING_DELETE_QUERY,
+ "useLegacySql": False,
+ }
+ },
+ )
+
+ # [START howto_sensor_bigquery_streaming_buffer_empty_defered]
+ check_streaming_buffer_empty_def = BigQueryStreamingBufferEmptySensor(
+ task_id="check_streaming_buffer_empty_def",
+ project_id=PROJECT_ID,
+ dataset_id=DATASET_NAME,
+ table_id=TABLE_NAME,
+ deferrable=True,
+ poke_interval=30,
+ timeout=5400, # 90 minutes - Google Cloud flushes streaming buffer within 90 minutes
+ )
+ # [END howto_sensor_bigquery_streaming_buffer_empty_defered]
+
+ # [START howto_sensor_bigquery_streaming_buffer_empty_async]
+ check_streaming_buffer_empty_async = BigQueryStreamingBufferEmptySensor(
+ task_id="check_streaming_buffer_empty_async",
+ project_id=PROJECT_ID,
+ dataset_id=DATASET_NAME,
+ table_id=TABLE_NAME,
+ poke_interval=30,
+ timeout=5400, # 90 minutes - Google Cloud flushes streaming buffer within 90 minutes
+ )
+ # [END howto_sensor_bigquery_streaming_buffer_empty_async]
Review Comment:
Hi @MaksYermak, thanks for the review!
You're right about the async example — it was just a duplicate of the sync task without **deferrable=True**.
I've added **deferrable=True**
to make it coherent. I also removed the deprecated
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement BigQueryStreamingBufferEmptySensor to handle DML operations on streaming tables [airflow]
Posted by "radhwene (via GitHub)" <gi...@apache.org>.
radhwene commented on PR #61148:
URL: https://github.com/apache/airflow/pull/61148#issuecomment-3876173310
## 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
<img width="704" height="300" alt="image" src="https://github.com/user-attachments/assets/88360e20-2c21-4398-93db-692eab42d159" />
### Tested on:
- Airflow 2.10.3 / Python 3.12
- google-cloud-bigquery SDK
- GCP project with BigQuery API enabled
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement BigQueryStreamingBufferEmptySensor to handle DML operations on streaming tables [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #61148:
URL: https://github.com/apache/airflow/pull/61148#issuecomment-4323610441
@radhwene This PR has been converted to **draft** because it does not yet meet our [Pull Request quality criteria](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-quality-criteria).
**Issues found:**
- :x: **Failing CI checks**: 5 checks failing past the 96h grace window — `CI image checks / Static checks`, `MyPy providers checks`, `provider distributions tests / Compat 2.11.1:P3.10:`, `provider distributions tests / Compat 3.0.6:P3.10:`, `provider distributions tests / Compat 3.1.8:P3.10:`. The static checks and mypy must be fixed locally before pushing — they are deterministic, not flaky.
- :x: **Unresolved review comments**: There are 11 unresolved review threads on this PR. For each one: either apply the suggestion in a follow-up commit, or reply in-thread with a brief explanation of why the feedback doesn't apply, or resolve the thread if the feedback is no longer relevant. **Once you believe a thread is resolved — whether by pushing a fix or by explaining why the suggestion doesn't apply — it is the author's responsibility to mark it as resolved by clicking the 'Resolve conversation' button at the bottom of the thread.** Reviewers don't auto-close their own threads, so unresolved threads read as 'still waiting on the author' and block the PR from moving forward.
**What to do next:**
- Run `prek run --from-ref main --stage pre-commit` and `breeze run mypy providers/google` locally to fix static checks and mypy.
- Address each of the 11 unresolved review threads.
- Mark the PR as "Ready for review" when CI is green and threads are resolved.
Converting a PR to draft is **not** a rejection — it is an invitation to bring the PR up to the project's standards so that maintainer review time is spent productively. There is no rush — take your time and work at your own pace. We appreciate your contribution and are happy to wait for updates. If you have questions, feel free to ask on the [Airflow Slack](https://s.apache.org/airflow-slack).
---
_Note: This comment was drafted by an AI-assisted triage tool and may contain mistakes. Once you have addressed the points above, an Apache Airflow maintainer — a real person — will take the next look at your PR. We use this [two-stage triage process](https://github.com/apache/airflow/blob/main/contributing-docs/25_maintainer_pr_triage.md#why-the-first-pass-is-automated) so that our maintainers' limited time is spent where it matters most: the conversation with you._
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement BigQueryStreamingBufferEmptySensor to handle DML operations on streaming tables [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #61148:
URL: https://github.com/apache/airflow/pull/61148#issuecomment-4043195338
@radhwene This PR has been converted to **draft** because it does not yet meet our [Pull Request quality criteria](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-quality-criteria).
**Issues found:**
- :x: **Other failing CI checks**: Failing: Postgres tests: core / DB-core:Postgres:14:3.10:API...Serialization, MySQL tests: core / DB-core:MySQL:8.0:3.10:API...Serialization, Sqlite tests: core / DB-core:Sqlite:3.10:API...Serialization, Special tests / Min SQLAlchemy test: core / DB-core:MinSQLAlchemy-Postgres:14:3.10:API...Serialization, Special tests / Latest Boto test: core / All-core:LatestBoto-Postgres:14:3.10:API...Serialization (+1 more). Run `prek run --from-ref main` locally to reproduce. See [static checks docs](https://github.com/apache/airflow/blob/main/contributing-docs/08_static_code_checks.rst).
> **Note:** Your branch is **360 commits behind `main`**. Some check failures may be caused by changes in the base branch rather than by your PR. Please rebase your branch and push again to get up-to-date CI results.
**What to do next:**
- The comment informs you what you need to do.
- Fix each issue, then mark the PR as "Ready for review" in the GitHub UI - but only after making sure that all the issues are fixed.
- Maintainers will then proceed with a normal review.
Converting a PR to draft is **not** a rejection — it is an invitation to bring the PR up to the project's standards so that maintainer review time is spent productively. If you have questions, feel free to ask on the [Airflow Slack](https://s.apache.org/airflow-slack).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement BigQueryStreamingBufferEmptySensor to handle DML operations on streaming tables [airflow]
Posted by "shahar1 (via GitHub)" <gi...@apache.org>.
shahar1 commented on PR #61148:
URL: https://github.com/apache/airflow/pull/61148#issuecomment-3875598803
> # 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
> ```python
> 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
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement BigQueryStreamingBufferEmptySensor to handle DML operations on streaming tables [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #61148:
URL: https://github.com/apache/airflow/pull/61148#issuecomment-4292935928
@radhwene — This PR has new commits since the last review requesting changes from @MaksYermak. If you believe you've addressed the feedback, please ping the reviewer(s) to request a re-review. Thanks!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement BigQueryStreamingBufferEmptySensor to handle DML operations on streaming tables [airflow]
Posted by "MaksYermak (via GitHub)" <gi...@apache.org>.
MaksYermak commented on code in PR #61148:
URL: https://github.com/apache/airflow/pull/61148#discussion_r2852449901
##########
providers/google/src/airflow/providers/google/cloud/sensors/bigquery.py:
##########
@@ -77,17 +76,7 @@ def __init__(
**kwargs,
) -> None:
if deferrable and "poke_interval" not in kwargs:
- # TODO: Remove once deprecated
- if "polling_interval" in kwargs:
- kwargs["poke_interval"] = kwargs["polling_interval"]
- warnings.warn(
- "Argument `poll_interval` is deprecated and will be removed "
- "in a future release. Please use `poke_interval` instead.",
- AirflowProviderDeprecationWarning,
- stacklevel=2,
- )
- else:
- kwargs["poke_interval"] = 5
Review Comment:
@radhwene could you please return this deprecated code back for the current operator? In google provider we have deprecation policy https://airflow.apache.org/docs/apache-airflow-providers-google/stable/deprecation-policy.html. We can not remove deprecated code without specifying planned removal date.
##########
providers/google/tests/system/google/cloud/bigquery/example_bigquery_sensors.py:
##########
@@ -152,22 +162,105 @@
)
# [END howto_sensor_bigquery_table_partition_async]
+ # [START howto_sensor_bigquery_streaming_buffer_empty]
+ check_streaming_buffer_empty = BigQueryStreamingBufferEmptySensor(
+ task_id="check_streaming_buffer_empty",
+ project_id=PROJECT_ID,
+ dataset_id=DATASET_NAME,
+ table_id=TABLE_NAME,
+ poke_interval=30,
+ timeout=5400, # 90 minutes - Google Cloud flushes streaming buffer within 90 minutes
+ )
+ # [END howto_sensor_bigquery_streaming_buffer_empty]
+
+ # Streaming operations: INSERT, UPDATE, DELETE
+ # These operations write data to the streaming buffer before being flushed to persistent storage
+ stream_insert = BigQueryInsertJobOperator(
+ task_id="stream_insert",
+ configuration={
+ "query": {
+ "query": STREAMING_INSERT_QUERY,
+ "useLegacySql": False,
+ }
+ },
+ )
+
+ stream_update = BigQueryInsertJobOperator(
+ task_id="stream_update",
+ configuration={
+ "query": {
+ "query": STREAMING_UPDATE_QUERY,
+ "useLegacySql": False,
+ }
+ },
+ )
+
+ stream_delete = BigQueryInsertJobOperator(
+ task_id="stream_delete",
+ configuration={
+ "query": {
+ "query": STREAMING_DELETE_QUERY,
+ "useLegacySql": False,
+ }
+ },
+ )
+
+ # [START howto_sensor_bigquery_streaming_buffer_empty_defered]
+ check_streaming_buffer_empty_def = BigQueryStreamingBufferEmptySensor(
+ task_id="check_streaming_buffer_empty_def",
+ project_id=PROJECT_ID,
+ dataset_id=DATASET_NAME,
+ table_id=TABLE_NAME,
+ deferrable=True,
+ poke_interval=30,
+ timeout=5400, # 90 minutes - Google Cloud flushes streaming buffer within 90 minutes
+ )
+ # [END howto_sensor_bigquery_streaming_buffer_empty_defered]
+
+ # [START howto_sensor_bigquery_streaming_buffer_empty_async]
+ check_streaming_buffer_empty_async = BigQueryStreamingBufferEmptySensor(
+ task_id="check_streaming_buffer_empty_async",
+ project_id=PROJECT_ID,
+ dataset_id=DATASET_NAME,
+ table_id=TABLE_NAME,
+ poke_interval=30,
+ timeout=5400, # 90 minutes - Google Cloud flushes streaming buffer within 90 minutes
+ )
+ # [END howto_sensor_bigquery_streaming_buffer_empty_async]
Review Comment:
@radhwene thank you for the explanation. If it is duplicate from the sync task then could you please remove this task from the system test and update the documentation?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement BigQueryStreamingBufferEmptySensor to handle DML operations on streaming tables [airflow]
Posted by "MaksYermak (via GitHub)" <gi...@apache.org>.
MaksYermak commented on code in PR #61148:
URL: https://github.com/apache/airflow/pull/61148#discussion_r2871841773
##########
providers/google/docs/operators/cloud/bigquery.rst:
##########
@@ -198,10 +198,6 @@ To fetch data from a BigQuery table you can use
Alternatively you can fetch data for selected columns if you pass fields to
``selected_fields``.
-.. note::
- The ``project_id`` parameter is **deprecated** and will be removed in a future
- release. Please use ``table_project_id`` instead.
-
Review Comment:
@radhwene why did you remove this note? Could you please return it back?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement BigQueryStreamingBufferEmptySensor to handle DML operations on streaming tables [airflow]
Posted by "radhwene (via GitHub)" <gi...@apache.org>.
radhwene commented on code in PR #61148:
URL: https://github.com/apache/airflow/pull/61148#discussion_r2876428523
##########
providers/google/docs/operators/cloud/bigquery.rst:
##########
@@ -198,10 +198,6 @@ To fetch data from a BigQuery table you can use
Alternatively you can fetch data for selected columns if you pass fields to
``selected_fields``.
-.. note::
- The ``project_id`` parameter is **deprecated** and will be removed in a future
- release. Please use ``table_project_id`` instead.
-
Review Comment:
It has been restored in the latest commit. Please let me know if anything else is missing.
<img width="3329" height="930" alt="Image" src="https://github.com/user-attachments/assets/1b0dcec5-1a84-4df7-ada8-ee45915f0b5b" />
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement BigQueryStreamingBufferEmptySensor to handle DML operations on streaming tables [airflow]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #61148:
URL: https://github.com/apache/airflow/pull/61148#discussion_r3066497835
##########
providers/google/tests/system/google/cloud/bigquery/example_bigquery_sensors.py:
##########
@@ -169,6 +232,11 @@
check_table_partition_exists_async,
check_table_partition_exists_def,
]
+ >> [stream_insert, stream_update, stream_delete]
+ >> [
+ check_streaming_buffer_empty,
+ check_streaming_buffer_empty_def,
+ ]
Review Comment:
The example describes ‘Streaming’ operations and implies it populates the streaming buffer, but `BigQueryInsertJobOperator` runs query jobs (including `INSERT ... VALUES`), which do not use the streaming buffer mechanism created by `tabledata.insertAll`. This makes the example misleading and may not demonstrate the real failure mode this sensor addresses. Consider updating the example to (a) perform an actual streaming insert (e.g., via the provider’s streaming insert operator/hook), then (b) run DML only after `BigQueryStreamingBufferEmptySensor` succeeds (i.e., `streaming_insert >> check_streaming_buffer_empty >> dml_update/delete`).
```suggestion
>> stream_insert
>> [
check_streaming_buffer_empty,
check_streaming_buffer_empty_def,
]
>> [stream_update, stream_delete]
```
##########
providers/google/src/airflow/providers/google/cloud/sensors/bigquery.py:
##########
@@ -256,3 +255,143 @@ def execute_complete(self, context: dict[str, Any], event: dict[str, str] | None
message = "No event received in trigger callback"
raise AirflowException(message)
+
+
+class BigQueryStreamingBufferEmptySensor(BaseSensorOperator):
+ """
+ Sensor for checking whether the streaming buffer in a BigQuery table is empty.
+
+ The BigQueryStreamingBufferEmptySensor waits for the streaming buffer in a specified
+ BigQuery table to be empty before proceeding. It can be used in ETL pipelines to ensure
+ that recent streamed data has been processed before continuing downstream tasks.
+
+ :ivar template_fields: Fields that can be templated in this operator.
+ :type template_fields: Sequence[str]
+ :ivar ui_color: Color of the operator in the Airflow UI.
+ :type ui_color: Str
+ :ivar project_id: The Google Cloud project ID where the BigQuery table resides.
+ :type project_id: Str
+ :ivar dataset_id: The ID of the dataset containing the BigQuery table.
+ :type dataset_id: Str
+ :ivar table_id: The ID of the BigQuery table to monitor.
+ :type table_id: Str
+ :ivar gcp_conn_id: The Airflow connection ID for GCP. Defaults to "google_cloud_default".
+ :type gcp_conn_id: Str
+ :ivar impersonation_chain: Optional array or string of service accounts to impersonate using short-term
+ credentials. If multiple accounts are provided, the service account must
+ grant the role `roles/iam.serviceAccountTokenCreator` on the next account
+ in the chain.
+ :type impersonation_chain: Str | Sequence[str] | None
+ :ivar deferrable: Indicates whether the operator supports deferrable execution. If True, the sensor
+ can defer instead of polling, leading to reduced resource use.
+ :type deferrable: Bool
+ """
+
+ template_fields: Sequence[str] = (
+ "project_id",
+ "dataset_id",
+ "table_id",
+ "impersonation_chain",
+ )
+
+ ui_color = "#f0eee4"
+
+ def __init__(
+ self,
+ *,
+ project_id: str,
+ dataset_id: str,
+ table_id: str,
+ gcp_conn_id: str = "google_cloud_default",
+ impersonation_chain: str | Sequence[str] | None = None,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ) -> None:
+ if deferrable and "poke_interval" not in kwargs:
+ kwargs["poke_interval"] = 30
+
+ super().__init__(**kwargs)
+
+ self.project_id = project_id
+ self.dataset_id = dataset_id
+ self.table_id = table_id
+ self.gcp_conn_id = gcp_conn_id
+ self.impersonation_chain = impersonation_chain
+ self.deferrable = deferrable
+
+ def execute(self, context: Context) -> None:
+ """
+ Executes the operator logic taking into account the `deferrable` attribute.
+ If not deferrable, it uses the base class `execute` method; otherwise, it
+ sets up deferral with a specific trigger to wait until the BigQuery Streaming
+ Buffer is empty.
+
+ :param context: The execution context provided by Airflow. It allows
+ access to metadata and runtime information of the task instance.
+ :type context: Context
+ :return: None. The method does not return anything but performs actions
+ relevant to the operator's execution.
+ """
+ if not self.deferrable:
+ super().execute(context)
+ else:
+ if not self.poke(context=context):
+ self.defer(
+ timeout=timedelta(seconds=self.timeout),
+ trigger=BigQueryStreamingBufferEmptyTrigger(
+ project_id=self.project_id,
+ dataset_id=self.dataset_id,
+ table_id=self.table_id,
+ poll_interval=self.poke_interval,
+ gcp_conn_id=self.gcp_conn_id,
+ hook_params={
+ "impersonation_chain": self.impersonation_chain,
+ },
+ ),
+ method_name="execute_complete",
+ )
Review Comment:
The sensor passes `impersonation_chain` only inside `hook_params`, but the trigger’s async hook construction uses the trigger’s `impersonation_chain` attribute (not `hook_params`). As a result, deferrable executions will ignore impersonation. Fix by passing `impersonation_chain=self.impersonation_chain` to `BigQueryStreamingBufferEmptyTrigger(...)` (and/or update the trigger to actually consume `hook_params`).
##########
providers/google/tests/unit/google/cloud/triggers/test_bigquery.py:
##########
@@ -888,3 +888,168 @@ def test_serialization_successfully(self):
"poll_interval": POLLING_PERIOD_SECONDS,
"hook_params": TEST_HOOK_PARAMS,
}
+
+
+@pytest.fixture
+def streaming_buffer_trigger():
+ return BigQueryStreamingBufferEmptyTrigger(
+ project_id=TEST_GCP_PROJECT_ID,
+ dataset_id=TEST_DATASET_ID,
+ table_id=TEST_TABLE_ID,
+ gcp_conn_id=TEST_GCP_CONN_ID,
+ hook_params=TEST_HOOK_PARAMS,
+ poll_interval=POLLING_PERIOD_SECONDS,
+ impersonation_chain=TEST_IMPERSONATION_CHAIN,
+ )
+
+
+class TestBigQueryStreamingBufferEmptyTrigger:
+ def test_serialization(self, streaming_buffer_trigger):
+ """Asserts that the trigger correctly serializes its arguments and classpath."""
+ classpath, kwargs = streaming_buffer_trigger.serialize()
+ assert (
+ classpath
+ == "airflow.providers.google.cloud.triggers.bigquery.BigQueryStreamingBufferEmptyTrigger"
+ )
+ assert kwargs == {
+ "project_id": TEST_GCP_PROJECT_ID,
+ "dataset_id": TEST_DATASET_ID,
+ "table_id": TEST_TABLE_ID,
+ "gcp_conn_id": TEST_GCP_CONN_ID,
+ "poll_interval": POLLING_PERIOD_SECONDS,
+ "hook_params": TEST_HOOK_PARAMS,
+ "impersonation_chain": TEST_IMPERSONATION_CHAIN,
+ }
+
+ @pytest.mark.asyncio
+ @mock.patch(
+ "airflow.providers.google.cloud.triggers.bigquery."
+ "BigQueryStreamingBufferEmptyTrigger._is_streaming_buffer_empty"
+ )
+ @mock.patch(
+ "airflow.providers.google.cloud.triggers.bigquery.BigQueryStreamingBufferEmptyTrigger._get_sync_hook"
+ )
+ async def test_run_buffer_empty_yields_success(
+ self, mock_get_hook, mock_is_empty, streaming_buffer_trigger
+ ):
+ """When buffer is empty, trigger yields a success TriggerEvent."""
+ mock_is_empty.return_value = True
+
+ generator = streaming_buffer_trigger.run()
+ actual = await generator.asend(None)
+
+ table_uri = f"{TEST_GCP_PROJECT_ID}:{TEST_DATASET_ID}.{TEST_TABLE_ID}"
+ assert actual == TriggerEvent(
+ {"status": "success", "message": f"Streaming buffer is empty for table: {table_uri}"}
+ )
+
+ @pytest.mark.asyncio
+ @mock.patch(
+ "airflow.providers.google.cloud.triggers.bigquery."
+ "BigQueryStreamingBufferEmptyTrigger._is_streaming_buffer_empty"
+ )
+ @mock.patch(
+ "airflow.providers.google.cloud.triggers.bigquery.BigQueryStreamingBufferEmptyTrigger._get_sync_hook"
+ )
+ async def test_run_buffer_not_empty_keeps_polling(
+ self, mock_get_hook, mock_is_empty, streaming_buffer_trigger
+ ):
+ """When buffer is not empty, trigger keeps polling (does not yield immediately)."""
+ mock_is_empty.return_value = False
+
+ task = asyncio.create_task(streaming_buffer_trigger.run().__anext__())
+ await asyncio.sleep(0.5)
+
+ # TriggerEvent was not returned yet
+ assert task.done() is False
+ asyncio.get_event_loop().stop()
Review Comment:
These tests don’t match the trigger implementation: (1) the trigger exposes `_get_async_hook`, not `_get_sync_hook`, so the patch target is invalid; (2) `_is_streaming_buffer_empty` currently uses `BigQueryTableAsyncHook.get_table_client()` + `await client.get()`, but the tests are written around `BigQueryHook.get_client().get_table()` and SDK `BQTable`. Update the tests to mock `BigQueryTableAsyncHook.get_table_client` and the returned client’s async `get()` response (including `streamingBuffer` presence/absence). Also avoid `asyncio.get_event_loop().stop()` in tests—cancel the task and assert via a timeout (e.g., `asyncio.wait_for(..., timeout=...)`) to prevent leaking tasks and interfering with the pytest-asyncio loop.
##########
providers/google/src/airflow/providers/google/cloud/triggers/bigquery.py:
##########
@@ -806,3 +806,158 @@ async def _partition_exists(self, hook: BigQueryAsyncHook, job_id: str | None, p
if records:
records = [row[0] for row in records]
return self.partition_id in records
+
+
+class BigQueryStreamingBufferEmptyTrigger(BaseTrigger):
+ """
+ Trigger that periodically checks if a BigQuery table's streaming buffer is empty.
+
+ This trigger continuously polls a BigQuery table to determine if its streaming buffer
+ has been fully processed and is now empty. It's particularly useful before running
+ DML operations (UPDATE, DELETE, MERGE) on tables populated via streaming inserts.
+
+ :param project_id: The Google Cloud Project in which to look for the table.
+ :param dataset_id: The dataset ID of the table to check.
+ :param table_id: The table ID of the table to check.
+ :param gcp_conn_id: Reference to Google Cloud connection ID.
+ :param hook_params: Additional parameters for hook initialization.
+ :param poll_interval: Polling period in seconds to check the streaming buffer status.
+ :param impersonation_chain: Optional service account to impersonate using short-term
+ credentials, or chained list of accounts required to get the access_token
+ of the last account in the list, which will be impersonated in the request.
+ If set as a string, the account must grant the originating account
+ the Service Account Token Creator IAM role.
+ If set as a sequence, the identities from the list must grant
+ Service Account Token Creator IAM role to the directly preceding identity, with first
+ account from the list granting this role to the originating account.
+ """
+
+ def __init__(
+ self,
+ project_id: str,
+ dataset_id: str,
+ table_id: str,
+ gcp_conn_id: str,
+ hook_params: dict[str, Any],
+ poll_interval: float = 30.0,
+ impersonation_chain: str | Sequence[str] | None = None,
+ ):
+ super().__init__()
+ self.project_id = project_id
+ self.dataset_id = dataset_id
+ self.table_id = table_id
+ self.gcp_conn_id = gcp_conn_id
+ self.poll_interval = poll_interval
+ self.hook_params = hook_params
+ self.impersonation_chain = impersonation_chain
+
+ def serialize(self) -> tuple[str, dict[str, Any]]:
+ """Serialize BigQueryStreamingBufferEmptyTrigger arguments and classpath."""
+ return (
+ "airflow.providers.google.cloud.triggers.bigquery.BigQueryStreamingBufferEmptyTrigger",
+ {
+ "project_id": self.project_id,
+ "dataset_id": self.dataset_id,
+ "table_id": self.table_id,
+ "gcp_conn_id": self.gcp_conn_id,
+ "poll_interval": self.poll_interval,
+ "hook_params": self.hook_params,
+ "impersonation_chain": self.impersonation_chain,
+ },
+ )
+
+ def _get_async_hook(self) -> BigQueryTableAsyncHook:
+ """Get the async hook for BigQuery table operations."""
+ return BigQueryTableAsyncHook(
+ gcp_conn_id=self.gcp_conn_id,
+ impersonation_chain=self.impersonation_chain,
+ )
+
+ async def run(self) -> AsyncIterator[TriggerEvent]:
+ """
+ Continuously check if the streaming buffer is empty.
+
+ Yields a TriggerEvent when the streaming buffer becomes empty or if an error occurs.
+ """
+ try:
+ hook = self._get_async_hook()
+ while True:
+ self.log.info(
+ "Checking streaming buffer for table %s.%s.%s",
+ self.project_id,
+ self.dataset_id,
+ self.table_id,
+ )
+
+ is_buffer_empty = await self._is_streaming_buffer_empty(
+ hook=hook,
+ project_id=self.project_id,
+ dataset_id=self.dataset_id,
+ table_id=self.table_id,
+ )
+
+ if is_buffer_empty:
+ table_uri = f"{self.project_id}:{self.dataset_id}.{self.table_id}"
+ message = f"Streaming buffer is empty for table: {table_uri}"
+ self.log.info(message)
+ yield TriggerEvent(
+ {
+ "status": "success",
+ "message": message,
+ }
+ )
+ return
+ else:
+ self.log.info(
+ "Streaming buffer still has data. Sleeping for %s seconds.",
+ self.poll_interval,
+ )
+ await asyncio.sleep(self.poll_interval)
+
+ except Exception as e:
+ self.log.exception(
+ "Exception occurred while checking streaming buffer for table %s.%s.%s",
+ self.project_id,
+ self.dataset_id,
+ self.table_id,
+ )
+ yield TriggerEvent({"status": "error", "message": str(e)})
+
+ async def _is_streaming_buffer_empty(
+ self,
+ hook: BigQueryTableAsyncHook,
+ project_id: str,
+ dataset_id: str,
+ table_id: str,
+ ) -> bool:
+ """
+ Check if the streaming buffer is empty for the specified table.
+
+ :param hook: BigQueryTableAsyncHook instance for async operations.
+ :param project_id: The Google Cloud Project ID.
+ :param dataset_id: The dataset ID containing the table.
+ :param table_id: The table ID to check.
+ :return: True if streaming buffer is empty or doesn't exist, False otherwise.
+ """
+ async with ClientSession() as session:
+ try:
+ client = await hook.get_table_client(
+ dataset=dataset_id,
+ table_id=table_id,
+ project_id=project_id,
+ session=session,
+ )
Review Comment:
`ClientSession()` is created and torn down on every poll. Since `run()` loops indefinitely, this adds unnecessary overhead and connection churn. Create a single `ClientSession` in `run()` (wrapping the whole polling loop) and pass it into `_is_streaming_buffer_empty(...)`, or refactor `_is_streaming_buffer_empty` to accept a pre-created session.
##########
providers/google/src/airflow/providers/google/cloud/sensors/bigquery.py:
##########
@@ -256,3 +255,143 @@ def execute_complete(self, context: dict[str, Any], event: dict[str, str] | None
message = "No event received in trigger callback"
raise AirflowException(message)
+
+
+class BigQueryStreamingBufferEmptySensor(BaseSensorOperator):
+ """
+ Sensor for checking whether the streaming buffer in a BigQuery table is empty.
+
+ The BigQueryStreamingBufferEmptySensor waits for the streaming buffer in a specified
+ BigQuery table to be empty before proceeding. It can be used in ETL pipelines to ensure
+ that recent streamed data has been processed before continuing downstream tasks.
+
+ :ivar template_fields: Fields that can be templated in this operator.
+ :type template_fields: Sequence[str]
+ :ivar ui_color: Color of the operator in the Airflow UI.
+ :type ui_color: Str
+ :ivar project_id: The Google Cloud project ID where the BigQuery table resides.
+ :type project_id: Str
+ :ivar dataset_id: The ID of the dataset containing the BigQuery table.
+ :type dataset_id: Str
+ :ivar table_id: The ID of the BigQuery table to monitor.
+ :type table_id: Str
+ :ivar gcp_conn_id: The Airflow connection ID for GCP. Defaults to "google_cloud_default".
+ :type gcp_conn_id: Str
+ :ivar impersonation_chain: Optional array or string of service accounts to impersonate using short-term
+ credentials. If multiple accounts are provided, the service account must
+ grant the role `roles/iam.serviceAccountTokenCreator` on the next account
+ in the chain.
+ :type impersonation_chain: Str | Sequence[str] | None
+ :ivar deferrable: Indicates whether the operator supports deferrable execution. If True, the sensor
+ can defer instead of polling, leading to reduced resource use.
+ :type deferrable: Bool
+ """
+
+ template_fields: Sequence[str] = (
+ "project_id",
+ "dataset_id",
+ "table_id",
+ "impersonation_chain",
+ )
+
+ ui_color = "#f0eee4"
+
+ def __init__(
+ self,
+ *,
+ project_id: str,
+ dataset_id: str,
+ table_id: str,
+ gcp_conn_id: str = "google_cloud_default",
+ impersonation_chain: str | Sequence[str] | None = None,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ) -> None:
+ if deferrable and "poke_interval" not in kwargs:
+ kwargs["poke_interval"] = 30
+
+ super().__init__(**kwargs)
+
+ self.project_id = project_id
+ self.dataset_id = dataset_id
+ self.table_id = table_id
+ self.gcp_conn_id = gcp_conn_id
+ self.impersonation_chain = impersonation_chain
+ self.deferrable = deferrable
+
+ def execute(self, context: Context) -> None:
+ """
+ Executes the operator logic taking into account the `deferrable` attribute.
+ If not deferrable, it uses the base class `execute` method; otherwise, it
+ sets up deferral with a specific trigger to wait until the BigQuery Streaming
+ Buffer is empty.
+
+ :param context: The execution context provided by Airflow. It allows
+ access to metadata and runtime information of the task instance.
+ :type context: Context
+ :return: None. The method does not return anything but performs actions
+ relevant to the operator's execution.
+ """
+ if not self.deferrable:
+ super().execute(context)
+ else:
+ if not self.poke(context=context):
+ self.defer(
+ timeout=timedelta(seconds=self.timeout),
+ trigger=BigQueryStreamingBufferEmptyTrigger(
+ project_id=self.project_id,
+ dataset_id=self.dataset_id,
+ table_id=self.table_id,
+ poll_interval=self.poke_interval,
+ gcp_conn_id=self.gcp_conn_id,
+ hook_params={
+ "impersonation_chain": self.impersonation_chain,
+ },
+ ),
+ method_name="execute_complete",
+ )
+
+ def execute_complete(self, context: dict[str, Any], event: dict[str, str] | None = None) -> str:
+ """
+ Act as a callback for when the trigger fires - returns immediately.
+
+ Relies on trigger to throw an exception, otherwise it assumes execution was successful.
+ """
+ table_uri = f"{self.project_id}:{self.dataset_id}.{self.table_id}"
+ self.log.info("Checking streaming buffer state for table: %s", table_uri)
+ if event:
+ if event["status"] == "success":
+ return event["message"]
+ raise AirflowException(event["message"])
+
+ message = "No event received in trigger callback"
+ raise AirflowException(message)
+
+ def poke(self, context: Context) -> bool:
+ """
+ Check if the BigQuery streaming buffer is empty for the specified table.
+
+ This method periodically checks the status of the BigQuery table's streaming buffer
+ to determine if it is empty. It is useful for ensuring that recent streamed data
+ has been fully processed before continuing with downstream tasks.
+ """
+ table_uri = f"{self.project_id}:{self.dataset_id}.{self.table_id}"
+ self.log.info("Checking streaming buffer state for table: %s", table_uri)
+
+ hook = BigQueryHook(
+ gcp_conn_id=self.gcp_conn_id,
+ impersonation_chain=self.impersonation_chain,
+ )
+ try:
+ table = hook.get_table(
+ project_id=self.project_id,
+ dataset_id=self.dataset_id,
+ table_id=self.table_id,
+ )
+ return table.get("streamingBuffer") is None
+ except Exception as err:
+ if "not found" in str(err):
+ raise AirflowException(
+ f"Table {self.project_id}.{self.dataset_id}.{self.table_id} not found"
+ ) from err
+ raise err
Review Comment:
This assumes `hook.get_table(...)` returns a dict-like object (using `.get(...)`). In many BigQuery client paths, `get_table` returns a `google.cloud.bigquery.table.Table` where the correct check is via `table.streaming_buffer` (or checking the API repr). Align the sync sensor with the trigger/tests by using the SDK Table object property (and catch `google.api_core.exceptions.NotFound` explicitly instead of string-matching `'not found'`, which is brittle and can misclassify errors).
##########
providers/google/src/airflow/providers/google/cloud/triggers/bigquery.py:
##########
@@ -806,3 +806,158 @@ async def _partition_exists(self, hook: BigQueryAsyncHook, job_id: str | None, p
if records:
records = [row[0] for row in records]
return self.partition_id in records
+
+
+class BigQueryStreamingBufferEmptyTrigger(BaseTrigger):
+ """
+ Trigger that periodically checks if a BigQuery table's streaming buffer is empty.
+
+ This trigger continuously polls a BigQuery table to determine if its streaming buffer
+ has been fully processed and is now empty. It's particularly useful before running
+ DML operations (UPDATE, DELETE, MERGE) on tables populated via streaming inserts.
+
+ :param project_id: The Google Cloud Project in which to look for the table.
+ :param dataset_id: The dataset ID of the table to check.
+ :param table_id: The table ID of the table to check.
+ :param gcp_conn_id: Reference to Google Cloud connection ID.
+ :param hook_params: Additional parameters for hook initialization.
+ :param poll_interval: Polling period in seconds to check the streaming buffer status.
+ :param impersonation_chain: Optional service account to impersonate using short-term
+ credentials, or chained list of accounts required to get the access_token
+ of the last account in the list, which will be impersonated in the request.
+ If set as a string, the account must grant the originating account
+ the Service Account Token Creator IAM role.
+ If set as a sequence, the identities from the list must grant
+ Service Account Token Creator IAM role to the directly preceding identity, with first
+ account from the list granting this role to the originating account.
+ """
+
+ def __init__(
+ self,
+ project_id: str,
+ dataset_id: str,
+ table_id: str,
+ gcp_conn_id: str,
+ hook_params: dict[str, Any],
+ poll_interval: float = 30.0,
+ impersonation_chain: str | Sequence[str] | None = None,
+ ):
+ super().__init__()
+ self.project_id = project_id
+ self.dataset_id = dataset_id
+ self.table_id = table_id
+ self.gcp_conn_id = gcp_conn_id
+ self.poll_interval = poll_interval
+ self.hook_params = hook_params
+ self.impersonation_chain = impersonation_chain
+
+ def serialize(self) -> tuple[str, dict[str, Any]]:
+ """Serialize BigQueryStreamingBufferEmptyTrigger arguments and classpath."""
+ return (
+ "airflow.providers.google.cloud.triggers.bigquery.BigQueryStreamingBufferEmptyTrigger",
+ {
+ "project_id": self.project_id,
+ "dataset_id": self.dataset_id,
+ "table_id": self.table_id,
+ "gcp_conn_id": self.gcp_conn_id,
+ "poll_interval": self.poll_interval,
+ "hook_params": self.hook_params,
+ "impersonation_chain": self.impersonation_chain,
+ },
+ )
+
+ def _get_async_hook(self) -> BigQueryTableAsyncHook:
+ """Get the async hook for BigQuery table operations."""
+ return BigQueryTableAsyncHook(
+ gcp_conn_id=self.gcp_conn_id,
+ impersonation_chain=self.impersonation_chain,
Review Comment:
`hook_params` is serialized/stored but never used by the trigger. This is confusing for API consumers and contributed to the impersonation bug in the sensor callsite. Either remove `hook_params` from the trigger API (and serialization), or apply it when building the hook (e.g., derive `impersonation_chain` from it, and/or pass supported hook args through consistently).
```suggestion
impersonation_chain = self.impersonation_chain
if impersonation_chain is None:
impersonation_chain = self.hook_params.get("impersonation_chain")
return BigQueryTableAsyncHook(
gcp_conn_id=self.gcp_conn_id,
impersonation_chain=impersonation_chain,
```
##########
providers/google/tests/unit/google/cloud/triggers/test_bigquery.py:
##########
@@ -888,3 +888,168 @@ def test_serialization_successfully(self):
"poll_interval": POLLING_PERIOD_SECONDS,
"hook_params": TEST_HOOK_PARAMS,
}
+
+
+@pytest.fixture
+def streaming_buffer_trigger():
+ return BigQueryStreamingBufferEmptyTrigger(
+ project_id=TEST_GCP_PROJECT_ID,
+ dataset_id=TEST_DATASET_ID,
+ table_id=TEST_TABLE_ID,
+ gcp_conn_id=TEST_GCP_CONN_ID,
+ hook_params=TEST_HOOK_PARAMS,
+ poll_interval=POLLING_PERIOD_SECONDS,
+ impersonation_chain=TEST_IMPERSONATION_CHAIN,
+ )
+
+
+class TestBigQueryStreamingBufferEmptyTrigger:
+ def test_serialization(self, streaming_buffer_trigger):
+ """Asserts that the trigger correctly serializes its arguments and classpath."""
+ classpath, kwargs = streaming_buffer_trigger.serialize()
+ assert (
+ classpath
+ == "airflow.providers.google.cloud.triggers.bigquery.BigQueryStreamingBufferEmptyTrigger"
+ )
+ assert kwargs == {
+ "project_id": TEST_GCP_PROJECT_ID,
+ "dataset_id": TEST_DATASET_ID,
+ "table_id": TEST_TABLE_ID,
+ "gcp_conn_id": TEST_GCP_CONN_ID,
+ "poll_interval": POLLING_PERIOD_SECONDS,
+ "hook_params": TEST_HOOK_PARAMS,
+ "impersonation_chain": TEST_IMPERSONATION_CHAIN,
+ }
+
+ @pytest.mark.asyncio
+ @mock.patch(
+ "airflow.providers.google.cloud.triggers.bigquery."
+ "BigQueryStreamingBufferEmptyTrigger._is_streaming_buffer_empty"
+ )
+ @mock.patch(
+ "airflow.providers.google.cloud.triggers.bigquery.BigQueryStreamingBufferEmptyTrigger._get_sync_hook"
+ )
+ async def test_run_buffer_empty_yields_success(
+ self, mock_get_hook, mock_is_empty, streaming_buffer_trigger
+ ):
Review Comment:
These tests don’t match the trigger implementation: (1) the trigger exposes `_get_async_hook`, not `_get_sync_hook`, so the patch target is invalid; (2) `_is_streaming_buffer_empty` currently uses `BigQueryTableAsyncHook.get_table_client()` + `await client.get()`, but the tests are written around `BigQueryHook.get_client().get_table()` and SDK `BQTable`. Update the tests to mock `BigQueryTableAsyncHook.get_table_client` and the returned client’s async `get()` response (including `streamingBuffer` presence/absence). Also avoid `asyncio.get_event_loop().stop()` in tests—cancel the task and assert via a timeout (e.g., `asyncio.wait_for(..., timeout=...)`) to prevent leaking tasks and interfering with the pytest-asyncio loop.
##########
providers/google/tests/system/google/cloud/bigquery/example_bigquery_sensors.py:
##########
@@ -149,9 +156,65 @@
project_id=PROJECT_ID,
dataset_id=DATASET_NAME,
table_id=TABLE_NAME,
+ deferrable=True,
)
# [END howto_sensor_bigquery_table_partition_async]
+ # [START howto_sensor_bigquery_streaming_buffer_empty]
+ check_streaming_buffer_empty = BigQueryStreamingBufferEmptySensor(
+ task_id="check_streaming_buffer_empty",
+ project_id=PROJECT_ID,
+ dataset_id=DATASET_NAME,
+ table_id=TABLE_NAME,
+ poke_interval=30,
+ timeout=5400, # 90 minutes - Google Cloud flushes streaming buffer within 90 minutes
+ )
+ # [END howto_sensor_bigquery_streaming_buffer_empty]
+
+ # Streaming operations: INSERT, UPDATE, DELETE
+ # These operations write data to the streaming buffer before being flushed to persistent storage
+ stream_insert = BigQueryInsertJobOperator(
+ task_id="stream_insert",
+ configuration={
+ "query": {
+ "query": STREAMING_INSERT_QUERY,
+ "useLegacySql": False,
+ }
+ },
+ )
Review Comment:
The example describes ‘Streaming’ operations and implies it populates the streaming buffer, but `BigQueryInsertJobOperator` runs query jobs (including `INSERT ... VALUES`), which do not use the streaming buffer mechanism created by `tabledata.insertAll`. This makes the example misleading and may not demonstrate the real failure mode this sensor addresses. Consider updating the example to (a) perform an actual streaming insert (e.g., via the provider’s streaming insert operator/hook), then (b) run DML only after `BigQueryStreamingBufferEmptySensor` succeeds (i.e., `streaming_insert >> check_streaming_buffer_empty >> dml_update/delete`).
##########
providers/google/docs/operators/cloud/bigquery.rst:
##########
@@ -518,6 +518,28 @@ Also you can use deferrable mode in this operator if you would like to free up t
:start-after: [START howto_sensor_bigquery_table_partition_async]
:end-before: [END howto_sensor_bigquery_table_partition_async]
+Check that the BigQuery Table Streaming Buffer is empty
+""""""""""""""""""""""""""""""""""""""""""""""""""""""""
+
+To check that the BigQuery streaming buffer of a table is empty you can use
+:class:`~airflow.providers.google.cloud.sensors.bigquery.BigQueryStreamingBufferEmptySensor`.
+This sensor is useful in ETL pipelines to ensure that recent streamed data has been fully
+processed before continuing downstream tasks.
+
+.. exampleinclude:: /../../google/tests/system/google/cloud/bigquery/example_bigquery_sensors.py
+ :language: python
+ :dedent: 4
+ :start-after: [START howto_sensor_bigquery_streaming_buffer_empty]
+ :end-before: [END howto_sensor_bigquery_streaming_buffer_empty]
+
+Also you can use deferrable mode in this operator if you would like to free up the worker slots while the sensor is running.
+
+.. exampleinclude:: /../../google/tests/system/google/cloud/bigquery/example_bigquery_sensors.py
+ :language: python
+ :dedent: 4
+ :start-after: [START howto_sensor_bigquery_streaming_buffer_empty_defered]
+ :end-before: [END howto_sensor_bigquery_streaming_buffer_empty_defered]
Review Comment:
Correct the tag spelling from 'defered' to 'deferred' to avoid propagating a typo into documentation references.
##########
providers/google/src/airflow/providers/google/cloud/triggers/bigquery.py:
##########
@@ -806,3 +806,158 @@ async def _partition_exists(self, hook: BigQueryAsyncHook, job_id: str | None, p
if records:
records = [row[0] for row in records]
return self.partition_id in records
+
+
+class BigQueryStreamingBufferEmptyTrigger(BaseTrigger):
+ """
+ Trigger that periodically checks if a BigQuery table's streaming buffer is empty.
+
+ This trigger continuously polls a BigQuery table to determine if its streaming buffer
+ has been fully processed and is now empty. It's particularly useful before running
+ DML operations (UPDATE, DELETE, MERGE) on tables populated via streaming inserts.
+
+ :param project_id: The Google Cloud Project in which to look for the table.
+ :param dataset_id: The dataset ID of the table to check.
+ :param table_id: The table ID of the table to check.
+ :param gcp_conn_id: Reference to Google Cloud connection ID.
+ :param hook_params: Additional parameters for hook initialization.
+ :param poll_interval: Polling period in seconds to check the streaming buffer status.
+ :param impersonation_chain: Optional service account to impersonate using short-term
+ credentials, or chained list of accounts required to get the access_token
+ of the last account in the list, which will be impersonated in the request.
+ If set as a string, the account must grant the originating account
+ the Service Account Token Creator IAM role.
+ If set as a sequence, the identities from the list must grant
+ Service Account Token Creator IAM role to the directly preceding identity, with first
+ account from the list granting this role to the originating account.
+ """
+
+ def __init__(
+ self,
+ project_id: str,
+ dataset_id: str,
+ table_id: str,
+ gcp_conn_id: str,
+ hook_params: dict[str, Any],
+ poll_interval: float = 30.0,
+ impersonation_chain: str | Sequence[str] | None = None,
+ ):
Review Comment:
`hook_params` is serialized/stored but never used by the trigger. This is confusing for API consumers and contributed to the impersonation bug in the sensor callsite. Either remove `hook_params` from the trigger API (and serialization), or apply it when building the hook (e.g., derive `impersonation_chain` from it, and/or pass supported hook args through consistently).
##########
providers/google/tests/unit/google/cloud/triggers/test_bigquery.py:
##########
@@ -888,3 +888,168 @@ def test_serialization_successfully(self):
"poll_interval": POLLING_PERIOD_SECONDS,
"hook_params": TEST_HOOK_PARAMS,
}
+
+
+@pytest.fixture
+def streaming_buffer_trigger():
+ return BigQueryStreamingBufferEmptyTrigger(
+ project_id=TEST_GCP_PROJECT_ID,
+ dataset_id=TEST_DATASET_ID,
+ table_id=TEST_TABLE_ID,
+ gcp_conn_id=TEST_GCP_CONN_ID,
+ hook_params=TEST_HOOK_PARAMS,
+ poll_interval=POLLING_PERIOD_SECONDS,
+ impersonation_chain=TEST_IMPERSONATION_CHAIN,
+ )
+
+
+class TestBigQueryStreamingBufferEmptyTrigger:
+ def test_serialization(self, streaming_buffer_trigger):
+ """Asserts that the trigger correctly serializes its arguments and classpath."""
+ classpath, kwargs = streaming_buffer_trigger.serialize()
+ assert (
+ classpath
+ == "airflow.providers.google.cloud.triggers.bigquery.BigQueryStreamingBufferEmptyTrigger"
+ )
+ assert kwargs == {
+ "project_id": TEST_GCP_PROJECT_ID,
+ "dataset_id": TEST_DATASET_ID,
+ "table_id": TEST_TABLE_ID,
+ "gcp_conn_id": TEST_GCP_CONN_ID,
+ "poll_interval": POLLING_PERIOD_SECONDS,
+ "hook_params": TEST_HOOK_PARAMS,
+ "impersonation_chain": TEST_IMPERSONATION_CHAIN,
+ }
+
+ @pytest.mark.asyncio
+ @mock.patch(
+ "airflow.providers.google.cloud.triggers.bigquery."
+ "BigQueryStreamingBufferEmptyTrigger._is_streaming_buffer_empty"
+ )
+ @mock.patch(
+ "airflow.providers.google.cloud.triggers.bigquery.BigQueryStreamingBufferEmptyTrigger._get_sync_hook"
+ )
+ async def test_run_buffer_empty_yields_success(
+ self, mock_get_hook, mock_is_empty, streaming_buffer_trigger
+ ):
+ """When buffer is empty, trigger yields a success TriggerEvent."""
+ mock_is_empty.return_value = True
+
+ generator = streaming_buffer_trigger.run()
+ actual = await generator.asend(None)
+
+ table_uri = f"{TEST_GCP_PROJECT_ID}:{TEST_DATASET_ID}.{TEST_TABLE_ID}"
+ assert actual == TriggerEvent(
+ {"status": "success", "message": f"Streaming buffer is empty for table: {table_uri}"}
+ )
+
+ @pytest.mark.asyncio
+ @mock.patch(
+ "airflow.providers.google.cloud.triggers.bigquery."
+ "BigQueryStreamingBufferEmptyTrigger._is_streaming_buffer_empty"
+ )
+ @mock.patch(
+ "airflow.providers.google.cloud.triggers.bigquery.BigQueryStreamingBufferEmptyTrigger._get_sync_hook"
+ )
+ async def test_run_buffer_not_empty_keeps_polling(
+ self, mock_get_hook, mock_is_empty, streaming_buffer_trigger
+ ):
+ """When buffer is not empty, trigger keeps polling (does not yield immediately)."""
+ mock_is_empty.return_value = False
+
+ task = asyncio.create_task(streaming_buffer_trigger.run().__anext__())
+ await asyncio.sleep(0.5)
Review Comment:
These tests don’t match the trigger implementation: (1) the trigger exposes `_get_async_hook`, not `_get_sync_hook`, so the patch target is invalid; (2) `_is_streaming_buffer_empty` currently uses `BigQueryTableAsyncHook.get_table_client()` + `await client.get()`, but the tests are written around `BigQueryHook.get_client().get_table()` and SDK `BQTable`. Update the tests to mock `BigQueryTableAsyncHook.get_table_client` and the returned client’s async `get()` response (including `streamingBuffer` presence/absence). Also avoid `asyncio.get_event_loop().stop()` in tests—cancel the task and assert via a timeout (e.g., `asyncio.wait_for(..., timeout=...)`) to prevent leaking tasks and interfering with the pytest-asyncio loop.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement BigQueryStreamingBufferEmptySensor to handle DML operations on streaming tables [airflow]
Posted by "radhwene (via GitHub)" <gi...@apache.org>.
radhwene commented on PR #61148:
URL: https://github.com/apache/airflow/pull/61148#issuecomment-3938889329
Hi @shahar1 ,
thanks for the feedback. I’ve fixed the failing unit tests: they were instantiating BigQueryHook() directly, which requires an Airflow connection (google_cloud_default) that isn’t available in the providers unit-test environment. The tests now fully mock the hook so they’re isolated and CI-friendly.
Could you please re-run the failed CI jobs when you have a moment?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement BigQueryStreamingBufferEmptySensor to handle DML operations on streaming tables [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk closed pull request #61148: Implement BigQueryStreamingBufferEmptySensor to handle DML operations on streaming tables
URL: https://github.com/apache/airflow/pull/61148
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org