Creating custom operators in Apache Airflow - Try It Yourself

Creating Custom Operators in Airflow

📖 Scenario: You are working as a data engineer. You want to automate a task in Airflow that prints a custom message. Instead of using the default operators, you decide to create your own custom operator to reuse it easily in many workflows.

🎯 Goal: Build a simple custom Airflow operator that prints a message. Then use it in a DAG to see the message printed when the task runs.

📋 What You'll Learn

Create a custom operator class inheriting from BaseOperator

Add an __init__ method to accept a message parameter

Override the execute method to print the message

Create a DAG that uses this custom operator

Run the DAG to see the printed message

💡 Why This Matters

🌍 Real World

Custom operators let you package repeated or complex tasks in Airflow so you can reuse them easily across many workflows.

💼 Career

Data engineers and DevOps professionals often create custom operators to automate and standardize data pipelines and workflows.

1

Create the custom operator class

Create a Python class called PrintMessageOperator that inherits from BaseOperator. Add an __init__ method that takes self, message, and **kwargs. Inside __init__, call super().__init__(**kwargs) and assign self.message = message.

Apache Airflow

from airflow.models import BaseOperator

# Define PrintMessageOperator class here
# Your code here

Need a hint?

Remember to inherit from BaseOperator and call super().__init__(**kwargs) inside your constructor.

2

Add the execute method

Inside the PrintMessageOperator class, add a method called execute that takes self and context. Inside this method, use print(self.message) to display the message.

Apache Airflow

from airflow.models import BaseOperator

class PrintMessageOperator(BaseOperator):
    def __init__(self, message, **kwargs):
        super().__init__(**kwargs)
        self.message = message

    # Add execute method here
    # Your code here

Need a hint?

The execute method is where the operator does its work. Use print(self.message) to show the message.

3

Create a DAG using the custom operator

Import DAG from airflow and datetime from datetime. Create a DAG object called dag with dag_id='print_message_dag', start_date=datetime(2024, 1, 1), and schedule_interval='@once'. Then create a task called print_task using PrintMessageOperator with task_id='print_task', message='Hello from custom operator!', and dag=dag.

Apache Airflow

from airflow.models import BaseOperator
from airflow import DAG
from datetime import datetime

class PrintMessageOperator(BaseOperator):
    def __init__(self, message, **kwargs):
        super().__init__(**kwargs)
        self.message = message

    def execute(self, context):
        print(self.message)

# Create DAG and task here
# Your code here

Need a hint?

Use DAG to create a workflow and assign your custom operator as a task inside it.

4

Run the DAG and print the message

Add a line to print the message by calling print_task.execute(context={}). This simulates running the task and shows the message.

Apache Airflow

from airflow.models import BaseOperator
from airflow import DAG
from datetime import datetime

class PrintMessageOperator(BaseOperator):
    def __init__(self, message, **kwargs):
        super().__init__(**kwargs)
        self.message = message

    def execute(self, context):
        print(self.message)


dag = DAG(dag_id='print_message_dag', start_date=datetime(2024, 1, 1), schedule_interval='@once')

print_task = PrintMessageOperator(task_id='print_task', message='Hello from custom operator!', dag=dag)

# Run the task to print the message
# Your code here

Need a hint?

Call execute on your task with an empty context dictionary to see the message printed.