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

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

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

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

Hint

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

Practice

1. What is the main purpose of creating a custom operator in Airflow?

easy

A. To install new Airflow plugins automatically

B. To change the Airflow webserver UI appearance

C. To configure the Airflow scheduler settings

D. To define reusable tasks with custom logic inside the execute method

Solution

1. Step 1: Understand what custom operators do

   Custom operators let you create tasks that run your own code inside Airflow workflows.

2. Step 2: Identify where custom logic is placed

   This logic is placed inside the execute method of the operator class.

3. Final Answer:

   To define reusable tasks with custom logic inside the execute method -> Option D

4. Quick Check:

   Custom operator = reusable task logic [OK]

Hint: Custom operators run your code in the execute() method [OK]

Common Mistakes:

• Confusing UI or scheduler settings with operators
• Thinking operators install plugins
• Believing custom operators change Airflow core behavior

2. Which of the following is the correct way to start defining a custom operator class in Airflow?

easy

A. operator MyOperator extends BaseOperator:

B. class MyOperator(BaseOperator):

C. class MyOperator():

D. def MyOperator():

Solution

1. Step 1: Recall Python class syntax for inheritance

   In Python, to inherit from a class, use class ClassName(BaseClass):.

2. Step 2: Match Airflow custom operator syntax

   Airflow custom operators inherit from BaseOperator, so the correct syntax is class MyOperator(BaseOperator):.

3. Final Answer:

   class MyOperator(BaseOperator): -> Option B

4. Quick Check:

   Python class inheritance = class MyOperator(BaseOperator): [OK]

Hint: Custom operators inherit BaseOperator with class syntax [OK]

Common Mistakes:

• Using function definition instead of class
• Omitting BaseOperator inheritance
• Using incorrect syntax like 'operator' keyword

3. Given this custom operator code snippet, what will be printed when the task runs?

from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults

class PrintOperator(BaseOperator):
    @apply_defaults
    def __init__(self, message: str, **kwargs):
        super().__init__(**kwargs)
        self.message = message

    def execute(self, context):
        print(f"Message: {self.message}")

# Task instance runs with message='Hello Airflow'

medium

A. Message: Hello Airflow

B. Hello Airflow

C. Message: message

D. Error: missing execute method

Solution

1. Step 1: Understand the execute method output

   The execute method prints the string with prefix "Message:" followed by the message attribute.

2. Step 2: Substitute the message value

   The task runs with message='Hello Airflow', so the print outputs "Message: Hello Airflow".

3. Final Answer:

   Message: Hello Airflow -> Option A

4. Quick Check:

   execute prints 'Message: ' + message [OK]

Hint: execute() prints 'Message: ' plus the message string [OK]

Common Mistakes:

• Forgetting the 'Message:' prefix in print
• Confusing attribute name with string literal
• Assuming execute method is missing

4. What is wrong with this custom operator code?

from airflow.models import BaseOperator

class BadOperator(BaseOperator):
    def __init__(self, param):
        self.param = param

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

medium

A. param should be a class variable, not instance variable

B. execute method should not have context parameter

C. Missing call to super().__init__() in __init__ method

D. BaseOperator cannot be inherited

Solution

1. Step 1: Check __init__ method for BaseOperator inheritance

   When inheriting BaseOperator, you must call super().__init__(**kwargs) to initialize the base class properly.

2. Step 2: Identify missing super() call

   The code does not call super().__init__()</>, so Airflow internal setup will fail.

3. Final Answer:

   Missing call to super().__init__() in __init__ method -> Option C

4. Quick Check:

   Always call super().__init__() in custom operator init [OK]

Hint: Always call super().__init__() in __init__ [OK]

Common Mistakes:

• Omitting super().__init__() call
• Removing context parameter from execute
• Misunderstanding inheritance rules

5. You want to create a custom operator that accepts a list of filenames and prints each filename during execution. Which code snippet correctly implements this behavior?

hard

A. class FilePrinterOperator(BaseOperator):
@apply_defaults
def __init__(self, files: list, **kwargs):
super().__init__(**kwargs)
self.files = files
def execute(self, context):
for f in self.files:
print(f)

B. class FilePrinterOperator(BaseOperator):
def __init__(self, files):
self.files = files
def execute(self):
for f in files:
print(f)

C. class FilePrinterOperator(BaseOperator):
@apply_defaults
def __init__(self, files: str, **kwargs):
super().__init__(**kwargs)
self.files = files
def execute(self, context):
print(self.files)

D. class FilePrinterOperator(BaseOperator):
@apply_defaults
def __init__(self, files: list, **kwargs):
super().__init__(**kwargs)
self.files = files
def execute(self, context):
print(files)

Solution

1. Step 1: Check constructor and inheritance

   class FilePrinterOperator(BaseOperator): @apply_defaults def __init__(self, files: list, **kwargs): super().__init__(**kwargs) self.files = files def execute(self, context): for f in self.files: print(f) correctly uses @apply_defaults, calls super().__init__(**kwargs), and stores the list in self.files.

2. Step 2: Verify execute method logic

   class FilePrinterOperator(BaseOperator): @apply_defaults def __init__(self, files: list, **kwargs): super().__init__(**kwargs) self.files = files def execute(self, context): for f in self.files: print(f) loops over self.files and prints each filename, matching the requirement.

3. Final Answer:

   Option A correctly implements the custom operator to print each filename -> Option A

4. Quick Check:

   Use self.files list and loop in execute() [OK]

Hint: Use self.files list and loop inside execute() [OK]

Common Mistakes:

• Not calling super().__init__()
• Using wrong parameter types
• Referencing variables without self.
• Missing context parameter in execute