diff --git a/providers/apache/iceberg/provider.yaml b/providers/apache/iceberg/provider.yaml index 913a22a0ff105..f45eee2d769d2 100644 --- a/providers/apache/iceberg/provider.yaml +++ b/providers/apache/iceberg/provider.yaml @@ -56,6 +56,11 @@ hooks: python-modules: - airflow.providers.apache.iceberg.hooks.iceberg +triggers: + - integration-name: Iceberg + python-modules: + - airflow.providers.apache.iceberg.triggers.iceberg + connection-types: - hook-class-name: airflow.providers.apache.iceberg.hooks.iceberg.IcebergHook hook-name: "Iceberg" diff --git a/providers/apache/iceberg/src/airflow/providers/apache/iceberg/get_provider_info.py b/providers/apache/iceberg/src/airflow/providers/apache/iceberg/get_provider_info.py index df1285577b9cc..900a8c4b167c7 100644 --- a/providers/apache/iceberg/src/airflow/providers/apache/iceberg/get_provider_info.py +++ b/providers/apache/iceberg/src/airflow/providers/apache/iceberg/get_provider_info.py @@ -40,6 +40,12 @@ def get_provider_info(): "python-modules": ["airflow.providers.apache.iceberg.hooks.iceberg"], } ], + "triggers": [ + { + "integration-name": "Iceberg", + "python-modules": ["airflow.providers.apache.iceberg.triggers.iceberg"], + } + ], "connection-types": [ { "hook-class-name": "airflow.providers.apache.iceberg.hooks.iceberg.IcebergHook", diff --git a/providers/apache/iceberg/src/airflow/providers/apache/iceberg/triggers/__init__.py b/providers/apache/iceberg/src/airflow/providers/apache/iceberg/triggers/__init__.py new file mode 100644 index 0000000000000..217e5db960782 --- /dev/null +++ b/providers/apache/iceberg/src/airflow/providers/apache/iceberg/triggers/__init__.py @@ -0,0 +1,17 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. 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 new file mode 100644 index 0000000000000..6777982eefb5a --- /dev/null +++ b/providers/apache/iceberg/src/airflow/providers/apache/iceberg/triggers/iceberg.py @@ -0,0 +1,170 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import asyncio +from typing import TYPE_CHECKING, Any + +from pyiceberg.exceptions import NoSuchNamespaceError, NoSuchTableError + +from airflow.providers.apache.iceberg.hooks.iceberg import IcebergHook +from airflow.providers.apache.iceberg.version_compat import AIRFLOW_V_3_0_PLUS + +if AIRFLOW_V_3_0_PLUS: + from airflow.triggers.base import BaseEventTrigger, TriggerEvent +else: + from airflow.triggers.base import ( # type: ignore[assignment] + BaseTrigger as BaseEventTrigger, + TriggerEvent, + ) + +if TYPE_CHECKING: + from collections.abc import AsyncIterator + +DEFAULT_BRANCH = "main" +WATERMARK_KEY = "snapshot_id" +# Raised by AssetStateStoreAccessors when the trigger is watched by more than one asset. +_MULTI_ASSET_ERROR = "concrete inlets and outlets" + + +class IcebergTableSnapshotTrigger(BaseEventTrigger): + """ + Fire an event whenever an Iceberg table gains a new snapshot. + + Polls the branch head through the catalog and emits an event when it points at a + snapshot the trigger has not reported yet, which makes a table commit usable as a + scheduling signal:: + + from airflow.sdk import Asset, AssetWatcher + + orders = Asset( + "orders", + watchers=[ + AssetWatcher( + name="orders_commits", + trigger=IcebergTableSnapshotTrigger(table="sales.orders"), + ) + ], + ) + + + @dag(schedule=[orders]) + def downstream(): ... + + The event carries ``table``, ``branch``, ``snapshot_id`` and ``previous_snapshot_id``, so a + task can scan the delta rather than the whole table. ``previous_snapshot_id`` is ``None`` on + the first event, which is also what a restart looks like where no watermark is kept, so a + task that must not run twice for one snapshot keys on ``snapshot_id``. + + :param table: Fully-qualified table name, ``namespace.table``. Nested namespaces are + written as ``a.b.table``. + :param iceberg_conn_id: Connection holding the catalog URI and credentials. + :param branch: Branch or tag to watch. Defaults to ``main``. + :param poll_interval: Seconds between polls. + :param last_seen_snapshot_id: Snapshot already reported. Leave unset to treat the + current head as the first event. + """ + + def __init__( + self, + *, + table: str, + iceberg_conn_id: str = IcebergHook.default_conn_name, + branch: str = DEFAULT_BRANCH, + poll_interval: float = 60, + last_seen_snapshot_id: int | None = None, + **kwargs: Any, + ) -> None: + super().__init__(**kwargs) + if "." not in table: + raise ValueError(f"Expected a fully-qualified table name (namespace.table), got: {table!r}") + self.table = table + self.iceberg_conn_id = iceberg_conn_id + self.branch = branch + self.poll_interval = poll_interval + self.last_seen_snapshot_id = last_seen_snapshot_id + + def serialize(self) -> tuple[str, dict[str, Any]]: + return ( + "airflow.providers.apache.iceberg.triggers.iceberg.IcebergTableSnapshotTrigger", + { + "table": self.table, + "iceberg_conn_id": self.iceberg_conn_id, + "branch": self.branch, + "poll_interval": self.poll_interval, + "last_seen_snapshot_id": self.last_seen_snapshot_id, + }, + ) + + def _head_snapshot_id(self) -> int | None: + """Return the snapshot the branch points at, or None if there is nothing to watch yet.""" + try: + table = IcebergHook(self.iceberg_conn_id).load_table(self.table) + except (NoSuchTableError, NoSuchNamespaceError): + # A watcher outlives the table it watches, and raising here would kill the trigger + # and have the triggerer restart it once per second until the table appears. + self.log.debug("%s does not exist yet; waiting", self.table) + return None + if ref := table.metadata.refs.get(self.branch): + return ref.snapshot_id + return None + + async def run(self) -> AsyncIterator[TriggerEvent]: + # serialize() is captured once when the trigger row is created, so a value mutated on + # self is lost when the triggerer restarts and the current head would be re-emitted as + # a new commit. The watermark survives that; the kwarg only seeds the first run. + # It postdates the Airflow versions this provider supports and is absent when several + # assets share the trigger, so polling carries on without it, losing only that cursor. + store = getattr(self, "asset_state_store", None) + if store is not None: + try: + stored = await asyncio.to_thread(store.get, 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 + # by hash(classpath, kwargs) while asset_watcher is many-to-many, so two assets + # watching this table with the same arguments share one trigger. There is then + # no single cursor to keep. A state store backend can raise ValueError too, and + # swallowing that would disable the watermark without saying so. + if _MULTI_ASSET_ERROR not in str(err): + raise + self.log.warning( + "%s is watched by more than one asset; not persisting a snapshot watermark, " + "so a triggerer restart may re-emit the current head.", + self.table, + ) + store = None + else: + if stored is not None: + self.last_seen_snapshot_id = int(stored) + + while True: + # pyiceberg is synchronous, so keep the catalog call off the event loop. + head = await asyncio.to_thread(self._head_snapshot_id) + 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) + yield TriggerEvent( + { + "table": self.table, + "branch": self.branch, + "snapshot_id": head, + "previous_snapshot_id": previous, + } + ) + await asyncio.sleep(self.poll_interval) diff --git a/providers/apache/iceberg/src/airflow/providers/apache/iceberg/version_compat.py b/providers/apache/iceberg/src/airflow/providers/apache/iceberg/version_compat.py index b05a7ff4d6b51..613a946fc90b5 100644 --- a/providers/apache/iceberg/src/airflow/providers/apache/iceberg/version_compat.py +++ b/providers/apache/iceberg/src/airflow/providers/apache/iceberg/version_compat.py @@ -32,8 +32,10 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: return airflow_version.major, airflow_version.minor, airflow_version.micro +AIRFLOW_V_3_0_PLUS: bool = get_base_airflow_version_tuple() >= (3, 0, 0) AIRFLOW_V_3_1_PLUS: bool = get_base_airflow_version_tuple() >= (3, 1, 0) __all__ = [ + "AIRFLOW_V_3_0_PLUS", "AIRFLOW_V_3_1_PLUS", ] diff --git a/providers/apache/iceberg/tests/system/apache/iceberg/example_iceberg_table_snapshot_watcher.py b/providers/apache/iceberg/tests/system/apache/iceberg/example_iceberg_table_snapshot_watcher.py new file mode 100644 index 0000000000000..dd212e451f0fa --- /dev/null +++ b/providers/apache/iceberg/tests/system/apache/iceberg/example_iceberg_table_snapshot_watcher.py @@ -0,0 +1,36 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +# [START howto_trigger_iceberg_table_snapshot] +from airflow.providers.apache.iceberg.triggers.iceberg import IcebergTableSnapshotTrigger +from airflow.providers.standard.operators.empty import EmptyOperator +from airflow.sdk import DAG, Asset, AssetWatcher + +trigger = IcebergTableSnapshotTrigger(table="default.orders", poll_interval=30) + +orders = Asset("iceberg_orders", watchers=[AssetWatcher(name="orders_commits", trigger=trigger)]) + +with DAG(dag_id="example_iceberg_table_snapshot_watcher", schedule=[orders]) as dag: + EmptyOperator(task_id="process_new_orders") +# [END howto_trigger_iceberg_table_snapshot] + + +from tests_common.test_utils.system_tests import get_test_run # noqa: E402 + +# Needed to run the example DAG with pytest (see: contributing-docs/testing/system_tests.rst) +test_run = get_test_run(dag) diff --git a/providers/apache/iceberg/tests/unit/apache/iceberg/triggers/__init__.py b/providers/apache/iceberg/tests/unit/apache/iceberg/triggers/__init__.py new file mode 100644 index 0000000000000..217e5db960782 --- /dev/null +++ b/providers/apache/iceberg/tests/unit/apache/iceberg/triggers/__init__.py @@ -0,0 +1,17 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. 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 new file mode 100644 index 0000000000000..1f44dac375cb8 --- /dev/null +++ b/providers/apache/iceberg/tests/unit/apache/iceberg/triggers/test_iceberg.py @@ -0,0 +1,275 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import asyncio +from contextlib import aclosing, suppress +from typing import TYPE_CHECKING +from unittest.mock import MagicMock, patch + +import pytest +from pyiceberg.exceptions import NoSuchNamespaceError, NoSuchTableError + +from airflow.providers.apache.iceberg.triggers.iceberg import IcebergTableSnapshotTrigger + +if TYPE_CHECKING: + from collections.abc import AsyncGenerator + +LOAD_TABLE = "airflow.providers.apache.iceberg.hooks.iceberg.IcebergHook.load_table" + +TRIGGER_PATH = "airflow.providers.apache.iceberg.triggers.iceberg.IcebergTableSnapshotTrigger" + + +def _table_at(snapshot_id: int | None, branch: str = "main") -> MagicMock: + """Build a table whose ``branch`` points at ``snapshot_id`` (None means no such ref).""" + table = MagicMock() + table.metadata.refs = {branch: MagicMock(snapshot_id=snapshot_id)} if snapshot_id is not None else {} + return table + + +async def _collect(trigger: IcebergTableSnapshotTrigger, count: int, timeout: float = 1.0) -> list[dict]: + """Pull up to ``count`` payloads off the trigger, giving up after ``timeout``.""" + payloads: list[dict] = [] + generator: AsyncGenerator = trigger.run() # type: ignore[assignment] + + async with aclosing(generator): + + async def pump() -> None: + async for event in generator: + payloads.append(event.payload) + if len(payloads) >= count: + return + + with suppress(asyncio.TimeoutError): + await asyncio.wait_for(pump(), timeout=timeout) + return payloads + + +def test_serialize_round_trip(): + trigger = IcebergTableSnapshotTrigger( + table="db.tbl", iceberg_conn_id="my_conn", branch="audit", poll_interval=5 + ) + classpath, kwargs = trigger.serialize() + + assert classpath == TRIGGER_PATH + assert kwargs == { + "table": "db.tbl", + "iceberg_conn_id": "my_conn", + "branch": "audit", + "poll_interval": 5, + "last_seen_snapshot_id": None, + } + assert IcebergTableSnapshotTrigger(**kwargs).serialize() == (classpath, kwargs) + + +def test_rejects_table_without_namespace(): + with pytest.raises(ValueError, match="fully-qualified table name"): + IcebergTableSnapshotTrigger(table="orders") + + +async def _never(): + """An empty shared stream, as the triggerer would hand to filter_shared_stream.""" + return + yield # pragma: no cover + + +def _shared_stream_key(trigger): + return getattr(trigger, "shared_stream_key", lambda: None)() + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"table": "db.tbl"}, id="defaults"), + pytest.param({"table": "db.tbl", "branch": "audit"}, id="branch"), + pytest.param({"table": "db.tbl", "poll_interval": 5}, id="interval"), + ], +) +def test_the_triggerer_reaches_run(kwargs): + """A non-None key sends the triggerer to filter_shared_stream instead of run(). + + Declaring one without also implementing open_shared_stream and filter_shared_stream + means the trigger raises NotImplementedError and never polls. getattr because + shared streams postdate the Airflow versions this provider supports. + """ + assert _shared_stream_key(IcebergTableSnapshotTrigger(**kwargs)) is None + + +@pytest.mark.asyncio +async def test_emits_current_head_on_first_poll(): + """With no watermark the current head is itself the first event.""" + with patch(LOAD_TABLE, return_value=_table_at(111)): + payloads = await _collect(IcebergTableSnapshotTrigger(table="db.tbl", poll_interval=0.01), 1) + + assert payloads == [ + {"table": "db.tbl", "branch": "main", "snapshot_id": 111, "previous_snapshot_id": None} + ] + + +@pytest.mark.asyncio +async def test_silent_while_head_is_unchanged(): + """A table that has not committed must not schedule anything.""" + with patch(LOAD_TABLE, return_value=_table_at(111)): + trigger = IcebergTableSnapshotTrigger(table="db.tbl", poll_interval=0.01, last_seen_snapshot_id=111) + payloads = await _collect(trigger, 1, timeout=0.2) + + assert payloads == [] + + +@pytest.mark.asyncio +async def test_emits_once_per_new_snapshot(): + """Each commit produces exactly one event carrying the snapshot it replaced.""" + with patch(LOAD_TABLE, side_effect=[_table_at(111), _table_at(222), _table_at(222), _table_at(333)]): + trigger = IcebergTableSnapshotTrigger(table="db.tbl", poll_interval=0.01, last_seen_snapshot_id=111) + payloads = await _collect(trigger, 2) + + assert [(p["previous_snapshot_id"], p["snapshot_id"]) for p in payloads] == [(111, 222), (222, 333)] + + +@pytest.mark.asyncio +async def test_silent_while_branch_is_absent(): + """Watching a branch that does not exist yet waits for it rather than failing.""" + with patch(LOAD_TABLE, return_value=_table_at(None)): + trigger = IcebergTableSnapshotTrigger(table="db.tbl", branch="audit", poll_interval=0.01) + payloads = await _collect(trigger, 1, timeout=0.2) + + assert payloads == [] + + +@pytest.mark.asyncio +async def test_watches_the_requested_branch(): + """The event reports the branch that was asked for, not main.""" + with patch(LOAD_TABLE, return_value=_table_at(999, branch="audit")): + trigger = IcebergTableSnapshotTrigger(table="db.tbl", branch="audit", poll_interval=0.01) + payloads = await _collect(trigger, 1) + + assert payloads[0]["branch"] == "audit" + assert payloads[0]["snapshot_id"] == 999 + + +@pytest.mark.asyncio +async def test_resumes_from_the_stored_watermark(): + """A restarted triggerer must not re-emit a snapshot it already reported. + + ``serialize()`` is captured once, so the kwarg still holds the value from when the trigger + row was written; only the stored watermark reflects what was actually emitted. + """ + store = MagicMock() + store.get.return_value = 222 + + trigger = IcebergTableSnapshotTrigger(table="db.tbl", poll_interval=0.01, last_seen_snapshot_id=111) + trigger.asset_state_store = store + + with patch(LOAD_TABLE, return_value=_table_at(222)): + payloads = await _collect(trigger, 1, timeout=0.2) + + assert payloads == [] + store.get.assert_called_once_with("snapshot_id") + + +@pytest.mark.asyncio +async def test_persists_the_watermark_on_each_event(): + store = MagicMock() + store.get.return_value = None + + 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)]): + payloads = await _collect(trigger, 2) + + 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)] + + +@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") + + trigger = IcebergTableSnapshotTrigger(table="db.tbl", poll_interval=0.01) + trigger.asset_state_store = store + + with patch(LOAD_TABLE, return_value=_table_at(111)): + payloads = await _collect(trigger, 1) + + assert [p["snapshot_id"] for p in payloads] == [111] + store.set.assert_not_called() + + +@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") + + trigger = IcebergTableSnapshotTrigger(table="db.tbl", poll_interval=0.01) + trigger.asset_state_store = store + + with patch(LOAD_TABLE, return_value=_table_at(111)): + with pytest.raises(ValueError, match="could not decode"): + await _collect(trigger, 1) + + +@pytest.mark.asyncio +async def test_runs_on_airflow_without_an_asset_state_store(): + """``asset_state_store`` postdates the oldest Airflow this provider supports.""" + trigger = IcebergTableSnapshotTrigger(table="db.tbl", poll_interval=0.01) + if hasattr(trigger, "asset_state_store"): + del trigger.asset_state_store + assert not hasattr(trigger, "asset_state_store") + + with patch(LOAD_TABLE, return_value=_table_at(111)): + payloads = await _collect(trigger, 1) + + assert [p["snapshot_id"] for p in payloads] == [111] + + +@pytest.mark.asyncio +async def test_the_triggerer_dispatch_polls_the_table(): + """Drive the branch triggerer_job_runner takes, rather than calling run() directly.""" + trigger = IcebergTableSnapshotTrigger(table="db.tbl", poll_interval=0.01) + + shared_key = _shared_stream_key(trigger) + with patch(LOAD_TABLE, return_value=_table_at(111)): + stream = trigger.filter_shared_stream(_never()) if shared_key is not None else trigger.run() + async with aclosing(stream) as events: + async for event in events: + assert event.payload["snapshot_id"] == 111 + break + + +@pytest.mark.asyncio +@pytest.mark.parametrize("error", [NoSuchTableError, NoSuchNamespaceError]) +async def test_waits_for_a_table_that_does_not_exist_yet(error): + """Raising kills the trigger, and the triggerer then restarts it once per second.""" + trigger = IcebergTableSnapshotTrigger(table="db.tbl", poll_interval=0.01) + + with patch(LOAD_TABLE, side_effect=error("Table does not exist: db.tbl")): + assert await _collect(trigger, 1, timeout=0.15) == [] + + +@pytest.mark.asyncio +async def test_fires_once_the_table_appears(): + trigger = IcebergTableSnapshotTrigger(table="db.tbl", poll_interval=0.01) + absent = NoSuchTableError("Table does not exist: db.tbl") + + with patch(LOAD_TABLE, side_effect=[absent, absent, _table_at(111)]): + payloads = await _collect(trigger, 1) + + assert [p["snapshot_id"] for p in payloads] == [111]