diff --git a/airflow-core/src/airflow/serialization/serialized_objects.py b/airflow-core/src/airflow/serialization/serialized_objects.py index 54bc3389c64ce..f42e51e28f7d0 100644 --- a/airflow-core/src/airflow/serialization/serialized_objects.py +++ b/airflow-core/src/airflow/serialization/serialized_objects.py @@ -41,7 +41,7 @@ from dateutil import relativedelta from pendulum.tz.timezone import FixedTimezone, Timezone -from airflow._shared.module_loading import import_string, qualname +from airflow._shared.module_loading import qualname from airflow._shared.timezones.timezone import from_timestamp, parse_timezone, utcnow from airflow.callbacks.callback_requests import DagCallbackRequest, TaskCallbackRequest from airflow.exceptions import AirflowException, DeserializationError, SerializationError @@ -242,6 +242,31 @@ def _decode_priority_weight_strategy(var: str) -> PriorityWeightStrategy: return priority_weight_strategy_class() +# Builtin exceptions a BASE_EXC_SER node can rebuild into. A user defined subclass +# (e.g. a custom KeyError) serializes to a name absent here and won't round-trip -- +# unchanged from before, since builtins never held it either. +_DESERIALIZABLE_BUILTIN_EXCEPTIONS: dict[str, type[BaseException]] = { + "KeyError": KeyError, + "AttributeError": AttributeError, +} + + +def _resolve_airflow_exception(exc_cls_name: str) -> type[AirflowException]: + """ + Resolve a stored ``AirflowException`` name without importing it. + + Read via ``vars()`` rather than ``getattr`` -- a module's deprecation-shim + ``__getattr__`` can still import on access -- which also lets pre-3.2.0 blobs + naming the old ``airflow.exceptions`` path keep resolving. + """ + module_name, _, attr_name = exc_cls_name.rpartition(".") + module = sys.modules.get(module_name) + exc_cls = vars(module).get(attr_name) if module is not None else None + if not (isinstance(exc_cls, type) and issubclass(exc_cls, AirflowException)): + raise DeserializationError(f"Refusing to deserialize unknown exception class {exc_cls_name!r}") + return exc_cls + + def _encode_start_trigger_args(var: StartTriggerArgs) -> dict[str, Any]: """Encode a StartTriggerArgs.""" @@ -664,9 +689,14 @@ def deserialize(cls, encoded_var: Any) -> Any: kwargs = deser["kwargs"] del deser if type_ == DAT.AIRFLOW_EXC_SER: - exc_cls = import_string(exc_cls_name) + exc_cls: type[BaseException] = _resolve_airflow_exception(exc_cls_name) else: - exc_cls = import_string(f"builtins.{exc_cls_name}") + builtin_exc_cls = _DESERIALIZABLE_BUILTIN_EXCEPTIONS.get(exc_cls_name) + if builtin_exc_cls is None: + raise DeserializationError( + f"Refusing to deserialize unsupported builtin exception {exc_cls_name!r}" + ) + exc_cls = builtin_exc_cls return exc_cls(*args, **kwargs) elif type_ == DAT.SET: return {cls.deserialize(v) for v in var} diff --git a/airflow-core/tests/unit/serialization/test_dag_serialization.py b/airflow-core/tests/unit/serialization/test_dag_serialization.py index 4708365e846e2..7852c25dc5ee8 100644 --- a/airflow-core/tests/unit/serialization/test_dag_serialization.py +++ b/airflow-core/tests/unit/serialization/test_dag_serialization.py @@ -52,6 +52,7 @@ from airflow.dag_processing.dagbag import DagBag from airflow.exceptions import ( AirflowException, + DeserializationError, ParamValidationError, SerializationError, ) @@ -77,7 +78,7 @@ from airflow.serialization.definitions.param import SerializedParam from airflow.serialization.definitions.xcom_arg import SchedulerPlainXComArg from airflow.serialization.encoders import ensure_serialized_asset -from airflow.serialization.enums import Encoding +from airflow.serialization.enums import DagAttributeTypes, Encoding from airflow.serialization.json_schema import load_dag_schema_dict from airflow.serialization.serialized_objects import ( BaseSerialization, @@ -519,6 +520,10 @@ def timetable_plugin(monkeypatch: pytest.MonkeyPatch): ) +class PluginContributedError(AirflowException): + """Defined outside airflow.exceptions on purpose: resolution must not be limited to exceptions Airflow itself ships.""" + + class TestStringifiedDAGs: """Unit tests for stringified DAGs.""" @@ -2806,6 +2811,120 @@ def test_create_dagrun_accepts_partition_key_for_partitioned_at_runtime_dag(self dr = dag_maker.create_dagrun(partition_key="runtime-key") assert dr.partition_key == "runtime-key" + @pytest.mark.parametrize( + ("module_name", "attr_name"), + [ + pytest.param("subprocess", "check_output", id="callable_in_loaded_module"), + pytest.param("os", "system", id="another_callable"), + pytest.param("builtins", "eval", id="builtin_callable"), + pytest.param("airflow.exceptions", "NoSuchThing", id="missing_attr"), + ], + ) + def test_airflow_exc_deserialization_rejects_a_name_in_a_loaded_module(self, module_name, attr_name): + """Having the module loaded is not enough -- the name must be an AirflowException subclass. + + The module is imported by the test first, so the rejection cannot be an artefact of the + module simply being absent. + """ + importlib.import_module(module_name) + encoded = BaseSerialization._encode( + BaseSerialization.serialize( + {"exc_cls_name": f"{module_name}.{attr_name}", "args": [], "kwargs": {}} + ), + type_=DagAttributeTypes.AIRFLOW_EXC_SER, + ) + with pytest.raises(DeserializationError, match="Refusing to deserialize unknown exception class"): + BaseSerialization.deserialize(encoded) + + @pytest.mark.parametrize( + "exc_cls_name", + [ + pytest.param("not.a.loaded.module.Thing", id="unloaded_module"), + pytest.param("NoModulePart", id="no_module_part"), + pytest.param("", id="empty"), + ], + ) + def test_airflow_exc_deserialization_rejects_an_unresolvable_name(self, exc_cls_name): + encoded = BaseSerialization._encode( + BaseSerialization.serialize({"exc_cls_name": exc_cls_name, "args": [], "kwargs": {}}), + type_=DagAttributeTypes.AIRFLOW_EXC_SER, + ) + with pytest.raises(DeserializationError, match="Refusing to deserialize unknown exception class"): + BaseSerialization.deserialize(encoded) + + def test_airflow_exc_deserialization_does_not_import_the_named_module(self): + """Resolution reads ``sys.modules``; it never imports what the blob names.""" + module_name = "airflow_exc_module_that_must_not_be_imported" + encoded = BaseSerialization._encode( + BaseSerialization.serialize({"exc_cls_name": f"{module_name}.Boom", "args": [], "kwargs": {}}), + type_=DagAttributeTypes.AIRFLOW_EXC_SER, + ) + with mock.patch.object( + importlib, "import_module", side_effect=AssertionError("imported"), autospec=True + ): + with pytest.raises(DeserializationError): + BaseSerialization.deserialize(encoded) + assert module_name not in sys.modules + + def test_airflow_exc_deserialization_roundtrips_airflow_exception(self): + result = BaseSerialization.deserialize(BaseSerialization.serialize(AirflowException("boom"))) + assert isinstance(result, AirflowException) + assert result.args == ("boom",) + + def test_airflow_exc_deserialization_resolves_a_subclass_outside_airflow(self): + """Resolves as long as its module is loaded, registration order doesn't matter.""" + result = BaseSerialization.deserialize(BaseSerialization.serialize(PluginContributedError("boom"))) + assert isinstance(result, PluginContributedError) + assert result.args == ("boom",) + + @pytest.mark.parametrize( + "exc_cls_name", + [ + "airflow.exceptions.AirflowException", + "airflow.exceptions.AirflowNotFoundException", + "airflow.exceptions.ParamValidationError", + "airflow.sdk.exceptions.AirflowException", + ], + ) + def test_airflow_exc_deserialization_accepts_the_pre_3_2_module_spelling(self, exc_cls_name): + """A blob written before these exceptions moved to ``airflow.sdk.exceptions`` still reads. + + 3.0/3.1 stored ``airflow.exceptions.``; 3.2.0 moved the classes and left a re-export + behind. Both spellings have to resolve, or upgrading strands every stored blob that carries + an exception node. + """ + encoded = BaseSerialization._encode( + BaseSerialization.serialize({"exc_cls_name": exc_cls_name, "args": ["boom"], "kwargs": {}}), + type_=DagAttributeTypes.AIRFLOW_EXC_SER, + ) + result = BaseSerialization.deserialize(encoded) + assert isinstance(result, AirflowException) + + def test_base_exc_deserialization_rejects_non_allowlisted_builtin(self): + """A BASE_EXC_SER name outside the {KeyError, AttributeError} the encoder emits is rejected.""" + # ``eval`` is the weaponisable case; ``ValueError`` is a harmless builtin the encoder + # never emits as BASE_EXC_SER -- both must be rejected. + for name in ("eval", "ValueError"): + encoded = BaseSerialization._encode( + BaseSerialization.serialize({"exc_cls_name": name, "args": ["1"], "kwargs": {}}), + type_=DagAttributeTypes.BASE_EXC_SER, + ) + with pytest.raises( + DeserializationError, match="Refusing to deserialize unsupported builtin exception" + ): + BaseSerialization.deserialize(encoded) + + @pytest.mark.parametrize("exc_type", [KeyError, AttributeError]) + def test_base_exc_serialize_deserialize_round_trip(self, exc_type): + """Pins the allow-list to the encode branch, by going through ``serialize`` rather than + hand-building the node: if that branch ever accepts another builtin, this fails.""" + result = BaseSerialization.deserialize(BaseSerialization.serialize(exc_type("boom"))) + + assert isinstance(result, exc_type) + # The encode branch stores ``[var.args]``, so the args arrive nested. Pre-existing shape, + # asserted as it is rather than as it ought to be. + assert result.args == (("boom",),) + def test_kubernetes_optional(): """Test that serialization module loads without kubernetes, but deserialization of PODs requires it"""