Creating custom operators in Apache Airflow - Step-by-Step CLI Walkthrough

Introduction

Sometimes the built-in tasks in Airflow do not fit your specific needs. Creating custom operators lets you define your own task behavior to automate unique workflows easily.

When you need to run a task that is not covered by Airflow's default operators.

When you want to reuse a complex task logic across multiple workflows.

When you want to simplify your DAG code by encapsulating task details.

When you want to add custom logging or error handling to a task.

When integrating with a new system or API that Airflow does not support out of the box.

Config File - my_custom_operator.py

my_custom_operator.py

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

class MyCustomOperator(BaseOperator):
    @apply_defaults
    def __init__(self, my_param, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.my_param = my_param

    def execute(self, context):
        print(f"Running custom operator with parameter: {self.my_param}")
        # Add your custom task logic here
        return f"Task completed with {self.my_param}"

This file defines a new operator class MyCustomOperator that inherits from Airflow's BaseOperator. The __init__ method sets up parameters, and the execute method contains the task's main logic. This operator can be imported and used in DAGs like any built-in operator.

Commands

Check the list of available DAGs to confirm Airflow is running and ready.

Terminal

airflow dags list

Expected Output

example_dag

Trigger a DAG that uses the custom operator to run the task immediately.

Terminal

airflow dags trigger example_custom_operator_dag

Expected Output

Created <DagRun example_custom_operator_dag @ 2024-06-01 12:00:00+00:00: manual__2024-06-01T12:00:00+00:00, externally triggered: True>

List all tasks in the DAG to verify the custom operator task is present.

Terminal

airflow tasks list example_custom_operator_dag

Expected Output

my_custom_task

Run the custom operator task locally for the given execution date to test its behavior.

Terminal

airflow tasks test example_custom_operator_dag my_custom_task 2024-06-01

Expected Output

Running custom operator with parameter: test_value

[2024-06-01 12:00:00,000] {taskinstance.py:XXXX} INFO - Task completed with test_value

Key Concept

If you remember nothing else from this pattern, remember: custom operators let you package unique task logic into reusable, clean components for your Airflow workflows.

Common Mistakes

Not calling super().__init__() in the custom operator's __init__ method

This causes Airflow to miss important setup steps, leading to errors or missing features.

Always call super().__init__(*args, **kwargs) inside your __init__ method.

Putting task logic outside the execute() method

Airflow runs the execute() method to perform the task; logic outside it won't run during task execution.

Put all task-specific code inside the execute() method.

Not importing the custom operator correctly in the DAG file

Airflow won't recognize the operator and will fail to parse the DAG.

Ensure the custom operator file is in the PYTHONPATH or the DAG folder and import it properly.

Summary

Create a Python class inheriting from BaseOperator to define a custom operator.

Implement the execute() method with the task's main logic.

Use airflow CLI commands to list DAGs, trigger runs, list tasks, and test tasks locally.

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