Apache Airflow (7) Custom Operator Plugins

One of key strengths of Airflow is its extensibility, which allows users to customize and extend Airflow’s functionality through plugins. Custom Operators are one of the most common types of plugins and are essential when you need to perform custom actions or integrate with third-party systems.

In this post, we will explain how to create and use custom operator plugins in Apache Airflow and walk through an example where we consolidate repeated code into reusable operator plugins.

What Are Custom Operators?

https://github.com/apache/airflow/tree/123dadda0e0648ef1412053d1743128333eecb63/providers/src/airflow/providers

Operators in Airflow are the building blocks of workflows (DAGs). They define the tasks that need to be executed. A custom operator is a class that extends an existing Airflow operator (or BaseOperator) to introduce additional functionality, such as performing custom logic or connecting to external systems.

Steps to Create a Custom Operator Plugin

Here’s a step-by-step guide on how to create a custom plugin and integrate it into your DAGs:

1. Create the Custom Operator Class

To create a custom operator, you need to define a Python class that extends Airflow’s BaseOperator class. Within this class, you will define the functionality of your operator, typically inside an execute method. This method is called when the operator is executed during the DAG run.

Example: Custom Operator to Check Data Quality

We will create a custom operator that checks if a table in Redshift contains more than zero rows. This custom operator will replace a more generic task in our DAG.

from airflow.models import BaseOperator
from airflow.hooks.postgres_hook import PostgresHook
from airflow.utils.decorators import apply_defaults
import logging

class HasRowsOperator(BaseOperator):
    """
    Custom operator to check that a table in Redshift contains rows.
    """
    @apply_defaults
    def __init__(self, table, redshift_conn_id, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.table = table
        self.redshift_conn_id = redshift_conn_id

    def execute(self, context):
        logging.info(f"Checking if table {self.table} has rows...")

        # Establish a connection to Redshift using the PostgresHook
        redshift_hook = PostgresHook(postgres_conn_id=self.redshift_conn_id)

        # Query the table to count rows
        records = redshift_hook.get_records(f"SELECT COUNT(*) FROM {self.table}")

        # Check the result
        if len(records) < 1 or len(records[0]) < 1:
            raise ValueError(f"Data quality check failed. {self.table} returned no results.")

        num_records = records[0][0]
        if num_records < 1:
            raise ValueError(f"Data quality check failed. {self.table} contained 0 rows.")

        logging.info(f"Data quality on table {self.table} check passed with {num_records} records.")

2. Add the Custom Operator to the Plugins Directory

Airflow allows you to extend its functionality by placing custom operator code into the plugins/ folder. Once placed here, Airflow will automatically detect and make these operators available for use in your DAGs.

Assuming you have a folder called custom_operators under the plugins/ directory, your file structure will look like this:

plugins/
    custom_operators/
        __init__.py
        has_rows_operator.py

3. Use the Custom Operator in Your DAG

Now that we’ve defined the custom operator, we can use it in our DAG. In the following example, we replace a simple data quality check with our custom HasRowsOperator to check the row count of tables in Redshift.

import pendulum
from airflow.decorators import dag, task
from airflow.operators.postgres_operator import PostgresOperator
from custom_operators.has_rows import HasRowsOperator
from project.common import sql_statements

@dag(
    start_date=pendulum.now(),
    max_active_runs=1
)
def demonstrate_custom_operators():

    # Task to create the 'traffics' table in Redshift
    create_traffics_table = PostgresOperator(
        task_id="create_traffics_table",
        postgres_conn_id="redshift",
        sql=sql_statements.CREATE_TRAFFICS_TABLE_SQL
    )

    # Task to load data into the 'traffics' table from S3 to Redshift
    copy_trips_task = S3ToRedshiftOperator(
        task_id="load_traffics_from_s3_to_redshift",
        table="traffics",
        redshift_conn_id="redshift",
        aws_credentials_id="aws_credentials",
        s3_bucket="your-bucket-name",
        s3_key="your-s3-key"
    )

    # Custom data quality check using the HasRowsOperator
    check_traffics_task = HasRowsOperator(
        task_id="count_traffics",
        table="traffics",
        redshift_conn_id="redshift",
    )

    create_traffics_table >> copy_traffics_task >> check_traffics_task

demonstrate_custom_operators_dag = demonstrate_custom_operators()

In the above DAG:

• We define a task to create the traffics table in Redshift.
• We load data from an S3 bucket into the traffics table.
• We then use the custom HasRowsOperator to ensure that the traffics table contains data.

4. Test and Extend Your Custom Operator

Once you’ve created the custom operator and added it to your DAG, you can test it by running the DAG. If any issues arise during the execution, you can adjust the execute method of your custom operator to handle errors more gracefully or extend its functionality.

Conclusion

Custom operators are a powerful feature in Apache Airflow that allows you to encapsulate complex logic into reusable components, improving code organization simplifying DAGs. When developing your own custom operators, be sure to:

• Leverage Airflow’s hooks for database and cloud integration
• Test your operator thoroughly to ensure it works across different workflows
• Keep you operator code modular and reusable to maintain clean, maintainable DAGs.

This process is applicable for many use cases and can be extended to handle various tasks like interacting with APIs, querying databases, or managing cloud resources.