Skip to content

Commit

Permalink
fix(providers/common/compat): add back add_input_dataset and add_outp…
Browse files Browse the repository at this point in the history
…ut_dataset to NoOpCollector (apache#44681)

* fix(providers/common/compat): add back add_input_dataset and add_output_dataset to NoOpCollector

* test(providers/common/compat): add test case test_compat_has_only_asset_methods and test_compat_has_asset_and_dataset_methods
  • Loading branch information
Lee-W authored Dec 5, 2024
1 parent 41a9f97 commit b9ab0b5
Show file tree
Hide file tree
Showing 2 changed files with 38 additions and 0 deletions.
10 changes: 10 additions & 0 deletions providers/src/airflow/providers/common/compat/lineage/hook.py
Original file line number Diff line number Diff line change
Expand Up @@ -89,17 +89,27 @@ def get_hook_lineage_collector():
if AIRFLOW_V_2_10_PLUS:
return _get_asset_compat_hook_lineage_collector()

# For the case that airflow has not yet upgraded to 2.10 or higher,
# but using the providers that already uses `get_hook_lineage_collector`
class NoOpCollector:
"""
NoOpCollector is a hook lineage collector that does nothing.
It is used when you want to disable lineage collection.
"""

# for providers that support asset rename
def add_input_asset(self, *_, **__):
pass

def add_output_asset(self, *_, **__):
pass

# for providers that do not support asset rename
def add_input_dataset(self, *_, **__):
pass

def add_output_dataset(self, *_, **__):
pass

return NoOpCollector()
28 changes: 28 additions & 0 deletions providers/tests/common/compat/lineage/test_hook.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,9 +16,37 @@
# under the License.
from __future__ import annotations

import pytest

from airflow.providers.common.compat.lineage.hook import get_hook_lineage_collector

from tests_common.test_utils.compat import AIRFLOW_V_3_0_PLUS


def test_that_compat_does_not_raise():
# On compat tests this goes into ImportError code path
assert get_hook_lineage_collector() is not None
assert get_hook_lineage_collector() is not None


@pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Test requires Airflow 3.0+")
def test_compat_has_only_asset_methods():
hook_lienage_collector = get_hook_lineage_collector()

assert hook_lienage_collector.add_input_asset is not None
assert hook_lienage_collector.add_output_asset is not None

with pytest.raises(AttributeError):
hook_lienage_collector.add_input_dataset
with pytest.raises(AttributeError):
hook_lienage_collector.add_output_dataset


@pytest.mark.skipif(AIRFLOW_V_3_0_PLUS, reason="Test requires Airflow < 3.0")
def test_compat_has_asset_and_dataset_methods():
hook_lienage_collector = get_hook_lineage_collector()

assert hook_lienage_collector.add_input_asset is not None
assert hook_lienage_collector.add_output_asset is not None
assert hook_lienage_collector.add_input_dataset is not None
assert hook_lienage_collector.add_output_dataset is not None

0 comments on commit b9ab0b5

Please sign in to comment.