pierrejeambrun commented on code in PR #49164:
URL: https://github.com/apache/airflow/pull/49164#discussion_r2269828119
##########
airflow-core/tests/unit/api_fastapi/common/test_exceptions.py:
##########
@@ -282,3 +294,102 @@ def test_handle_multiple_columns_unique_constraint_error(
assert response_detail == expected_detail
assert "INSERT INTO dag_run" in actual_statement
+ assert exeinfo_response_error.value.detail == expected_exception.detail
+
+
+class TestDagErrorHandler:
+ dag_error_handler = DagErrorHandler()
+
+ def test_handle_deserialization_error_with_value_error(self):
+ error_message = "Missing Dag ID in serialized Dag"
+ deserialization_error = DeserializationError("test_dag_id")
+
+ value_error = ValueError(error_message)
+ deserialization_error.__cause__ = value_error
+
+ expected_exception = HTTPException(
+ status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+ detail=f"An error occurred while trying to deserialize Dag:
{deserialization_error}",
+ )
+
+ with pytest.raises(HTTPException) as exeinfo_response_error:
+ self.dag_error_handler.exception_handler(None,
deserialization_error)
+
+ assert exeinfo_response_error.value.status_code ==
expected_exception.status_code
+ assert exeinfo_response_error.value.detail == expected_exception.detail
+
+ def test_handle_deserialization_error_with_runtime_error(self):
+ error_message = "Error during Dag serialization process"
+ deserialization_error = DeserializationError("test_dag_id")
+
+ runtime_error = RuntimeError(error_message)
+ deserialization_error.__cause__ = runtime_error
+
+ expected_exception = HTTPException(
+ status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+ detail=f"An error occurred while trying to deserialize Dag:
{deserialization_error}",
+ )
+
+ with pytest.raises(HTTPException) as exeinfo_response_error:
+ self.dag_error_handler.exception_handler(None,
deserialization_error)
+
+ assert exeinfo_response_error.value.status_code ==
expected_exception.status_code
+ assert exeinfo_response_error.value.detail == expected_exception.detail
+
+ def test_handle_deserialization_error_with_key_error(self):
+ key = "required_field"
+ deserialization_error = DeserializationError("test_dag_id")
+
+ key_error = KeyError(key)
+ deserialization_error.__cause__ = key_error
+
+ expected_exception = HTTPException(
+ status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+ detail=f"An error occurred while trying to deserialize Dag:
{deserialization_error}",
+ )
+
+ with pytest.raises(HTTPException) as exeinfo_response_error:
+ self.dag_error_handler.exception_handler(None,
deserialization_error)
+
+ assert exeinfo_response_error.value.status_code ==
expected_exception.status_code
+ assert exeinfo_response_error.value.detail == expected_exception.detail
+
+ def test_handle_real_dag_deserialization_error(self):
+ """Test handling a real Dag deserialization error with actual
serialized Dag."""
+
+ # Create a test DAG
+ dag_id = "test_dag_for_error"
+ test_dag = DAG(dag_id=dag_id, start_date=datetime(2025, 4, 15),
schedule="@once")
+ EmptyOperator(task_id="test_task", dag=test_dag)
+ test_dag.sync_to_db()
+ SerializedDagModel.write_dag(test_dag, bundle_name=dag_id)
+
+ with create_session() as session:
+ dag_model =
session.scalar(select(SerializedDagModel).where(SerializedDagModel.dag_id ==
dag_id))
+ if not dag_model:
+ pytest.fail("Failed to find serialized Dag in database")
Review Comment:
Test will fail later anyway
```suggestion
```
##########
airflow-core/tests/unit/api_fastapi/common/test_exceptions.py:
##########
@@ -282,3 +294,102 @@ def test_handle_multiple_columns_unique_constraint_error(
assert response_detail == expected_detail
assert "INSERT INTO dag_run" in actual_statement
+ assert exeinfo_response_error.value.detail == expected_exception.detail
+
+
+class TestDagErrorHandler:
+ dag_error_handler = DagErrorHandler()
+
+ def test_handle_deserialization_error_with_value_error(self):
+ error_message = "Missing Dag ID in serialized Dag"
+ deserialization_error = DeserializationError("test_dag_id")
+
+ value_error = ValueError(error_message)
+ deserialization_error.__cause__ = value_error
+
+ expected_exception = HTTPException(
+ status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+ detail=f"An error occurred while trying to deserialize Dag:
{deserialization_error}",
+ )
+
+ with pytest.raises(HTTPException) as exeinfo_response_error:
Review Comment:
Use `matches` param for `pytest.raises` instead to match the exception
message.
Same for all tests actually, this will make the code shorted and remove the
need for a manual extra assertion on the exception message.
##########
airflow-core/src/airflow/exceptions.py:
##########
@@ -517,6 +517,23 @@ class UnknownExecutorException(ValueError):
"""Raised when an attempt is made to load an executor which is not
configured."""
+class DeserializationError(Exception):
+ """
+ Raised when a Dag cannot be deserialized.
+
+ This exception should be raised using exception chaining:
+ `raise DeserializationError(dag_id) from original_exception`
+ """
+
+ def __init__(self, dag_id: str = "unknown"):
Review Comment:
```suggestion
def __init__(self, dag_id: str | None = None):
```
##########
airflow-core/src/airflow/serialization/serialized_objects.py:
##########
@@ -1935,81 +1935,154 @@ def serialize_dag(cls, dag: DAG) -> dict:
@classmethod
def deserialize_dag(cls, encoded_dag: dict[str, Any]) -> SerializedDAG:
"""Deserializes a DAG from a JSON object."""
- if "dag_id" not in encoded_dag:
- raise RuntimeError(
- "Encoded dag object has no dag_id key. You may need to run
`airflow dags reserialize`."
- )
-
- dag = SerializedDAG(dag_id=encoded_dag["dag_id"], schedule=None)
-
- for k, v in encoded_dag.items():
- if k == "_downstream_task_ids":
- v = set(v)
- elif k == "tasks":
- SerializedBaseOperator._load_operator_extra_links =
cls._load_operator_extra_links
- tasks = {}
- for obj in v:
- if obj.get(Encoding.TYPE) == DAT.OP:
- deser =
SerializedBaseOperator.deserialize_operator(obj[Encoding.VAR])
- tasks[deser.task_id] = deser
- k = "task_dict"
- v = tasks
- elif k == "timezone":
- v = cls._deserialize_timezone(v)
- elif k == "dagrun_timeout":
- v = cls._deserialize_timedelta(v)
- elif k.endswith("_date"):
- v = cls._deserialize_datetime(v)
- elif k == "edge_info":
- # Value structure matches exactly
- pass
- elif k == "timetable":
- v = decode_timetable(v)
- elif k == "weight_rule":
- v = decode_priority_weight_strategy(v)
- elif k in cls._decorated_fields:
- v = cls.deserialize(v)
- elif k == "params":
- v = cls._deserialize_params_dict(v)
- elif k == "tags":
- v = set(v)
- # else use v as it is
-
- object.__setattr__(dag, k, v)
+ try:
+ if "dag_id" not in encoded_dag:
+ raise RuntimeError(
+ "Encoded dag object has no dag_id key. You may need to
run `airflow dags reserialize`."
+ )
- # Set _task_group
- if "task_group" in encoded_dag:
- tg = TaskGroupSerialization.deserialize_task_group(
- encoded_dag["task_group"],
- None,
- dag.task_dict,
- dag,
+ dag = SerializedDAG(dag_id=encoded_dag["dag_id"], schedule=None)
+
+ for k, v in encoded_dag.items():
+ if k == "_downstream_task_ids":
+ v = set(v)
+ elif k == "tasks":
+ SerializedBaseOperator._load_operator_extra_links =
cls._load_operator_extra_links
+ tasks = {}
+ for obj in v:
+ if obj.get(Encoding.TYPE) == DAT.OP:
+ deser =
SerializedBaseOperator.deserialize_operator(obj[Encoding.VAR])
+ tasks[deser.task_id] = deser
+ k = "task_dict"
+ v = tasks
+ elif k == "timezone":
+ v = cls._deserialize_timezone(v)
+ elif k == "dagrun_timeout":
+ v = cls._deserialize_timedelta(v)
+ elif k.endswith("_date"):
+ v = cls._deserialize_datetime(v)
+ elif k == "edge_info":
+ # Value structure matches exactly
+ pass
+ elif k == "timetable":
+ v = decode_timetable(v)
+ elif k == "weight_rule":
+ v = decode_priority_weight_strategy(v)
+ elif k in cls._decorated_fields:
+ v = cls.deserialize(v)
+ elif k == "params":
+ v = cls._deserialize_params_dict(v)
+ elif k == "tags":
+ v = set(v)
+ # else use v as it is
+
+ object.__setattr__(dag, k, v)
+
+ # Set _task_group
+ if "task_group" in encoded_dag:
+ tg = TaskGroupSerialization.deserialize_task_group(
+ encoded_dag["task_group"],
+ None,
+ dag.task_dict,
+ dag,
+ )
+ object.__setattr__(dag, "task_group", tg)
+ else:
+ # This must be old data that had no task_group. Create a root
TaskGroup and add
+ # all tasks to it.
+ object.__setattr__(dag, "task_group",
TaskGroup.create_root(dag))
+ for task in dag.tasks:
+ dag.task_group.add(task)
+
+ # Set has_on_*_callbacks to True if they exist in Serialized blob
as False is the default
+ if "has_on_success_callback" in encoded_dag:
+ dag.has_on_success_callback = True
+ if "has_on_failure_callback" in encoded_dag:
+ dag.has_on_failure_callback = True
+
+ keys_to_set_none = (
+ dag.get_serialized_fields() - encoded_dag.keys() -
cls._CONSTRUCTOR_PARAMS.keys()
)
- object.__setattr__(dag, "task_group", tg)
- else:
- # This must be old data that had no task_group. Create a root
TaskGroup and add
- # all tasks to it.
- object.__setattr__(dag, "task_group", TaskGroup.create_root(dag))
- for task in dag.tasks:
- dag.task_group.add(task)
-
- # Set has_on_*_callbacks to True if they exist in Serialized blob as
False is the default
- if "has_on_success_callback" in encoded_dag:
- dag.has_on_success_callback = True
- if "has_on_failure_callback" in encoded_dag:
- dag.has_on_failure_callback = True
-
- if "deadline" in encoded_dag and encoded_dag["deadline"] is not None:
- dag.deadline =
DeadlineAlert.deserialize_deadline_alert(encoded_dag["deadline"])
+ for k in keys_to_set_none:
+ setattr(dag, k, None)
+
+ for task in dag.task_dict.values():
+ SerializedBaseOperator.set_task_dag_references(task, dag)
+
+ for k, v in encoded_dag.items():
+ if k == "_downstream_task_ids":
+ v = set(v)
+ elif k == "tasks":
+ SerializedBaseOperator._load_operator_extra_links =
cls._load_operator_extra_links
+ tasks = {}
+ for obj in v:
+ if obj.get(Encoding.TYPE) == DAT.OP:
+ deser =
SerializedBaseOperator.deserialize_operator(obj[Encoding.VAR])
+ tasks[deser.task_id] = deser
+ k = "task_dict"
+ v = tasks
+ elif k == "timezone":
+ v = cls._deserialize_timezone(v)
+ elif k == "dagrun_timeout":
+ v = cls._deserialize_timedelta(v)
+ elif k.endswith("_date"):
+ v = cls._deserialize_datetime(v)
+ elif k == "edge_info":
+ # Value structure matches exactly
+ pass
+ elif k == "timetable":
+ v = decode_timetable(v)
+ elif k == "weight_rule":
+ v = decode_priority_weight_strategy(v)
+ elif k in cls._decorated_fields:
+ v = cls.deserialize(v)
+ elif k == "params":
+ v = cls._deserialize_params_dict(v)
+ elif k == "tags":
+ v = set(v)
+ # else use v as it is
+
+ object.__setattr__(dag, k, v)
+
+ # Set _task_group
+ if "task_group" in encoded_dag:
+ tg = TaskGroupSerialization.deserialize_task_group(
+ encoded_dag["task_group"],
+ None,
+ dag.task_dict,
+ dag,
+ )
+ object.__setattr__(dag, "task_group", tg)
+ else:
+ # This must be old data that had no task_group. Create a root
TaskGroup and add
+ # all tasks to it.
+ object.__setattr__(dag, "task_group",
TaskGroup.create_root(dag))
+ for task in dag.tasks:
+ dag.task_group.add(task)
+
+ # Set has_on_*_callbacks to True if they exist in Serialized blob
as False is the default
+ if "has_on_success_callback" in encoded_dag:
+ dag.has_on_success_callback = True
+ if "has_on_failure_callback" in encoded_dag:
+ dag.has_on_failure_callback = True
+
+ if "deadline" in encoded_dag and encoded_dag["deadline"] is not
None:
+ dag.deadline =
DeadlineAlert.deserialize_deadline_alert(encoded_dag["deadline"])
+
+ keys_to_set_none = (
+ dag.get_serialized_fields() - encoded_dag.keys() -
cls._CONSTRUCTOR_PARAMS.keys()
+ )
+ for k in keys_to_set_none:
+ setattr(dag, k, None)
- keys_to_set_none = dag.get_serialized_fields() - encoded_dag.keys() -
cls._CONSTRUCTOR_PARAMS.keys()
- for k in keys_to_set_none:
- setattr(dag, k, None)
+ for task in dag.task_dict.values():
+ SerializedBaseOperator.set_task_dag_references(task, dag)
- for task in dag.task_dict.values():
- SerializedBaseOperator.set_task_dag_references(task, dag)
+ return dag
- return dag
+ except (RuntimeError, ValueError, KeyError) as err:
+ dag_id = encoded_dag.get("dag_id", "unknown")
Review Comment:
Make dag_id optional, default to 'None' and use that to deduce that the
`dag_id` is unknown, this will avoid using a hardcoded "unknown" string
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
