diff --git a/airflow/datasets/__init__.py b/airflow/datasets/__init__.py index d4305eeb04940..55d947544c1d2 100644 --- a/airflow/datasets/__init__.py +++ b/airflow/datasets/__init__.py @@ -206,12 +206,20 @@ def iter_dag_dependencies(self, *, source: str, target: str) -> Iterator[DagDepe raise NotImplementedError -@attr.define(unsafe_hash=False) +@attr.define() class DatasetAlias(BaseDataset): """A represeation of dataset alias which is used to create dataset during the runtime.""" name: str + def __eq__(self, other: Any) -> bool: + if isinstance(other, DatasetAlias): + return self.name == other.name + return NotImplemented + + def __hash__(self) -> int: + return hash(self.name) + def iter_dag_dependencies(self, *, source: str, target: str) -> Iterator[DagDependency]: """ Iterate a dataset alias as dag dependency. @@ -233,7 +241,7 @@ class DatasetAliasEvent(TypedDict): dest_dataset_uri: str -@attr.define(unsafe_hash=False) +@attr.define() class Dataset(os.PathLike, BaseDataset): """A representation of data dependencies between workflows.""" @@ -241,13 +249,21 @@ class Dataset(os.PathLike, BaseDataset): converter=_sanitize_uri, validator=[attr.validators.min_len(1), attr.validators.max_len(3000)], ) - extra: dict[str, Any] = attr.field(factory=dict) + extra: dict[str, Any] | None = None __version__: ClassVar[int] = 1 def __fspath__(self) -> str: return self.uri + def __eq__(self, other: Any) -> bool: + if isinstance(other, self.__class__): + return self.uri == other.uri + return NotImplemented + + def __hash__(self) -> int: + return hash(self.uri) + @property def normalized_uri(self) -> str | None: """ diff --git a/airflow/lineage/hook.py b/airflow/lineage/hook.py index 45227c5248180..4ff35e4d9ce82 100644 --- a/airflow/lineage/hook.py +++ b/airflow/lineage/hook.py @@ -112,7 +112,7 @@ def create_dataset( """ if uri: # Fallback to default factory using the provided URI - return Dataset(uri=uri, extra=dataset_extra or {}) + return Dataset(uri=uri, extra=dataset_extra) if not scheme: self.log.debug( diff --git a/airflow/models/dag.py b/airflow/models/dag.py index 56f7dc89d25b1..54ebce7392b3b 100644 --- a/airflow/models/dag.py +++ b/airflow/models/dag.py @@ -2786,12 +2786,12 @@ def bulk_write_to_db( curr_outlet_references = curr_orm_dag and curr_orm_dag.task_outlet_dataset_references for task in dag.tasks: dataset_outlets: list[Dataset] = [] - dataset_alias_outlets: list[DatasetAlias] = [] + dataset_alias_outlets: set[DatasetAlias] = set() for outlet in task.outlets: if isinstance(outlet, Dataset): dataset_outlets.append(outlet) elif isinstance(outlet, DatasetAlias): - dataset_alias_outlets.append(outlet) + dataset_alias_outlets.add(outlet) if not dataset_outlets: if curr_outlet_references: diff --git a/newsfragments/42054.significant.rst b/newsfragments/42054.significant.rst deleted file mode 100644 index aebf70757fa07..0000000000000 --- a/newsfragments/42054.significant.rst +++ /dev/null @@ -1,4 +0,0 @@ -Dataset and DatasetAlias are no longer hashable - -This means they can no longer be used as dict keys or put into a set. Dataset's -equality logic is also tweaked slightly to consider the extra dict. diff --git a/tests/datasets/test_dataset.py b/tests/datasets/test_dataset.py index 940e445669cce..017c874763495 100644 --- a/tests/datasets/test_dataset.py +++ b/tests/datasets/test_dataset.py @@ -120,6 +120,12 @@ def test_not_equal_when_different_uri(): assert dataset1 != dataset2 +def test_hash(): + uri = "s3://example/dataset" + dataset = Dataset(uri=uri) + hash(dataset) + + def test_dataset_logic_operations(): result_or = dataset1 | dataset2 assert isinstance(result_or, DatasetAny) @@ -181,10 +187,10 @@ def test_datasetbooleancondition_evaluate_iter(): assert all_condition.evaluate({"s3://bucket1/data1": True, "s3://bucket2/data2": False}) is False # Testing iter_datasets indirectly through the subclasses - datasets_any = dict(any_condition.iter_datasets()) - datasets_all = dict(all_condition.iter_datasets()) - assert datasets_any == {"s3://bucket1/data1": dataset1, "s3://bucket2/data2": dataset2} - assert datasets_all == {"s3://bucket1/data1": dataset1, "s3://bucket2/data2": dataset2} + datasets_any = set(any_condition.iter_datasets()) + datasets_all = set(all_condition.iter_datasets()) + assert datasets_any == {("s3://bucket1/data1", dataset1), ("s3://bucket2/data2", dataset2)} + assert datasets_all == {("s3://bucket1/data1", dataset1), ("s3://bucket2/data2", dataset2)} @pytest.mark.parametrize( diff --git a/tests/lineage/test_hook.py b/tests/lineage/test_hook.py index 16f386b6849ea..97e160e1f971b 100644 --- a/tests/lineage/test_hook.py +++ b/tests/lineage/test_hook.py @@ -69,7 +69,7 @@ def test_add_input_dataset(self, mock_dataset): self.collector.add_input_dataset(hook, uri="test_uri") assert next(iter(self.collector._inputs.values())) == (dataset, hook) - mock_dataset.assert_called_once_with(uri="test_uri", extra={}) + mock_dataset.assert_called_once_with(uri="test_uri", extra=None) def test_grouping_datasets(self): hook_1 = MagicMock() @@ -96,7 +96,7 @@ def test_grouping_datasets(self): @patch("airflow.lineage.hook.ProvidersManager") def test_create_dataset(self, mock_providers_manager): def create_dataset(arg1, arg2="default", extra=None): - return Dataset(uri=f"myscheme://{arg1}/{arg2}", extra=extra or {}) + return Dataset(uri=f"myscheme://{arg1}/{arg2}", extra=extra) mock_providers_manager.return_value.dataset_factories = {"myscheme": create_dataset} assert self.collector.create_dataset( diff --git a/tests/timetables/test_datasets_timetable.py b/tests/timetables/test_datasets_timetable.py index b456b9bf5dc9c..b055f0d34dc90 100644 --- a/tests/timetables/test_datasets_timetable.py +++ b/tests/timetables/test_datasets_timetable.py @@ -134,7 +134,7 @@ def test_serialization(dataset_timetable: DatasetOrTimeSchedule, monkeypatch: An "timetable": "mock_serialized_timetable", "dataset_condition": { "__type": "dataset_all", - "objects": [{"__type": "dataset", "uri": "test_dataset", "extra": {}}], + "objects": [{"__type": "dataset", "uri": "test_dataset", "extra": None}], }, }