I want to build something where I need to capture all of the leaf tasks and add a downstream dependency to them to make a job complete in our database. Is there an easy way to find all the leaf nodes of a DAG in Airflow?
Use upstream_task_ids and downstream_task_ids #property from BaseOperator
def get_start_tasks(dag: DAG) -> List[BaseOperator]:
# returns list of "head" / "root" tasks of DAG
return [task for task in dag.tasks if not task.upstream_task_ids]
def get_end_tasks(dag: DAG) -> List[BaseOperator]:
# returns list of "leaf" tasks of DAG
return [task for task in dag.tasks if not task.downstream_task_ids]
Type-Annotations from Python 3.6+
UPDATE-1
Now Airflow DAG model has powerful #property functions like
leaves
roots
topological_sort
Related
Now, I create multiple tasks using a variable like this and it works fine.
with DAG(....) as dag:
body = Variable.get("config_table", deserialize_json=True)
for i in range(len(body.keys())):
simple_task = Operator(
task_id = 'task_' + str(i),
.....
But I need to use XCOM value for some reason instead of using a variable.
Is it possible to dynamically create tasks with XCOM pull value?
I try to set value like this and it's not working
body = "{{ ti.xcom_pull(key='config_table', task_ids='get_config_table') }}"
It's possible to dynamically create tasks from XComs generated from a previous task, there are more extensive discussions on this topic, for example in this question. One of the suggested approaches follows this structure, here is a working example I made:
sample_file.json:
{
"cities": [ "London", "Paris", "BA", "NY" ]
}
Get your data from an API or file or any source. Push it as XCom.
def _process_obtained_data(ti):
list_of_cities = ti.xcom_pull(task_ids='get_data')
Variable.set(key='list_of_cities',
value=list_of_cities['cities'], serialize_json=True)
def _read_file():
with open('dags/sample_file.json') as f:
data = json.load(f)
# push to XCom using return
return data
with DAG('dynamic_tasks_example', schedule_interval='#once',
start_date=days_ago(2),
catchup=False) as dag:
get_data = PythonOperator(
task_id='get_data',
python_callable=_read_file)
Add a second task which will pull from pull from XCom and set a Variable with the data you will use to iterate later on.
preparation_task = PythonOperator(
task_id='preparation_task',
python_callable=_process_obtained_data)
*Of course, if you want you can merge both tasks into one. I prefer not to because usually, I take a subset of the fetched data to create the Variable.
Read from that Variable and later iterate on it. It's critical to define default_var.
end = DummyOperator(
task_id='end',
trigger_rule='none_failed')
# Top-level code within DAG block
iterable_list = Variable.get('list_of_cities',
default_var=['default_city'],
deserialize_json=True)
Declare dynamic tasks and their dependencies within a loop. Make the task_id uniques. TaskGroup is optional, helps you sorting the UI.
with TaskGroup('dynamic_tasks_group',
prefix_group_id=False,
) as dynamic_tasks_group:
if iterable_list:
for index, city in enumerate(iterable_list):
say_hello = PythonOperator(
task_id=f'say_hello_from_{city}',
python_callable=_print_greeting,
op_kwargs={'city_name': city, 'greeting': 'Hello'}
)
say_goodbye = PythonOperator(
task_id=f'say_goodbye_from_{city}',
python_callable=_print_greeting,
op_kwargs={'city_name': city, 'greeting': 'Goodbye'}
)
# TaskGroup level dependencies
say_hello >> say_goodbye
# DAG level dependencies
get_data >> preparation_task >> dynamic_tasks_group >> end
DAG Graph View:
Imports:
import json
from airflow import DAG
from airflow.utils.dates import days_ago
from airflow.models import Variable
from airflow.operators.python_operator import PythonOperator
from airflow.operators.dummy import DummyOperator
from airflow.utils.task_group import TaskGroup
Things to keep in mind:
If you have simultaneous dag_runs of this same DAG, all of them will use the same variable, so you may need to make it 'unique' by differentiating their names.
You must set the default value while reading the Variable, otherwise, the first execution may not be processable to the Scheduler.
The Airflow Graph View UI may not refresh the changes immediately. Happens especially in the first run after adding or removing items from the iterable on which the dynamic task generation is created.
If you need to read from many variables, it's important to remember that it's recommended to store them in one single JSON value to avoid constantly create connections to the metadata database (example in this article).
Good luck!
Edit:
Another important point to take into consideration:
With this approach, the call to Variable.get() method is top-level code, so is read by the scheduler every 30 seconds (default of min_file_process_interval setting). This means that a connection to the metadata DB will happen each time.
Edit:
Added if clause to handle emtpy iterable_list case.
This is not possible, and in general dynamic tasks are not recommended:
The way the Airflow scheduler works is by reading the dag file, loading the tasks into the memory and then checks which dags and which tasks it need to schedule, while xcom are a runtime values that are related to a specific dag run, so the scheduler cannot relay on xcom values.
When using dynamic tasks you're making debug much harder for yourself, as the values you use for creating the dag can change and you'll lose access to logs without even understanding why.
What you can do is use branch operator, to have those tasks always and just skip them based on the xcom value.
For example:
def branch_func(**context)
return f"task_{context['ti'].xcom_pull(key=key)}"
branch = BranchPythonOperator(
task_id="branch",
python_callback=branch_func
)
tasks = [BaseOperator(task_id=f"task_{i}") for i in range(3)]
branch >> tasks
In some cases it's also not good to use this method (for example when I've 100 possible tasks), in those cases I'd recommend writing your own operator or use a single PythonOperator.
We've made extensive use of [ExternalTaskSensor][1] to the point where the quantity of cross-dag dependencies have become difficult to track. As such we would like a method of extracting all tasks that use this sensor as well as the parameters passed to these tasks such as external_dag_id and external_task_id. Extracting this info would allow us to create a list of dependencies (and maybe a graph if we want it).
Approach:
So far we've been able to use the list_dags cli option to get a list of all dags. For each dag we then run the list_tasks option with the -t parameter to get a list of tasks and the operator used. The next step is to retrieve the parameters passed to these tasks, this is where we are stuck. Are there any official or non-official methods of scraping this data?
Info:
We are running Airflow 1.10.9 and Composer 1.11.0. Our script so far is written in python3.
[1]: https://airflow.readthedocs.io/en/stable/_modules/airflow/sensors/external_task_sensor.html
You can do it this way:
dag_models = session.query(DagModel).filter(DagModel.is_active.is_(True)).all()
for dag_model in dag_models:
dag = dag_model.get_dag()
for task in dag.task_dict.values():
if isinstance(task, ExternalTaskSensor):
do_smth(task.external_dag_id, task.external_task_id)
You can exploit Airflow's metadb for this.
either query directly
SELECT operator
FROM task_instance
WHERE dag_id = 'my_dag'
AND task_id = 'my_task';```
or use SQLAlchemy
from airflow.utils.session import provide_session
from airflow.models import TaskInstance
#provide_session
def get_operator_name(my_dag_id: str, my_task_id: str, session=None) -> str:
"""Fetch TaskInstance from the database using pickling"""
task_instance: TaskInstance = session.query(TaskInstance).filter(TaskInstance.dag_id == my_dag_id).filter(TaskInstance.task_id == my_task_id).first()
return task_instance.operator
The downside of this approach is that it wouldn't work until task has run at least once (and it's entry has been created in TaskInstance table)
Reference
cli.py
I'm running a script that checks the status of my database before a DAG runs and compares it to after the DAG finished running.
def pre_dag_db
pass
def run_dag
pass
def post_dag_db
pass
Is there a way for me to know when the DAG finished running so that my script knows when to run post_dag_db? The idea is that my post_dag_db runs after my DAG finished running because the DAG manipulates the db.
The easiest way to do this would be to just run the script as last task in your dag, maybe using a BashOperator.
Other options would be to trigger a separate dag (TriggerDagRunOperator) and there implement a dag that calls your script.
If you really cannot call your script from Airflow itself, you might want to check the REST APIs https://airflow.apache.org/docs/stable/api.html and use them to retrieve information about the dag_run. But this seems overly-complicated to me.
Quick and easy way is to add one task in DAG which will work/run as last task of the DAG, this will work like magic for you.
you can use any of the operator like (PythonOperator, BashOperator, etc).
I think you can use the following code:
dag = get_dag(args)
dr = DagRun.find(dag.dag_id, execution_date=args.execution_date)
print(dr[0].state if len(dr) > 0 else None)
This code is taken from airflow cli.
Make a custom class that inherits from dag, and whose dependencies are the same as your dag.
something like (custom_dag.py)
from airflow.models.dag import DAG
class PreAndPostDAG(DAG):
#property
def tasks(self) -> List[BaseOperator]:
return [self.pre_graph] + list(self.task_dict.values()) + [self.post_graph]
#property
def pre_graph(self):
#whatever crazy things you want to do here, before DAG starts
pass
#property
def post_graph(self):
#whatever crazy things you want to do here, AFTER DAG finishes
pass
That's the easiest I can think of, then you just import it when defining your dags:
from custom_dag import PreAndPostDAG
with PreAndPostDAG(
'LS',
default_args=default_args,
description='A simple tutorial DAG',
schedule_interval=timedelta(days=1),
start_date=days_ago(2),
tags=['example'],
) as dag:
# t1, t2 and t3 are examples of tasks created by instantiating operators
t1 = BashOperator(
task_id='list',
bash_command='ls',
)
You get the rest, hope that helps
I'm new to Airflow and I'm currently building a DAG that will execute a PythonOperator, a BashOperator, and then another PythonOperator structured like this:
def authenticate_user(**kwargs):
...
list_prev = [...]
AUTHENTICATE_USER = PythonOperator(
task_id='AUTHENTICATE_USER',
python_callable=authenticate_user,
provide_context=True,
dag=dag)
CHANGE_ROLE = BashOperator(
task_id='CHANGE_ROLE',
bash_command='...',
dag=dag)
def calculations(**kwargs):
list_prev
...
CALCULATIONS = PythonOperator(
task_id='CALCULATIONS',
python_callable=calculations,
provide_context=True,
dag=dag)
My issue is, I create a list of variables in the first PythonOperator (AUTHENTICATE_USER) that I would like to use later in my second PythonOperator (CALCULATIONS) after executing the BashOperator (CHANGE_ROLE). Is there a way for me to carry over that created list into other PythonOperators in my current DAG?
Thank you
I can think of 3 possible ways (to avoid confusion with the Airflow's concept of Variable, I'll call the data that you want to share between tasks as values)
Airflow XCOMs: Push your values from AUTHENTICATE_USER task and pull them in your CALCULATIONS task. You can either publish and access each value separately or wrap them all into a Python dict or list (better as it reduces db reads and writes)
External system: Persist your values from 1st task into some external system such as database, files or S3-objects and access them from downstream tasks when needed
Airflow Variables: This is a specific case of point (2) above (as Variables are stored in Airflow's backend meta-db). You can programmatically create, modify or delete Variables by exploiting the underlying SQLAlchemy model. See this for hints.
Lets say I have two DAG, where dag2 executed dag1 as part of it's flow using TriggerDagRunOperator as follows:
dag1: task1 > task2 > task3
dag2: task4 > dag1 > task5
Now lets say dag2 is scheduled for once a day at 5PM.
Is there a way for me to get the execution timestamp for dag2 (the parent DAG) while I'm running dag1?
Is there any built-in parameter that holds that value?
And if something happened and dag2 was triggered later than usual, lets say 6PM same day, then I still want to get the original scheduling time - that is 5PM while I'm in dag1.
Pass a function to the python_callable argument of TriggerDagRunOperator that injects the execution_date into the triggered DAG:
def inject_execution_date(context, dag_run_obj):
dag_run_obj.payload = {"parent_execution_date": context["execution_date"]}
return dag_run_obj
[...]
trigger_dro = TriggerDagRunOperator(python_callable=inject_execution_date, [...])
You can access this in the child DAG with context["conf"]["parent_execution_date"]