Writing a Custom Operator

Concepts covered: DAG definition

Promoting a function to an operator

A @task function covers most needs: write Python, return a value, done. A custom operator is worth writing when the same body with the same 2 or 3 parameters is being pasted into DAG after DAG, when the task should look and behave like a built-in in the graph with its own color and rendered fields, or when stopping it from outside must clean something up. Write the function first, and promote it to an operator the 3rd time you paste it. The example below counts the rows in a CSV file and fails the task when the count is below a floor, a check 14 feed loads all need.

import csv

from airflow.sdk import BaseOperator


class CsvRowCountOperator(BaseOperator):
    template_fields = ("path",)
    ui_color = "#dbeafe"

    def __init__(self, *, path, min_rows=1, **kwargs):
        super().__init__(**kwargs)
        self.path = path
        self.min_rows = min_rows

    def execute(self, context):
        with open(self.path, newline="") as f:
            count = sum(1 for _ in csv.reader(f)) - 1
        self.log.info("%s has %d rows", self.path, count)
        if count < self.min_rows:
            raise ValueError(f"{self.path}: {count} rows, need {self.min_rows}")
        return count

The class extends BaseOperator from airflow.sdk, which is what makes it a task the scheduler can plan and a worker can run. __init__ takes its own arguments as keywords and stores them on self. Every other keyword, task_id included, goes to super().__init__. It opens no file and reaches no service, because __init__ runs in the DAG processor on every parse cycle, 2 times a minute for as long as the file exists, and again on the worker each time it loads the file to run a task instance, once per try. A lookup placed there is paid by both processes on every one of those occasions. execute is the body: it runs once per try on a worker, receives the same context dict a @task function gets, and whatever it returns flows downstream the way a @task return value does. Raising any exception fails the try.

The contract a built-in operator keeps

• __init__(self, *, ..., kwargs): Keyword-only arguments, super().__init__(kwargs) first, attributes stored, nothing opened.
• execute(self, context): The whole body. Return a value for downstream tasks, raise to fail the try.
• template_fields: Strings stored under these names have their {{ }} placeholders filled in before execute runs, and the UI shows the rendered result.
• ui_color: The box color in the graph, so the eye finds every row-count check at once.
• on_kill(self): Close what execute opened: a subprocess, a remote job, a lock.

Using it, once and 40 times

A DAG file calls the operator with CsvRowCountOperator(task_id="check_acme", path="feeds/acme.csv", min_rows=100), just as it would a built-in. Mapping it works the same way too. The file below checks every file list_files returns, which ties this section to the first one: partial carries task_id and min_rows, expand carries the paths, and the grid shows 1 row named check with the count of copies beside it.

from airflow.sdk import dag, task
from operators.csv_ops import CsvRowCountOperator


@dag(schedule=None)
def check_feeds():
    @task
    def list_files():
        return ["feeds/acme.csv", "feeds/globex.csv"]

    CsvRowCountOperator.partial(task_id="check", min_rows=100).expand(
        path=list_files()
    )


check_feeds()

Where the module lives matters because 2 different processes import it. The DAG processor imports it to build the task; the worker imports it to run execute. A module inside the bundle, dags/operators/csv_ops.py with an __init__.py beside it, reaches both, because both fetch the bundle and the bundle's root is on the import path. A package installed with pip on every host reaches both as well and is the right home once 2 bundles share the operator. A module that only the processor can see produces a DAG that parses and then fails on the worker with an import error at the top of every try.

Do

• Extend airflow.sdk.BaseOperator and call super().__init__(**kwargs) first
• List every attribute that holds a path or query in template_fields
• Return the result from execute so downstream tasks can use it
• Put the module where the processor and every worker can import it

Avoid

• Open a file in __init__, or a socket or session
• Compute anything in __init__ that depends on the run
• Give the operator a positional argument
• Catch every exception in execute; a swallowed failure shows as success

About this lesson

This section is part of the DAG Anatomy: Advanced lesson on DataDriven, the home of data engineering. Lessons cover SQL, Python, Spark, data modeling and pipeline architecture, and each section pairs its explanation with worked examples and practice you run as you go.