Make dataset_events access safer

The key is not always available (mostly in tests). Create a placeholder
if it's not.

airflow/models/baseoperator.py

@@ -91,7 +91,7 @@
 from airflow.ti_deps.deps.prev_dagrun_dep import PrevDagrunDep
 from airflow.ti_deps.deps.trigger_rule_dep import TriggerRuleDep
 from airflow.utils import timezone
-from airflow.utils.context import Context
+from airflow.utils.context import Context, context_get_dataset_events
 from airflow.utils.decorators import fixup_decorator_warning_stack
 from airflow.utils.edgemodifier import EdgeModifier
 from airflow.utils.helpers import validate_key
@@ -1272,8 +1272,11 @@ def pre_execute(self, context: Any):
         """Execute right before self.execute() is called."""
         if self._pre_execute_hook is None:
             return
-        runner = ExecutionCallableRunner(self._pre_execute_hook, context["dataset_events"], logger=self.log)
-        runner.run(context)
+        ExecutionCallableRunner(
+            self._pre_execute_hook,
+            context_get_dataset_events(context),
+            logger=self.log,
+        ).run(context)
 
     def execute(self, context: Context) -> Any:
         """
@@ -1294,8 +1297,11 @@ def post_execute(self, context: Any, result: Any = None):
         """
         if self._post_execute_hook is None:
             return
-        runner = ExecutionCallableRunner(self._post_execute_hook, context["dataset_events"], logger=self.log)
-        runner.run(context, result)
+        ExecutionCallableRunner(
+            self._post_execute_hook,
+            context_get_dataset_events(context),
+            logger=self.log,
+        ).run(context, result)
 
     def on_kill(self) -> None:
         """

airflow/models/taskinstance.py

@@ -111,6 +111,7 @@
     Context,
     DatasetEventAccessors,
     VariableAccessor,
+    context_get_dataset_events,
     context_merge,
 )
 from airflow.utils.email import send_email
@@ -437,8 +438,11 @@ def _execute_callable(context: Context, **execute_callable_kwargs):
             # Print a marker for log grouping of details before task execution
             log.info("::endgroup::")
 
-            runner = ExecutionCallableRunner(execute_callable, context["dataset_events"], logger=log)
-            return runner.run(context=context, **execute_callable_kwargs)
+            return ExecutionCallableRunner(
+                execute_callable,
+                context_get_dataset_events(context),
+                logger=log,
+            ).run(context=context, **execute_callable_kwargs)
         except SystemExit as e:
             # Handle only successful cases here. Failure cases will be handled upper
             # in the exception chain.

airflow/operators/python.py

@@ -52,7 +52,7 @@
 from airflow.models.variable import Variable
 from airflow.operators.branch import BranchMixIn
 from airflow.utils import hashlib_wrapper
-from airflow.utils.context import DatasetEventAccessors, context_copy_partial, context_merge
+from airflow.utils.context import context_copy_partial, context_get_dataset_events, context_merge
 from airflow.utils.file import get_unique_dag_module_name
 from airflow.utils.operator_helpers import ExecutionCallableRunner, KeywordParameters
 from airflow.utils.process_utils import execute_in_subprocess
@@ -231,11 +231,8 @@ def __init__(
     def execute(self, context: Context) -> Any:
         context_merge(context, self.op_kwargs, templates_dict=self.templates_dict)
         self.op_kwargs = self.determine_kwargs(context)
+        self._dataset_events = context_get_dataset_events(context)
 
-        try:
-            self._dataset_events = context["dataset_events"]
-        except KeyError:
-            self._dataset_events = DatasetEventAccessors()
         return_value = self.execute_callable()
         if self.show_return_value_in_logs:
             self.log.info("Done. Returned value was: %s", return_value)

airflow/utils/context.py

@@ -355,3 +355,10 @@ def _create_value(k: str, v: Any) -> Any:
         return lazy_object_proxy.Proxy(factory)
 
     return {k: _create_value(k, v) for k, v in source._context.items()}
+
+
+def context_get_dataset_events(context: Context) -> DatasetEventAccessors:
+    try:
+        return context["dataset_events"]
+    except KeyError:
+        return DatasetEventAccessors()

airflow/utils/context.pyi

@@ -127,3 +127,4 @@ def context_merge(context: Context, **kwargs: Any) -> None: ...
 def context_update_for_unmapped(context: Context, task: BaseOperator) -> None: ...
 def context_copy_partial(source: Context, keys: Container[str]) -> Context: ...
 def lazy_mapping_from_context(source: Context) -> Mapping[str, Any]: ...
+def context_get_dataset_events(context: Context) -> DatasetEventAccessors: ...

tests/operators/test_python.py

@@ -1121,7 +1121,6 @@ def f(
             conf,
             dag,
             dag_run,
-            dataset_events,
             task,
             # other
             **context,