Check dataset_alias exists in inlets when use it to retrieve inlet_ev…

…ents (apache#41043)
eladkal · Jul 26, 2024 · f0ef691 · f0ef691
1 parent 277e746
commit f0ef691
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/airflow/utils/context.py b/airflow/utils/context.py
@@ -280,8 +280,9 @@ def __getitem__(self, key: int | str | Dataset | DatasetAlias) -> LazyDatasetEve
             obj = key
 
         if isinstance(obj, DatasetAlias):
+            dataset_alias = self._dataset_aliases[obj.name]
             join_clause = DatasetEvent.source_aliases
-            where_clause = DatasetAliasModel.name == obj.name
+            where_clause = DatasetAliasModel.name == dataset_alias.name
         elif isinstance(obj, (Dataset, str)):
             dataset = self._datasets[extract_event_key(obj)]
             join_clause = DatasetEvent.dataset

diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py
@@ -2732,6 +2732,8 @@ def read(*, inlet_events):
 
                 with pytest.raises(KeyError):
                     inlet_events["does_not_exist"]
+                with pytest.raises(KeyError):
+                    inlet_events[DatasetAlias("does_not_exist")]
                 with pytest.raises(IndexError):
                     inlet_events[DatasetAlias(dsa_name)][5]