SameerMesiah97 commented on code in PR #69118:
URL: https://github.com/apache/airflow/pull/69118#discussion_r3501983398
##########
providers/google/tests/unit/google/cloud/sensors/test_bigquery.py:
##########
@@ -343,6 +344,38 @@ def test_poke_returns_false_when_buffer_present(self,
mock_hook):
assert sensor.poke(mock.MagicMock()) is False
+ @mock.patch("airflow.providers.google.cloud.sensors.bigquery.BigQueryHook")
+ def test_poke_requires_consecutive_empty_confirmations(self, mock_hook):
+ sensor = _make_streaming_sensor(empty_confirmations=2)
+ 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 it must not
+ # report empty; only the second consecutive empty reading does.
+ assert sensor.poke(mock.MagicMock()) is False
+ assert sensor.poke(mock.MagicMock()) is True
Review Comment:
I would parametrize this test to cover non-default configurations for
`empty_confirmations`:
```
@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
```
##########
providers/google/src/airflow/providers/google/cloud/sensors/bigquery.py:
##########
@@ -333,20 +333,26 @@ class
BigQueryStreamingBufferEmptySensor(BaseSensorOperator):
``UPDATE/MERGE/DELETE statement over table ... would affect rows in the
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
+ The ``table.streaming_buffer`` metadata BigQuery exposes is eventually
+ consistent: for a short window right after a streaming insert the rows are
+ in the buffer but the metadata still reads absent. A single absent reading
+ is therefore ambiguous (truly empty vs. metadata lag), so the sensor only
+ reports empty after ``empty_confirmations`` consecutive empty readings,
each
+ one ``poke_interval`` apart. This spans the eventual-consistency window
+ without hanging when the table is genuinely empty or the buffer flushes
+ between two pokes.
Review Comment:
I would keep the warning and change the this part of the docstring to this:
```
.. warning::
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.
```
We want to document suprising behaviour here.
##########
providers/google/src/airflow/providers/google/cloud/sensors/bigquery.py:
##########
@@ -426,4 +439,16 @@ def poke(self, context: Context) -> bool:
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,
+ )
+ return self._consecutive_empty >= self.empty_confirmations
Review Comment:
nit: this could be more explicit:
```
if self._consecutive_empty >= self.empty_confirmations:
return True
return False
```
##########
providers/google/src/airflow/providers/google/cloud/sensors/bigquery.py:
##########
@@ -333,20 +333,26 @@ class
BigQueryStreamingBufferEmptySensor(BaseSensorOperator):
``UPDATE/MERGE/DELETE statement over table ... would affect rows in the
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
+ The ``table.streaming_buffer`` metadata BigQuery exposes is eventually
+ consistent: for a short window right after a streaming insert the rows are
+ in the buffer but the metadata still reads absent. A single absent reading
+ is therefore ambiguous (truly empty vs. metadata lag), so the sensor only
+ reports empty after ``empty_confirmations`` consecutive empty readings,
each
+ one ``poke_interval`` apart. This spans the eventual-consistency window
+ without hanging when the table is genuinely empty or the buffer flushes
+ between two pokes.
:param project_id: Google Cloud project containing the table.
:param dataset_id: Dataset of the table to monitor.
:param table_id: Table to monitor.
: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 (each
+ ``poke_interval`` apart) required before reporting the buffer empty.
+ Must be at least 1; values above 1 guard against BigQuery's
+ eventually-consistent streaming-buffer metadata reporting empty too
+ early after a streaming insert.
Review Comment:
This docstring is too long. I would do something like this instead:
```
:param empty_confirmations: Number of consecutive empty readings required
before considering the streaming buffer empty. Must be at least 1.
Defaults to 2.
```
It is good to document the defaults for users.
##########
providers/google/src/airflow/providers/google/cloud/triggers/bigquery.py:
##########
@@ -891,13 +891,21 @@ class BigQueryStreamingBufferEmptyTrigger(BaseTrigger):
Used by
:class:`~airflow.providers.google.cloud.sensors.bigquery.BigQueryStreamingBufferEmptySensor`
in deferrable mode.
+ The ``streamingBuffer`` table metadata BigQuery returns is eventually
+ consistent, so a single absent reading can be a false "empty" right after a
+ streaming insert. The trigger therefore yields success only after
+ ``empty_confirmations`` consecutive empty polls, each ``poll_interval``
+ apart.
Review Comment:
This part you added belongs in the comment near the modifications you have
made in the `run` method.
##########
providers/google/tests/unit/google/cloud/triggers/test_bigquery.py:
##########
@@ -1070,12 +1081,54 @@ def test_async_hook_receives_impersonation_chain(self,
mock_hook_cls, streaming_
@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
+
@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, streaming_buffer_trigger
+ ):
+ # A single empty reading must not yield success (default
empty_confirmations=2);
+ # success only after the second consecutive empty poll.
+ mock_is_empty.side_effect = [True, True]
+ actual = await streaming_buffer_trigger.run().asend(None)
+
+ assert mock_is_empty.await_count == 2
+ 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}"}
Review Comment:
I would parametrize this test to cover non-default configurations for
`empty_confirmations`:
```
@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,
):
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}"}
)
```
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
