Cutting-Edge Data Engineering with Airflow ETL Pipelines#

Step 1: Environment and Dependencies#

Airflow runs on Python. It’s recommended to create a clean virtual environment:

1
# Create a new virtual environment
2
python -m venv airflow_env
3
source airflow_env/bin/activate
4

5
# Install or upgrade pip
6
pip install --upgrade pip setuptools wheel

Step 2: Install Airflow#

You can install Airflow from PyPI. It’s best to pin specific versions to ensure consistency if you’re working across multiple teams.

1
pip install apache-airflow==2.4.*

Step 3: Initialize the Airflow Database#

Airflow requires a metadata database to store information about DAG runs, tasks, variables, and logs. You can initialize a local SQLite database for development, but for production, it is recommended to use PostgreSQL or MySQL.

1
# Initialize the database
2
airflow db init

Step 4: Create a User#

1
airflow users create \
2
    --username admin \
3
    --firstname John \
4
    --lastname Doe \
5
    --role Admin \
6
    --email admin@example.com

Step 5: Start Airflow Scheduler and Webserver#

To kick things off, run the scheduler and web server:

1
# Terminal 1: Start scheduler
2
airflow scheduler
3

4
# Terminal 2: Start webserver
5
airflow webserver --port 8080

You can now access the Airflow UI at http://localhost:8080.

DAGs#

The Directed Acyclic Graph (DAG) is Airflow’s core concept. It is a logical organization of tasks that must be executed in a certain order without any loops or cycles. DAGs are defined in Python files, which are scanned and loaded into Airflow.

1
from datetime import datetime
2
from airflow import DAG
3
from airflow.operators.bash import BashOperator
4

5
default_args = {
6
    'start_date': datetime(2023, 1, 1),
7
    'retries': 1,
8
}
9

10
with DAG(
11
    dag_id='example_dag',
12
    default_args=default_args,
13
    schedule_interval='@daily',
14
    catchup=False
15
) as dag:
16

17
    task_1 = BashOperator(
18
        task_id='bash_task',
19
        bash_command='echo "Hello, Airflow!"'
20
    )
21

22
    task_1

Tasks and Operators#

Within a DAG, there are operators that define specific actions, such as running Python functions, executing SQL queries, or triggering external services. Tasks are the instances of operators.

Commonly Used Operators:

• BashOperator: Runs a bash command or script.
• PythonOperator: Executes a specific Python callable function.
• TriggerDagRunOperator: Triggers other DAGs.
• SqlSensor, PostgresOperator, HiveOperator, SnowflakeOperator: Database-related tasks.

Task Dependencies#

One of the major features of Airflow is its intuitive approach to specifying dependencies. You can define tasks that run sequentially (e.g., task_1 >> task_2) or in parallel.

1
from airflow.operators.empty import EmptyOperator
2

3
start = EmptyOperator(task_id='start', dag=dag)
4
middle = EmptyOperator(task_id='middle', dag=dag)
5
end = EmptyOperator(task_id='end', dag=dag)
6

7
start >> middle >> end  # This sets the execution order

Hooks and Connections#

Airflow provides Hooks to abstract connections to external systems like databases, cloud services, or APIs. For instance, PostgresHook or MySqlHook handle authentication details and provide a structured way to run queries.

You can manage connections securely in the Airflow UI under the Admin > Connections section.

Extract#

1. List new files in the S3 bucket
2. Download these files to a local or temporary directory

1
from airflow import DAG
2
from airflow.providers.amazon.aws.operators.s3 import S3ListOperator, S3DownloadOperator
3
from datetime import datetime
4

5
default_args = {
6
    'start_date': datetime(2023, 1, 1),
7
    'retries': 1,
8
}
9

10
with DAG(
11
    dag_id='simple_etl',
12
    default_args=default_args,
13
    schedule_interval='@daily',
14
    catchup=False
15
) as dag:
16

17
    list_s3_files = S3ListOperator(
18
        task_id='list_s3_files',
19
        bucket='my-data-bucket',
20
        prefix='incoming_csv/'
21
    )
22

23
    download_s3_files = S3DownloadOperator(
24
        task_id='download_s3_files',
25
        bucket='my-data-bucket',
26
        s3_keys="{{ task_instance.xcom_pull(task_ids='list_s3_files') }}",
27
        local_filepath='/tmp/'
28
    )

Transform#

Data transformation can be handled in multiple ways. You could use a PythonOperator, BashOperator to call a Spark job, or a specialized plugin. Here, we’ll do a quick example in Python.

1
from airflow.operators.python import PythonOperator
2
import pandas as pd
3
import os
4

5
def transform_data(**context):
6
    s3_file_list = context['ti'].xcom_pull(task_ids='list_s3_files')
7
    for file_name in s3_file_list:
8
        full_path = os.path.join('/tmp', file_name.split('/')[-1])
9
        df = pd.read_csv(full_path)
10
        # Example transformation: dropping rows with null values
11
        df.dropna(inplace=True)
12
        df.to_csv(full_path, index=False)
13

14
transform_task = PythonOperator(
15
    task_id='transform_data',
16
    python_callable=transform_data,
17
    provide_context=True
18
)

Load#

For loading into PostgreSQL, you could use a PostgresHook or PostgresOperator. Here’s a simple illustration:

1
from airflow.providers.postgres.operators.postgres import PostgresOperator
2

3
create_table = PostgresOperator(
4
    task_id='create_table',
5
    postgres_conn_id='my_postgres_conn',
6
    sql="""
7
    CREATE TABLE IF NOT EXISTS my_data (
8
        column1 VARCHAR(255),
9
        column2 VARCHAR(255),
10
        ...
11
    );
12
    """
13
)
14

15
load_data = PostgresOperator(
16
    task_id='load_data',
17
    postgres_conn_id='my_postgres_conn',
18
    sql="""
19
    COPY my_data FROM '/tmp/my_data.csv' DELIMITER ',' CSV HEADER;
20
    """
21
)

Link the Tasks#

Putting it all together:

1
list_s3_files >> download_s3_files >> transform_task >> create_table >> load_data

And that’s your basic ETL flow. This example can be expanded upon or refactored based on your data complexities or organizational requirements.

Airflow Executors#

Executors define how Airflow distributes the work:

Choosing the executor depends on your infrastructure, team size, and the complexity of your pipelines.

Templating and Macros#

Airflow uses the Jinja templating engine. This allows you to dynamically insert parameters:

1
from airflow.operators.bash import BashOperator
2

3
bash_task = BashOperator(
4
    task_id='templated_task',
5
    bash_command="""
6
        echo "Current execution date is {{ ds }}"
7
        echo "Previous execution date is {{ prev_ds }}"
8
    """,
9
    dag=dag
10
)

You can also leverage custom macros and environment variables to parameterize your tasks.

XCom#

XCom (cross-communication) is Airflow’s way to share data between tasks. Instead of storing large data sets, it’s intended for small pieces of metadata (e.g., file names, record counts).

1
def push_function(**context):
2
    context['ti'].xcom_push(key='filename', value='file_20230101.csv')
3

4
def pull_function(**context):
5
    data = context['ti'].xcom_pull(key='filename', task_ids='push_task')
6

7
push_task = PythonOperator(
8
    task_id='push_task',
9
    python_callable=push_function
10
)
11

12
pull_task = PythonOperator(
13
    task_id='pull_task',
14
    python_callable=pull_function
15
)
16

17
push_task >> pull_task

Variables#

Sometimes you need static or semi-static values such as directory paths, connection strings, or configuration flags. Airflow Variables allow you to store and retrieve these values globally.

1
from airflow.models import Variable
2

3
my_var = Variable.get("MY_VARIABLE_NAME")

Airflow Sensors#

Sensors are specialized operators that wait for a certain criterion to be met before proceeding. For example, an S3KeySensor can wait until a file appears in a bucket:

1
from airflow.providers.amazon.aws.sensors.s3 import S3KeySensor
2

3
wait_for_s3 = S3KeySensor(
4
    task_id='wait_for_s3_file',
5
    bucket_key='incoming_csv/{{ ds }}/file_{{ ds_nodash }}.csv',
6
    bucket_name='my-data-bucket',
7
    aws_conn_id='aws_default',
8
    poke_interval=60,  # how often to check in seconds
9
    timeout=3600       # how long to wait for file
10
)

Celery Executor#

For larger production environments, the CeleryExecutor is often used. It employs a queue-based architecture where a broker (e.g., RabbitMQ or Redis) handles task dispatching, and multiple worker processes pull tasks from the queue.

1. Airflow Scheduler enqueues tasks.
2. Celery Workers pick up tasks from the queue.
3. Result Backend (e.g., Redis or database) stores completed task statuses.

Kubernetes Executor#

If your organization runs on Kubernetes, the KubernetesExecutor allows tasks to be run inside isolated pods, providing fine-grained resource control. Each task can run in its own environment, simplifying dependency management and improving security.

Monitoring and Alerting#

Reliable monitoring is critical for production. Airflow integrates with external services (PagerDuty, Slack, email) to send alerts on task failures or SLA misses. You can also configure Airflow’s metrics to feed into solutions like Prometheus and Grafana.

Observability and Logging#

• Centralized Logging: Store logs in a central location (S3, GCS, or a logging service) to ensure easy access and searchability.
• Structured Logging: Log key details (file names, transaction IDs, etc.) to enable better filtering and correlation.
• Task Duration Metrics: Keep an eye on average run times to proactively detect performance bottlenecks.

Testing and CI/CD#

• Unit Tests: For PythonOperators, ensure the logic is tested in isolation.
• DAG Validation: Your CI pipeline can run airflow dags test to validate DAG integrity.
• Linting: Tools like flake8 and black can ensure code quality.

Version Control and Environments#

• Separate Environments: Maintain staging, development, and production Airflow environments.
• GitOps: Store DAG files and configuration in Git. Automate deployment to orchestrate changes seamlessly.
• Tagging: Tag releases to revert or identify pipeline states quickly.