TaskFlow API for cleaner XCom in Apache Airflow - Commands & Configuration

Introduction

Airflow's TaskFlow API helps you write workflows where tasks can easily share data without messy code. It solves the problem of passing data between tasks in a clean and readable way.

When you want to pass results from one task to another without using manual XCom push and pull calls.

When you want your workflow code to look like regular Python functions with clear inputs and outputs.

When you want to avoid writing extra code to handle XCom serialization and deserialization.

When you want to improve readability and maintainability of your Airflow DAGs.

When you want to use decorators to define tasks simply and clearly.

Config File - taskflow_dag.py

taskflow_dag.py

from airflow.decorators import dag, task
from airflow.utils.dates import days_ago

@dag(start_date=days_ago(1), schedule_interval='@daily', catchup=False)
def example_taskflow_dag():

    @task
def extract():
        return {'name': 'Airflow', 'version': 2}

    @task
def transform(data):
        return f"{data['name']} version {data['version']}"

    @task
def load(message):
        print(f"Loading message: {message}")

    data = extract()
    message = transform(data)
    load(message)

example_taskflow_dag()

This Python file defines an Airflow DAG using the TaskFlow API.

@dag decorates the main function to define the workflow.

@task decorates Python functions to make them Airflow tasks that automatically handle XComs.

Data flows naturally as function return values and arguments, making the code clean and easy to follow.

Commands

List all available DAGs to verify that the new DAG is recognized by Airflow.

Terminal

airflow dags list

Expected Output

example_taskflow_dag

Trigger the DAG run manually to start the workflow and see the tasks execute.

Terminal

airflow dags trigger example_taskflow_dag

Expected Output

Created <DagRun example_taskflow_dag @ 2024-06-01T00:00:00+00:00: manual__2024-06-01T00:00:00+00:00, externally triggered: True>

List all tasks in the DAG to confirm the tasks are defined correctly.

Terminal

airflow tasks list example_taskflow_dag

Expected Output

extract

transform

load

Run the 'extract' task locally for the given date to test its output and XCom push.

Terminal

airflow tasks test example_taskflow_dag extract 2024-06-01

Expected Output

[2024-06-01 00:00:00,000] {taskinstance.py:876} INFO - Executing <Task(TaskFlow) object: extract> on 2024-06-01

[2024-06-01 00:00:00,100] {taskinstance.py:1123} INFO - Marking task as SUCCESS. dag_id=example_taskflow_dag, task_id=extract, execution_date=2024-06-01

Key Concept

If you remember nothing else from this pattern, remember: TaskFlow API lets you pass data between tasks as simple Python function inputs and outputs without manual XCom handling.

Common Mistakes

Trying to manually push or pull XComs inside TaskFlow API tasks.

TaskFlow API automatically handles XComs for you, so manual calls cause confusion and errors.

Return values from @task functions and pass them as arguments to other @task functions.

Not decorating functions with @task inside the DAG function.

Without @task, Airflow does not treat the function as a task, so no XCom or scheduling happens.

Always use @task decorator on functions that represent tasks.

Defining tasks outside the DAG function decorated with @dag.

Tasks must be defined inside the DAG function to be part of that DAG's context.

Define all @task functions inside the @dag decorated function.

Summary

Use @dag to define your workflow and @task to define tasks as Python functions.

Return values from tasks automatically become XComs and can be passed as inputs to other tasks.

Trigger and test your DAG using airflow CLI commands to verify task execution and data passing.

Practice

1. What is the main benefit of using the TaskFlow API in Airflow for passing data between tasks?

easy

A. It automatically handles data sharing between tasks using return values.

B. It requires manual XCom push and pull commands for data sharing.

C. It disables task dependencies to speed up execution.

D. It only works with BashOperator tasks.

Solution

1. Step 1: Understand TaskFlow API purpose

   The TaskFlow API simplifies task communication by using Python functions and decorators.

2. Step 2: Recognize automatic XCom handling

   Return values from decorated tasks are automatically shared with downstream tasks via XCom without manual commands.

3. Final Answer:

   It automatically handles data sharing between tasks using return values. -> Option A

4. Quick Check:

   TaskFlow API = automatic XCom via return [OK]

Hint: TaskFlow API shares data automatically via return values [OK]

Common Mistakes:

• Thinking manual XCom push/pull is needed
• Believing it disables dependencies
• Assuming it only works with BashOperator

2. Which of the following is the correct way to define a task using the TaskFlow API in Airflow?

easy

A. def my_task(): pass

B. @dag
def my_task(): pass

C. @task
def my_task(): pass

D. my_task = PythonOperator(task_id='my_task')

Solution

1. Step 1: Identify TaskFlow API syntax

   Tasks are defined as Python functions decorated with @task to enable TaskFlow features.

2. Step 2: Compare options

   @task def my_task(): pass uses @task decorator correctly; others either lack decorator or use wrong decorators/operators.

3. Final Answer:

   @task\ndef my_task(): pass -> Option C

4. Quick Check:

   TaskFlow task = function + @task decorator [OK]

Hint: Use @task decorator above function to define TaskFlow task [OK]

Common Mistakes:

• Omitting @task decorator
• Confusing @dag with @task
• Using PythonOperator instead of TaskFlow

3. Given the following code snippet using TaskFlow API, what will be printed?

@task
def add(x, y):
    return x + y

@task
def print_result(value):
    print(f'Result is {value}')

with DAG('example_dag', start_date=datetime(2023,1,1)) as dag:
    sum_value = add(3, 4)
    print_result(sum_value)

medium

A. No output printed

B. Result is

C. SyntaxError

D. Result is 7

Solution

1. Step 1: Understand TaskFlow return and XCom

   The add task returns 7, which is passed automatically to print_result.

2. Step 2: Trace print statement execution

   print_result prints the string with the value 7, so output is 'Result is 7'.

3. Final Answer:

   Result is 7 -> Option D

4. Quick Check:

   Returned value passed and printed = 'Result is 7' [OK]

Hint: Returned values flow automatically and print as expected [OK]

Common Mistakes:

• Expecting raw XCom object printed
• Confusing syntax error with correct code
• Thinking no output prints without manual XCom pull

4. What is wrong with this TaskFlow API code snippet?

@task
def multiply(x, y):
    print(x * y)

with DAG('dag', start_date=datetime(2023,1,1)) as dag:
    result = multiply(2, 5)
    print(f'Result is {result}')

medium

A. The DAG is missing the schedule_interval parameter.

B. The multiply task should return the result instead of printing it.

C. The print statement inside the DAG is invalid syntax.

D. The @task decorator is missing parentheses.

Solution

1. Step 1: Identify how TaskFlow passes data

   TaskFlow passes data via return values, not print statements inside tasks.

2. Step 2: Analyze code behavior

   Since multiply only prints and returns None, result is None, so printing it shows 'Result is None'.

3. Final Answer:

   The multiply task should return the result instead of printing it. -> Option B

4. Quick Check:

   Return value needed for XCom, not print [OK]

Hint: Return values to share data, not print inside tasks [OK]

Common Mistakes:

• Thinking print shares data between tasks
• Confusing DAG parameters with task errors
• Believing @task needs parentheses always

5. You want to create a DAG using TaskFlow API where task A returns a list of numbers, and task B receives this list and returns the sum. Which code snippet correctly implements this pattern?

hard

A. @task
def task_a():
return [1, 2, 3]

@task
def task_b(numbers):
return sum(numbers)

with DAG('dag', start_date=datetime(2023,1,1)) as dag:
nums = task_a()
total = task_b(nums)

B. @task
def task_a():
print([1, 2, 3])

@task
def task_b(numbers):
print(sum(numbers))

with DAG('dag', start_date=datetime(2023,1,1)) as dag:
nums = task_a()
total = task_b(nums)

C. @task
def task_a():
return '1,2,3'

@task
def task_b(numbers):
return sum(numbers)

with DAG('dag', start_date=datetime(2023,1,1)) as dag:
nums = task_a()
total = task_b(nums)

D. @task
def task_a():
return [1, 2, 3]

@task
def task_b(numbers):
return sum(numbers)

with DAG('dag', start_date=datetime(2023,1,1)) as dag:
total = task_b(task_a())

Solution

1. Step 1: Check task A returns a list

   @task
   def task_a():
       return [1, 2, 3]
   
   @task
   def task_b(numbers):
       return sum(numbers)
   
   with DAG('dag', start_date=datetime(2023,1,1)) as dag:
       nums = task_a()
       total = task_b(nums)

   returns a list [1, 2, 3] correctly;

   @task
   def task_a():
       print([1, 2, 3])
   
   @task
   def task_b(numbers):
       print(sum(numbers))
   
   with DAG('dag', start_date=datetime(2023,1,1)) as dag:
       nums = task_a()
       total = task_b(nums)

   prints instead of returning;

   @task
   def task_a():
       return '1,2,3'
   
   @task
   def task_b(numbers):
       return sum(numbers)
   
   with DAG('dag', start_date=datetime(2023,1,1)) as dag:
       nums = task_a()
       total = task_b(nums)

   returns a string, not list.

2. Step 2: Check task B sums the list argument

   @task
   def task_a():
       return [1, 2, 3]
   
   @task
   def task_b(numbers):
       return sum(numbers)
   
   with DAG('dag', start_date=datetime(2023,1,1)) as dag:
       nums = task_a()
       total = task_b(nums)

   passes the list to task B which sums it correctly;

   @task
   def task_a():
       return [1, 2, 3]
   
   @task
   def task_b(numbers):
       return sum(numbers)
   
   with DAG('dag', start_date=datetime(2023,1,1)) as dag:
       total = task_b(task_a())

   skips assigning intermediate variable but is valid but less clear.

3. Final Answer:

   Option A correctly returns a list and passes it to sum in task B. -> Option A

4. Quick Check:

   Return list and pass it to sum task [OK]

Hint: Return data structures, pass as args, sum in next task [OK]

Common Mistakes:

• Using print instead of return in tasks
• Passing string instead of list to sum
• Ignoring task dependencies or variable assignment