Revert "Try to make dataset objects totally unhashable (apache#42054)" (

apache#42065)

This reverts commit 1c4a00b.

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,21 +241,29 @@ 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."""
 
     uri: str = attr.field(
         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:
         """

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(

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:

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(

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(

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}],
         },
     }