From f0e3df26fec7ac7025748a7bb286143fb33ad622 Mon Sep 17 00:00:00 2001 From: 1fanwang <1fannnw@gmail.com> Date: Thu, 27 Aug 2026 16:04:54 -0400 Subject: [PATCH 1/2] Use async asset state access in Iceberg trigger The trigger already runs on the event loop, so synchronous state access needlessly consumes worker threads and blocks async-only state stores. Signed-off-by: 1fanwang <1fannnw@gmail.com> --- .../apache/iceberg/triggers/iceberg.py | 4 ++-- .../apache/iceberg/triggers/test_iceberg.py | 18 ++++++++++-------- 2 files changed, 12 insertions(+), 10 deletions(-) diff --git a/providers/apache/iceberg/src/airflow/providers/apache/iceberg/triggers/iceberg.py b/providers/apache/iceberg/src/airflow/providers/apache/iceberg/triggers/iceberg.py index 6777982eefb5a..34e885584bee0 100644 --- a/providers/apache/iceberg/src/airflow/providers/apache/iceberg/triggers/iceberg.py +++ b/providers/apache/iceberg/src/airflow/providers/apache/iceberg/triggers/iceberg.py @@ -132,7 +132,7 @@ async def run(self) -> AsyncIterator[TriggerEvent]: store = getattr(self, "asset_state_store", None) if store is not None: try: - stored = await asyncio.to_thread(store.get, WATERMARK_KEY) + stored = await store.aget(WATERMARK_KEY) except ValueError as err: # The accessor serves one asset at a time, so it refuses to guess when this # trigger is watched by several. That happens because triggers are deduplicated @@ -158,7 +158,7 @@ async def run(self) -> AsyncIterator[TriggerEvent]: if head is not None and head != self.last_seen_snapshot_id: previous, self.last_seen_snapshot_id = self.last_seen_snapshot_id, head if store is not None: - await asyncio.to_thread(store.set, WATERMARK_KEY, head) + await store.aset(WATERMARK_KEY, head) yield TriggerEvent( { "table": self.table, diff --git a/providers/apache/iceberg/tests/unit/apache/iceberg/triggers/test_iceberg.py b/providers/apache/iceberg/tests/unit/apache/iceberg/triggers/test_iceberg.py index 15877d3d9d0b2..4ca24392000dc 100644 --- a/providers/apache/iceberg/tests/unit/apache/iceberg/triggers/test_iceberg.py +++ b/providers/apache/iceberg/tests/unit/apache/iceberg/triggers/test_iceberg.py @@ -19,7 +19,7 @@ import asyncio from contextlib import aclosing, suppress from typing import TYPE_CHECKING -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest from pyiceberg.exceptions import NoSuchNamespaceError, NoSuchTableError @@ -171,7 +171,7 @@ async def test_resumes_from_the_stored_watermark(): row was written; only the stored watermark reflects what was actually emitted. """ store = MagicMock() - store.get.return_value = 222 + store.aget = AsyncMock(return_value=222) trigger = IcebergTableSnapshotTrigger(table="db.tbl", poll_interval=0.01, last_seen_snapshot_id=111) trigger.asset_state_store = store @@ -180,13 +180,14 @@ async def test_resumes_from_the_stored_watermark(): payloads = await _collect(trigger, 1, timeout=0.2) assert payloads == [] - store.get.assert_called_once_with("snapshot_id") + store.aget.assert_awaited_once_with("snapshot_id") @pytest.mark.asyncio async def test_persists_the_watermark_on_each_event(): store = MagicMock() - store.get.return_value = None + store.aget = AsyncMock(return_value=None) + store.aset = AsyncMock() trigger = IcebergTableSnapshotTrigger(table="db.tbl", poll_interval=0.01) trigger.asset_state_store = store @@ -197,14 +198,15 @@ async def test_persists_the_watermark_on_each_event(): payloads = await _collect(trigger, 2, timeout=3.0) assert [p["snapshot_id"] for p in payloads] == [111, 222] - assert [c.args for c in store.set.call_args_list] == [("snapshot_id", 111), ("snapshot_id", 222)] + assert [c.args for c in store.aset.await_args_list] == [("snapshot_id", 111), ("snapshot_id", 222)] @pytest.mark.asyncio async def test_runs_without_a_watermark_when_several_assets_watch_it(): """More than one watched asset leaves no single cursor, so it degrades instead of raising.""" store = MagicMock() - store.get.side_effect = ValueError("Task has 2 concrete inlets and outlets") + store.aget = AsyncMock(side_effect=ValueError("Task has 2 concrete inlets and outlets")) + store.aset = AsyncMock() trigger = IcebergTableSnapshotTrigger(table="db.tbl", poll_interval=0.01) trigger.asset_state_store = store @@ -213,14 +215,14 @@ async def test_runs_without_a_watermark_when_several_assets_watch_it(): payloads = await _collect(trigger, 1) assert [p["snapshot_id"] for p in payloads] == [111] - store.set.assert_not_called() + store.aset.assert_not_awaited() @pytest.mark.asyncio async def test_a_state_store_failure_is_not_mistaken_for_several_assets(): """A pluggable backend can raise ValueError too, and hiding it would disable the watermark.""" store = MagicMock() - store.get.side_effect = ValueError("could not decode the stored reference") + store.aget = AsyncMock(side_effect=ValueError("could not decode the stored reference")) trigger = IcebergTableSnapshotTrigger(table="db.tbl", poll_interval=0.01) trigger.asset_state_store = store From 0f05b196a494bfaa1507e040ed7e856297e66aa0 Mon Sep 17 00:00:00 2001 From: 1fanwang <1fannnw@gmail.com> Date: Thu, 27 Aug 2026 16:19:52 -0400 Subject: [PATCH 2/2] Keep Iceberg test timeout rationale accurate The previous wording still counted state-store calls as thread-pool work after the async migration. Signed-off-by: 1fanwang <1fannnw@gmail.com> --- .../iceberg/tests/unit/apache/iceberg/triggers/test_iceberg.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/providers/apache/iceberg/tests/unit/apache/iceberg/triggers/test_iceberg.py b/providers/apache/iceberg/tests/unit/apache/iceberg/triggers/test_iceberg.py index 4ca24392000dc..2099d2767ede5 100644 --- a/providers/apache/iceberg/tests/unit/apache/iceberg/triggers/test_iceberg.py +++ b/providers/apache/iceberg/tests/unit/apache/iceberg/triggers/test_iceberg.py @@ -193,8 +193,7 @@ async def test_persists_the_watermark_on_each_event(): trigger.asset_state_store = store with patch(LOAD_TABLE, side_effect=[_table_at(111), _table_at(222), _table_at(222)]): - # Gathering 2 events runs several real asyncio.to_thread calls (head lookup + store - # get/set); the default 1s budget is too tight under CI thread-pool scheduling latency. + # Multiple real thread-pool head lookups can exceed the default 1s budget under CI latency. payloads = await _collect(trigger, 2, timeout=3.0) assert [p["snapshot_id"] for p in payloads] == [111, 222]