Create custom DAG classes with custom `clear` etc.
- Dominant language
- Python
- Stars
- 46.9k
- Forks
- 17.8k
- Avg merge
- 2d 10h
- Merged PRs (30d)
- 483
Description
### Description
I have a use case where we have a task which creates a Spark cluster for all other tasks to use. This task is the root of all tasks in the DAG so runs first when the DAG is kicked off. If anything fails in the DAG, I need to clear the failed task as well as the root task to create a cluster which can be reused. This results in me having to clear two tasks each time.
I would love to be able to override the `clear` method in DAG https://github.com/apache/airflow/blob/main/airflow/models/dag.py#L1821 in a custom DAG class to also always clear the root task. I tried doing this in my custom task but it didnt work. After some looking around, the reason is because of the way the DAGs are serialized to be stored in the metadata DB. They lose information about their original type so we are unable to overwrite any methods.
This happens in this method: https://github.com/apache/airflow/blob/main/airflow/serialization/serialized_objects.py#L1000
Anyone have an easier way to accomplishing this?
One way to solve this would be to store the information about the class in this object - this code is adding `serialize_dag["_class"] = dag.__class__.__module__ + '.' + dag.__class__.__name__` in the `serialize_dag` and updating https://github.com/apache/airflow/blob/main/airflow/serialization/schema.json#L109 to include `_class`. Then the issue comes down to deserializing the DAG which Im struggling with.
We can run
```
from pydoc import locate
dag_class = locate(encoded_dag["_class"])
dag_class(dag.__dict__)
```
here https://github.com/apache/airflow/blob/main/airflow/serialization/serialized_objects.py#L1120 but Im getting tripped up on how to define the attributes we pass to the `dag_class` constructor. Any help is appreciated 🙏
### Use case/motivation
I would like to be able to overwrite methods within the DAG class to customize the behaviour of my custom DAG classes which inherit from DAG
### Related issues
I didnt find any :(
### Are you willing to submit a PR?
- [X] Yes I am willing to submit a PR!
### Code of Conduct
- [X] I agree to follow this project's [Code of Conduct](https://github.com/apache/airflow/blob/main/CODE_OF_CONDUCT.md)
Contributor guide
Research direction
Start by reading the DAG clear method in airflow/models/dag.py and the serialization and deserialization paths in airflow/serialization/serialized_objects.py, including serialize_dag and the code around line 1120. Check airflow/serialization/schema.json for the serialized shape. Done means a custom DAG class retains its behavior after serialization and deserialization, with tests covering the supported attributes and custom clear behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- data-engineering
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100