This is an automated email from the ASF dual-hosted git repository.
shahar1 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 1cbcc306be6 Stop BigQueryStreamingBufferEmptySensor reporting empty
too early (#69118)
1cbcc306be6 is described below
commit 1cbcc306be68472ecbfdd1a6c488f2d507b67d94
Author: Anas Khan <[email protected]>
AuthorDate: Sat Oct 31 12:00:23 2026 +0530
Stop BigQueryStreamingBufferEmptySensor reporting empty too early (#69118)
Require consecutive empty readings before the BigQuery streaming buffer
sensor succeeds, and defer reschedule confirmations to the trigger with the
remaining timeout.
---------
Co-authored-by: Shahar Epstein <[email protected]>
---
.../providers/google/cloud/sensors/bigquery.py | 73 +++++++++--
.../providers/google/cloud/triggers/bigquery.py | 33 ++++-
.../example_bigquery_streaming_buffer_sensor.py | 19 +--
.../unit/google/cloud/sensors/test_bigquery.py | 140 ++++++++++++++++++++-
.../unit/google/cloud/triggers/test_bigquery.py | 82 +++++++++++-
5 files changed, 313 insertions(+), 34 deletions(-)
diff --git
a/providers/google/src/airflow/providers/google/cloud/sensors/bigquery.py
b/providers/google/src/airflow/providers/google/cloud/sensors/bigquery.py
index 65531b857bb..877517b31a6 100644
--- a/providers/google/src/airflow/providers/google/cloud/sensors/bigquery.py
+++ b/providers/google/src/airflow/providers/google/cloud/sensors/bigquery.py
@@ -28,15 +28,18 @@ from google.api_core.exceptions import NotFound
from google.cloud.bigquery import DatasetReference, TableReference
from airflow.exceptions import AirflowProviderDeprecationWarning
-from airflow.providers.common.compat.sdk import AirflowException,
BaseSensorOperator, conf
+from airflow.providers.common.compat.sdk import AirflowException,
BaseSensorOperator, conf, timezone
from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook
from airflow.providers.google.cloud.triggers.bigquery import (
BigQueryStreamingBufferEmptyTrigger,
BigQueryTableExistenceTrigger,
BigQueryTablePartitionExistenceTrigger,
)
+from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
if TYPE_CHECKING:
+ from datetime import datetime
+
from airflow.providers.common.compat.sdk import Context
@@ -339,12 +342,13 @@ class
BigQueryStreamingBufferEmptySensor(BaseSensorOperator):
streaming buffer`` errors.
.. warning::
- The sensor reads ``table.streaming_buffer`` from BigQuery's table
- metadata, which is eventually consistent. For a short window right
- after a streaming insert the buffer metadata is still absent, so the
- sensor may report the buffer empty before it actually is. Known
- limitation tracked at
- https://github.com/apache/airflow/issues/66963
+ BigQuery's ``table.streaming_buffer`` metadata is eventually
consistent.
+ The sensor mitigates this by requiring multiple consecutive empty
+ observations before reporting the buffer empty (configurable via
+ ``empty_confirmations``), but this may delay success by up to one
+ ``poke_interval`` compared to earlier releases. ``mode="reschedule"``
with
+ ``empty_confirmations > 1`` needs a running triggerer once an empty
reading is seen;
+ ``empty_confirmations=1`` keeps the worker-only behaviour.
:param project_id: Google Cloud project containing the table.
:param dataset_id: Dataset of the table to monitor.
@@ -352,6 +356,9 @@ class
BigQueryStreamingBufferEmptySensor(BaseSensorOperator):
:param gcp_conn_id: Airflow connection ID for GCP.
:param impersonation_chain: Optional service account to impersonate, or a
chained list of accounts. See the Google provider docs for details.
+ :param empty_confirmations: Number of consecutive empty readings required
+ before considering the streaming buffer empty. Must be at least 1.
+ Defaults to 2.
:param deferrable: Run in deferrable mode using
:class:`BigQueryStreamingBufferEmptyTrigger`.
"""
@@ -374,6 +381,7 @@ class
BigQueryStreamingBufferEmptySensor(BaseSensorOperator):
table_id: str,
gcp_conn_id: str = "google_cloud_default",
impersonation_chain: str | Sequence[str] | None = None,
+ empty_confirmations: int = 2,
deferrable: bool = conf.getboolean("operators", "default_deferrable",
fallback=False),
**kwargs,
) -> None:
@@ -382,12 +390,17 @@ class
BigQueryStreamingBufferEmptySensor(BaseSensorOperator):
super().__init__(**kwargs)
+ if empty_confirmations < 1:
+ raise ValueError("empty_confirmations must be at least 1")
+
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.empty_confirmations = empty_confirmations
self.deferrable = deferrable
+ self._consecutive_empty = 0
def execute(self, context: Context) -> None:
if not self.deferrable:
@@ -395,8 +408,11 @@ class
BigQueryStreamingBufferEmptySensor(BaseSensorOperator):
return
if self.poke(context=context):
return
+ self._defer_to_trigger(timeout=self.timeout)
+
+ def _defer_to_trigger(self, timeout: float) -> None:
self.defer(
- timeout=timedelta(seconds=self.timeout),
+ timeout=timedelta(seconds=timeout),
trigger=BigQueryStreamingBufferEmptyTrigger(
project_id=self.project_id,
dataset_id=self.dataset_id,
@@ -404,6 +420,7 @@ class
BigQueryStreamingBufferEmptySensor(BaseSensorOperator):
poll_interval=self.poke_interval,
gcp_conn_id=self.gcp_conn_id,
impersonation_chain=self.impersonation_chain,
+ empty_confirmations=self.empty_confirmations,
),
method_name="execute_complete",
)
@@ -432,4 +449,42 @@ class
BigQueryStreamingBufferEmptySensor(BaseSensorOperator):
table =
hook.get_client(project_id=self.project_id).get_table(table_ref)
except NotFound as err:
raise ValueError(f"Table {table_uri} not found") from err
- return table.streaming_buffer is None
+
+ if table.streaming_buffer is not None:
+ self._consecutive_empty = 0
+ return False
+
+ self._consecutive_empty += 1
+ self.log.info(
+ "Streaming buffer reported empty (%s/%s confirmations) for table:
%s",
+ self._consecutive_empty,
+ self.empty_confirmations,
+ table_uri,
+ )
+ if self._consecutive_empty >= self.empty_confirmations:
+ return True
+ if self.reschedule and not self.deferrable:
+ # A reschedule rebuilds the operator and Airflow 3 clears XComs on
each run start, so the
+ # confirmation count cannot survive it; the trigger repeats the
full confirmation count instead.
+ first_reschedule_date = self._get_first_reschedule_date(context)
+ elapsed = (
+ (timezone.utcnow() - first_reschedule_date).total_seconds() if
first_reschedule_date else 0
+ )
+ if self.timeout - elapsed > 0:
+ self._defer_to_trigger(timeout=self.timeout - elapsed)
+ return False
+
+ def _get_first_reschedule_date(self, context: Context) -> datetime | None:
+ if AIRFLOW_V_3_0_PLUS:
+ return context["ti"].get_first_reschedule_date(context)
+ # Same lookup as the Airflow 2.11 BaseSensorOperator.execute;
`_orig_start_date` only exists in Airflow 2.
+ from airflow.sensors.base import _orig_start_date # type:
ignore[attr-defined]
+
+ ti = context["ti"]
+ return _orig_start_date(
+ dag_id=ti.dag_id,
+ task_id=ti.task_id,
+ run_id=ti.run_id,
+ map_index=ti.map_index,
+ try_number=(ti.max_tries or 0) - (self.retries or 0) + 1,
+ )
diff --git
a/providers/google/src/airflow/providers/google/cloud/triggers/bigquery.py
b/providers/google/src/airflow/providers/google/cloud/triggers/bigquery.py
index 359caa3a07a..25b5fc529f2 100644
--- a/providers/google/src/airflow/providers/google/cloud/triggers/bigquery.py
+++ b/providers/google/src/airflow/providers/google/cloud/triggers/bigquery.py
@@ -902,6 +902,8 @@ class BigQueryStreamingBufferEmptyTrigger(BaseTrigger):
:param poll_interval: Seconds between polls.
:param impersonation_chain: Optional service account to impersonate, or a
chained list of accounts.
+ :param empty_confirmations: Number of consecutive empty polls required
+ before reporting the buffer empty. Must be at least 1.
"""
def __init__(
@@ -912,14 +914,18 @@ class BigQueryStreamingBufferEmptyTrigger(BaseTrigger):
gcp_conn_id: str,
poll_interval: float = 30.0,
impersonation_chain: str | Sequence[str] | None = None,
+ empty_confirmations: int = 2,
):
super().__init__()
+ if empty_confirmations < 1:
+ raise ValueError("empty_confirmations must be at least 1")
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.impersonation_chain = impersonation_chain
+ self.empty_confirmations = empty_confirmations
def serialize(self) -> tuple[str, dict[str, Any]]:
return (
@@ -931,6 +937,7 @@ class BigQueryStreamingBufferEmptyTrigger(BaseTrigger):
"gcp_conn_id": self.gcp_conn_id,
"poll_interval": self.poll_interval,
"impersonation_chain": self.impersonation_chain,
+ "empty_confirmations": self.empty_confirmations,
},
)
@@ -942,6 +949,12 @@ class BigQueryStreamingBufferEmptyTrigger(BaseTrigger):
async def run(self) -> AsyncIterator[TriggerEvent]:
table_uri = f"{self.project_id}:{self.dataset_id}.{self.table_id}"
+ # The ``streamingBuffer`` table metadata BigQuery returns is eventually
+ # consistent, so a single absent reading can be a false "empty" right
+ # after a streaming insert. Yield success only after
+ # ``empty_confirmations`` consecutive empty polls, each
``poll_interval``
+ # apart.
+ consecutive_empty = 0
try:
hook = self._get_async_hook()
async with ClientSession() as session:
@@ -955,11 +968,21 @@ class BigQueryStreamingBufferEmptyTrigger(BaseTrigger):
table_id=self.table_id,
)
if is_empty:
- message = f"Streaming buffer is empty for table:
{table_uri}"
- self.log.info(message)
- yield TriggerEvent({"status": "success", "message":
message})
- return
- self.log.info("Streaming buffer not empty, sleeping %ss",
self.poll_interval)
+ consecutive_empty += 1
+ if consecutive_empty >= self.empty_confirmations:
+ message = f"Streaming buffer is empty for table:
{table_uri}"
+ self.log.info(message)
+ yield TriggerEvent({"status": "success",
"message": message})
+ return
+ self.log.info(
+ "Streaming buffer reported empty (%s/%s
confirmations), sleeping %ss",
+ consecutive_empty,
+ self.empty_confirmations,
+ self.poll_interval,
+ )
+ else:
+ consecutive_empty = 0
+ self.log.info("Streaming buffer not empty, sleeping
%ss", self.poll_interval)
await asyncio.sleep(self.poll_interval)
except Exception as e:
self.log.exception("Error while checking streaming buffer for
table %s", table_uri)
diff --git
a/providers/google/tests/system/google/cloud/bigquery/example_bigquery_streaming_buffer_sensor.py
b/providers/google/tests/system/google/cloud/bigquery/example_bigquery_streaming_buffer_sensor.py
index 68035ac944e..3b07f5c5062 100644
---
a/providers/google/tests/system/google/cloud/bigquery/example_bigquery_streaming_buffer_sensor.py
+++
b/providers/google/tests/system/google/cloud/bigquery/example_bigquery_streaming_buffer_sensor.py
@@ -26,7 +26,6 @@ test can run for a long time end-to-end and is therefore
opt-in: set
from __future__ import annotations
import os
-import time
from datetime import datetime
import pytest
@@ -65,7 +64,7 @@ TABLE_NAME =
f"partitioned_table_{DAG_ID}_{ENV_ID}".replace("-", "_")
# DML on rows still in the streaming buffer is rejected by BigQuery, hence the
# sensor in the streaming-insert -> sensor -> DML chain below.
STREAMING_UPDATE_QUERY = f"UPDATE {DATASET_NAME}.{TABLE_NAME} SET value = 200
WHERE value = 100"
-STREAMING_DELETE_QUERY = f"DELETE FROM {DATASET_NAME}.{TABLE_NAME} WHERE value
= 200"
+STREAMING_DELETE_QUERY = f"DELETE FROM {DATASET_NAME}.{TABLE_NAME} WHERE value
IN (100, 200)"
SCHEMA = [
{"name": "value", "type": "INTEGER", "mode": "REQUIRED"},
@@ -109,21 +108,10 @@ with DAG(
rows=[{"value": 100, "ds": ds}],
fail_on_error=True,
)
- # BigQuery's streamingBuffer table metadata is eventually consistent:
for
- # a few seconds after a streaming insert the row is in the buffer but
- # table.streaming_buffer is still None. Wait for the metadata to catch
up
- # so check_streaming_buffer_empty does not falsely report "empty"
before
- # the buffer is reported at all. Remove once the sensor handles this
- # itself; tracked at https://github.com/apache/airflow/issues/66963
- client = hook.get_client(project_id=PROJECT_ID)
- table_uri = f"{PROJECT_ID}.{DATASET_NAME}.{TABLE_NAME}"
- for _ in range(30):
- if client.get_table(table_uri).streaming_buffer is not None:
- return
- time.sleep(2)
- raise RuntimeError("BigQuery streaming buffer metadata did not appear
within 60s")
streaming_insert_task = streaming_insert()
+ # A fresh insert right before the deferrable sensor, so it also sees the
just-inserted state.
+ streaming_insert_def_task =
streaming_insert.override(task_id="streaming_insert_def")()
# [START howto_sensor_bigquery_streaming_buffer_empty]
check_streaming_buffer_empty = BigQueryStreamingBufferEmptySensor(
@@ -181,6 +169,7 @@ with DAG(
>> streaming_insert_task
>> check_streaming_buffer_empty
>> stream_update
+ >> streaming_insert_def_task
>> check_streaming_buffer_empty_def
>> stream_delete
>> delete_dataset
diff --git a/providers/google/tests/unit/google/cloud/sensors/test_bigquery.py
b/providers/google/tests/unit/google/cloud/sensors/test_bigquery.py
index eb953d6af03..7435b55e215 100644
--- a/providers/google/tests/unit/google/cloud/sensors/test_bigquery.py
+++ b/providers/google/tests/unit/google/cloud/sensors/test_bigquery.py
@@ -16,13 +16,15 @@
# under the License.
from __future__ import annotations
+from datetime import UTC, datetime, timedelta
from unittest import mock
import pytest
+import time_machine
from google.api_core.exceptions import NotFound
from google.cloud.bigquery import TableReference
-from airflow.providers.common.compat.sdk import AirflowException, TaskDeferred
+from airflow.providers.common.compat.sdk import AirflowException,
AirflowSensorTimeout, TaskDeferred
from airflow.providers.google.cloud.sensors.bigquery import (
BigQueryRoutineExistenceSensor,
BigQueryStreamingBufferEmptySensor,
@@ -35,6 +37,8 @@ from airflow.providers.google.cloud.triggers.bigquery import (
BigQueryTablePartitionExistenceTrigger,
)
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
+
TEST_PROJECT_ID = "test_project"
TEST_DATASET_ID = "test_dataset"
TEST_TABLE_ID = "test_table"
@@ -350,6 +354,7 @@ class TestBigQueryStreamingBufferEmptySensor:
sensor = _make_streaming_sensor(
gcp_conn_id=TEST_GCP_CONN_ID,
impersonation_chain=TEST_IMPERSONATION_CHAIN,
+ empty_confirmations=1,
)
mock_table = mock.MagicMock(streaming_buffer=None)
mock_hook.return_value.get_client.return_value.get_table.return_value
= mock_table
@@ -380,6 +385,135 @@ class TestBigQueryStreamingBufferEmptySensor:
assert sensor.poke(mock.MagicMock()) is False
+ @pytest.mark.parametrize(
+ ("empty_confirmations", "expected_results"),
+ [
+ pytest.param(2, [False, True], id="default"),
+ pytest.param(3, [False, False, True], id="three_confirmations"),
+ ],
+ )
+ @mock.patch("airflow.providers.google.cloud.sensors.bigquery.BigQueryHook")
+ def test_poke_requires_consecutive_empty_confirmations(
+ self,
+ mock_hook,
+ empty_confirmations,
+ expected_results,
+ ):
+ sensor =
_make_streaming_sensor(empty_confirmations=empty_confirmations)
+ mock_hook.return_value.get_client.return_value.get_table.return_value
= mock.MagicMock(
+ streaming_buffer=None
+ )
+
+ # A single absent reading is ambiguous (metadata lag), so the sensor
only
+ # reports empty after the configured number of consecutive empty
readings.
+ for expected in expected_results:
+ assert sensor.poke(mock.MagicMock()) is expected
+
+ @mock.patch("airflow.providers.google.cloud.sensors.bigquery.BigQueryHook")
+ def test_poke_resets_confirmations_when_buffer_reappears(self, mock_hook):
+ sensor = _make_streaming_sensor(empty_confirmations=2)
+ get_table = mock_hook.return_value.get_client.return_value.get_table
+ get_table.side_effect = [
+ mock.MagicMock(streaming_buffer=None),
+ mock.MagicMock(streaming_buffer={"estimatedRows": 5}),
+ mock.MagicMock(streaming_buffer=None),
+ mock.MagicMock(streaming_buffer=None),
+ ]
+
+ assert sensor.poke(mock.MagicMock()) is False # 1st empty
+ assert sensor.poke(mock.MagicMock()) is False # buffer reappears,
counter resets
+ assert sensor.poke(mock.MagicMock()) is False # empty again, count
restarts at 1
+ assert sensor.poke(mock.MagicMock()) is True # 2nd consecutive empty
+
+ @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS,
reason="get_first_reschedule_date is Airflow 3 only")
+ @pytest.mark.parametrize(
+ ("elapsed", "expected_timeout"),
+ [
+ pytest.param(None, 100, id="first_try"),
+ pytest.param(30, 70, id="after_reschedules"),
+ ],
+ )
+ @mock.patch("airflow.providers.google.cloud.sensors.bigquery.BigQueryHook")
+ def
test_execute_defers_to_trigger_with_remaining_timeout_in_reschedule_mode(
+ self, mock_hook, elapsed, expected_timeout
+ ):
+ now = datetime(2026, 1, 1, tzinfo=UTC)
+ sensor = _make_streaming_sensor(
+ mode="reschedule", poke_interval=10, timeout=100,
empty_confirmations=3
+ )
+ mock_hook.return_value.get_client.return_value.get_table.return_value
= mock.MagicMock(
+ streaming_buffer=None
+ )
+ ti = mock.MagicMock(spec=["get_first_reschedule_date"])
+ ti.get_first_reschedule_date.return_value = (
+ None if elapsed is None else now - timedelta(seconds=elapsed)
+ )
+
+ with time_machine.travel(now, tick=False), pytest.raises(TaskDeferred)
as exc:
+ sensor.execute({"ti": ti})
+
+ trigger = exc.value.trigger
+ assert isinstance(trigger, BigQueryStreamingBufferEmptyTrigger)
+ assert trigger.empty_confirmations == 3
+ assert trigger.poll_interval == 10
+ assert exc.value.timeout == timedelta(seconds=expected_timeout)
+ assert exc.value.method_name == "execute_complete"
+
+ @pytest.mark.skipif(AIRFLOW_V_3_0_PLUS, reason="_orig_start_date only
exists in Airflow 2")
+ @mock.patch("airflow.sensors.base._orig_start_date", autospec=True)
+ @mock.patch("airflow.providers.google.cloud.sensors.bigquery.BigQueryHook")
+ def
test_poke_defers_with_remaining_timeout_in_reschedule_mode_on_airflow_2(
+ self, mock_hook, mock_orig_start_date
+ ):
+ now = datetime(2026, 1, 1, tzinfo=UTC)
+ mock_orig_start_date.return_value = now - timedelta(seconds=30)
+ sensor = _make_streaming_sensor(mode="reschedule", poke_interval=10,
timeout=100, retries=2)
+ mock_hook.return_value.get_client.return_value.get_table.return_value
= mock.MagicMock(
+ streaming_buffer=None
+ )
+ ti = mock.MagicMock(
+ spec=["dag_id", "task_id", "run_id", "map_index", "max_tries"],
+ dag_id="dag",
+ task_id="task-id",
+ run_id="run",
+ map_index=-1,
+ max_tries=3,
+ )
+
+ with time_machine.travel(now, tick=False), pytest.raises(TaskDeferred)
as exc:
+ sensor.poke({"ti": ti})
+
+ mock_orig_start_date.assert_called_once_with(
+ dag_id="dag", task_id="task-id", run_id="run", map_index=-1,
try_number=2
+ )
+ assert exc.value.timeout == timedelta(seconds=70)
+
+ @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS,
reason="get_first_reschedule_date is Airflow 3 only")
+ @mock.patch(
+
"airflow.providers.google.cloud.sensors.bigquery.BigQueryStreamingBufferEmptySensor.defer",
+ autospec=True,
+ )
+ @mock.patch("airflow.providers.google.cloud.sensors.bigquery.BigQueryHook")
+ def
test_execute_times_out_instead_of_deferring_when_reschedule_timeout_elapsed(
+ self, mock_hook, mock_defer
+ ):
+ sensor = _make_streaming_sensor(mode="reschedule", poke_interval=10,
timeout=100)
+ mock_hook.return_value.get_client.return_value.get_table.return_value
= mock.MagicMock(
+ streaming_buffer=None
+ )
+ now = datetime(2026, 1, 1, tzinfo=UTC)
+ ti = mock.MagicMock(spec=["get_first_reschedule_date"])
+ ti.get_first_reschedule_date.return_value = now -
timedelta(seconds=500)
+
+ with time_machine.travel(now, tick=False),
pytest.raises(AirflowSensorTimeout):
+ sensor.execute({"ti": ti})
+
+ mock_defer.assert_not_called()
+
+ def test_init_rejects_non_positive_empty_confirmations(self):
+ with pytest.raises(ValueError, match="empty_confirmations must be at
least 1"):
+ _make_streaming_sensor(empty_confirmations=0)
+
@mock.patch("airflow.providers.google.cloud.sensors.bigquery.BigQueryHook")
def test_poke_raises_value_error_when_table_not_found(self, mock_hook):
mock_hook.return_value.get_client.return_value.get_table.side_effect =
NotFound("missing")
@@ -401,7 +535,7 @@ class TestBigQueryStreamingBufferEmptySensor:
streaming_buffer=None
)
- _make_streaming_sensor(deferrable=True).execute(mock.MagicMock())
+ _make_streaming_sensor(deferrable=True,
empty_confirmations=1).execute(mock.MagicMock())
mock_defer.assert_not_called()
@@ -411,6 +545,7 @@ class TestBigQueryStreamingBufferEmptySensor:
gcp_conn_id=TEST_GCP_CONN_ID,
impersonation_chain=TEST_IMPERSONATION_CHAIN,
deferrable=True,
+ empty_confirmations=3,
)
mock_hook.return_value.get_client.return_value.get_table.return_value
= mock.MagicMock(
streaming_buffer={"estimatedRows": 1}
@@ -428,6 +563,7 @@ class TestBigQueryStreamingBufferEmptySensor:
assert trigger.project_id == TEST_PROJECT_ID
assert trigger.dataset_id == TEST_DATASET_ID
assert trigger.table_id == TEST_TABLE_ID
+ assert trigger.empty_confirmations == 3
def test_execute_complete_returns_message_on_success(self):
sensor = _make_streaming_sensor(deferrable=True)
diff --git a/providers/google/tests/unit/google/cloud/triggers/test_bigquery.py
b/providers/google/tests/unit/google/cloud/triggers/test_bigquery.py
index 2c23d3050e0..27e93522ed9 100644
--- a/providers/google/tests/unit/google/cloud/triggers/test_bigquery.py
+++ b/providers/google/tests/unit/google/cloud/triggers/test_bigquery.py
@@ -1178,8 +1178,19 @@ class TestBigQueryStreamingBufferEmptyTrigger:
"gcp_conn_id": TEST_GCP_CONN_ID,
"poll_interval": POLLING_PERIOD_SECONDS,
"impersonation_chain": TEST_IMPERSONATION_CHAIN,
+ "empty_confirmations": 2,
}
+ def test_init_rejects_non_positive_empty_confirmations(self):
+ with pytest.raises(ValueError, match="empty_confirmations must be at
least 1"):
+ BigQueryStreamingBufferEmptyTrigger(
+ project_id=TEST_GCP_PROJECT_ID,
+ dataset_id=TEST_DATASET_ID,
+ table_id=TEST_TABLE_ID,
+ gcp_conn_id=TEST_GCP_CONN_ID,
+ empty_confirmations=0,
+ )
+
@mock.patch("airflow.providers.google.cloud.triggers.bigquery.BigQueryTableAsyncHook")
def test_async_hook_receives_impersonation_chain(self, mock_hook_cls,
streaming_buffer_trigger):
streaming_buffer_trigger._get_async_hook()
@@ -1191,12 +1202,77 @@ class TestBigQueryStreamingBufferEmptyTrigger:
@pytest.mark.asyncio
@mock.patch(f"{_TRIGGER_PATH}._is_streaming_buffer_empty")
@mock.patch(f"{_TRIGGER_PATH}._get_async_hook")
- async def test_run_yields_success_when_buffer_empty(
- self, _mock_hook, mock_is_empty, streaming_buffer_trigger
- ):
+ async def test_run_yields_success_when_buffer_empty(self, _mock_hook,
mock_is_empty):
+ trigger = BigQueryStreamingBufferEmptyTrigger(
+ 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,
+ impersonation_chain=TEST_IMPERSONATION_CHAIN,
+ empty_confirmations=1,
+ )
mock_is_empty.return_value = True
+ actual = await trigger.run().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
+ @pytest.mark.parametrize(
+ ("empty_confirmations", "side_effect", "expected_polls"),
+ [
+ pytest.param(2, [True, True], 2, id="default"),
+ pytest.param(3, [True, True, True], 3, id="three_confirmations"),
+ ],
+ )
+
@mock.patch("airflow.providers.google.cloud.triggers.bigquery.asyncio.sleep",
new_callable=AsyncMock)
+ @mock.patch(f"{_TRIGGER_PATH}._is_streaming_buffer_empty")
+ @mock.patch(f"{_TRIGGER_PATH}._get_async_hook")
+ async def test_run_waits_for_consecutive_empty_confirmations(
+ self,
+ _mock_hook,
+ mock_is_empty,
+ mock_sleep,
+ empty_confirmations,
+ side_effect,
+ expected_polls,
+ ):
+ # A single empty reading must not yield success; success only after
+ # ``empty_confirmations`` consecutive empty polls.
+ mock_is_empty.side_effect = side_effect
+
+ trigger = BigQueryStreamingBufferEmptyTrigger(
+ project_id=TEST_GCP_PROJECT_ID,
+ dataset_id=TEST_DATASET_ID,
+ table_id=TEST_TABLE_ID,
+ gcp_conn_id=TEST_GCP_CONN_ID,
+ empty_confirmations=empty_confirmations,
+ )
+
+ actual = await trigger.run().asend(None)
+
+ assert mock_is_empty.await_count == expected_polls
+
+ 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.asyncio.sleep",
new_callable=AsyncMock)
+ @mock.patch(f"{_TRIGGER_PATH}._is_streaming_buffer_empty")
+ @mock.patch(f"{_TRIGGER_PATH}._get_async_hook")
+ async def test_run_resets_confirmations_when_buffer_reappears(
+ self, _mock_hook, mock_is_empty, mock_sleep, streaming_buffer_trigger
+ ):
+ # empty, then non-empty (resets), then two consecutive empties ->
success.
+ mock_is_empty.side_effect = [True, False, True, True]
actual = await streaming_buffer_trigger.run().asend(None)
+ assert mock_is_empty.await_count == 4
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}"}
