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..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 @@ -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,31 +180,32 @@ 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 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] - 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 +214,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