Airflow - Change single task color

user_4800907 0
Viewed 2197

Do someone know if it is possible to change a single task color in the Airflow UI? I've read that it's possible to change an operator color (after importing it you can override the class method), but what if I would have two instances of a PythonOperator with different colors?

Thanks in advance!

airflow google-cloud-platform python
1 Answers

Currently we can set the operator color only (docs). All the task instances of an operator will have the same color. Since class variables are used (e.g. PythonOperator.ui_color, PythonOperatpr.ui_fgcolor) to set color its not possible to change color based on task instance.

One way is to inherit the PythonOperator for each color as follows:

class BluePythonOperator(PythonOperator):
    ui_color = "blue"
    ui_fgcolor = "white"


class BlackPythonOperator(PythonOperator):
    ui_color = "black"
    ui_fgcolor = "white"


with DAG(dag_id='my_example',
         default_args=default_args) as dag0:
    t0 = BluePythonOperator(
        task_id="blue_colored_task",
        python_callable=(lambda x: print(x)))
    t1 = BlackPythonOperator(
        task_id="black_colored_task",
        python_callable=(lambda x: print(x)))

However, if there are more colors we will have to create many inherited classes.

So I tried to create a class factory sort of thing that returns a subclass of PythonOperator with colors applied based on task_id.

def colorful_python_operator(task_id):
    tid_to_colors = {
        "blue_colored_task": {"ui_color": "blue", "ui_fgcolor": "white"},
        "black_colored_task": {"ui_color": "black", "ui_fgcolor": "white"}
    }

    class ColorfulPythonOperator(PythonOperator):
        value = tid_to_colors[task_id]
        ui_color = value["ui_color"]
        ui_fgcolor = value["ui_fgcolor"]
        pass

    return ColorfulPythonOperator


with DAG(dag_id='my_another_example',
         default_args=default_args) as dag1:
    t2 = colorful_python_operator("blue_colored_task")(
        task_id="blue_colored_task",
        python_callable=(lambda x: print(x)))

    t3 = colorful_python_operator("black_colored_task")(
        task_id="black_colored_task",
        python_callable=(lambda x: print(x)))
user_45936540
Related