From Raw to Refined: Building ETL Pipelines with Airflow#

Prerequisites#

1. Python 3.7 or higher.
2. A virtual environment manager like venv or conda (recommended).
3. pip (to install Airflow and its dependencies).

Though Airflow can be deployed in various ways (locally, Docker, Kubernetes), starting locally is often the simplest for beginners.

Installation Steps#

1. Create and activate a virtual environment.
   Example with venv:

   1
   python -m venv airflow-env
   2
   source airflow-env/bin/activate
   

2. Install Airflow via pip.

   1
   pip install "apache-airflow[postgres,google]>=2.5,<3"
   

   The [postgres,google] part indicates additional dependencies for using Postgres and Google integrations, but you can tailor it to your environment.

3. Initialize Airflow’s metadata database.

   1
   airflow db init
   

4. Create an admin user (for the web interface).

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

5. Start the Airflow scheduler and webserver.

   1
   airflow scheduler &
   2
   airflow webserver
   

   By default, the Airflow UI will be accessible at http://localhost:8080.

Post-Installation Configuration#

• airflow.cfg file: You can change default settings like executor, sql_alchemy_conn, or web_server_port.
• Environment variables: Airflow can read environment variables for settings, which helps when deploying to different environments (e.g., dev, staging, prod).

Directed Acyclic Graph (DAG)#

A DAG is a collection of tasks with dependencies that define how they’re organized and in what order they should run. Each node (task) in the DAG represents work to be done, and edges represent the required sequence.

Tasks#

Within a DAG, a task is a parameterized instance of one of Airflow’s classes that tells Airflow what to do (run a Python function, execute a Bash command, move files, etc.).

Operators#

Operators are templates for tasks. Airflow comes with many operators out of the box:

• BashOperator: Executes a bash command.
• PythonOperator: Runs a Python function.
• PostgresOperator: Executes SQL statements in a Postgres database.
• S3ToRedshiftOperator: Moves data from S3 to Redshift.
• And many others for AWS, GCP, Spark, etc.

In addition, there are advanced concepts like Hooks (wrappers around external APIs or databases) and Sensors (tasks that wait for a certain condition to be met).

Explanation#

• default_args: Arguments passed to each task (owner, start date, retries, etc.).
• dag_id: Unique identifier for the DAG.
• schedule_interval: Defines how often the DAG runs (@daily, @hourly, a cron expression, etc.).
• task_hello >> task_date: The “>>” sets the execution order, ensuring task_hello completes before task_date.

Place this Python file in the dags folder (usually located in ~/airflow/dags by default), and Airflow will automatically detect it. You’ll see it in the Airflow web UI and can trigger or schedule it as needed.

Common Data Sources#

• Relational Databases (Postgres, MySQL, Oracle)
• Cloud Storage (AWS S3, GCS)
• SaaS APIs (Salesforce, Stripe)
• Flat Files (CSV, Parquet)

Airflow offers a variety of operators and hooks to connect to these services. Below is a short table that summarizes Airflow operators/hook pairs for some typical sources:

Example: Extract Data from a Database#

Say you have a Postgres database storing user data. You want to export a table to a CSV file in an S3 bucket daily. Below is a simplified outline of how you’d do it with Airflow:

1
import datetime
2
from airflow import DAG
3
from airflow.operators.python_operator import PythonOperator
4
from airflow.providers.postgres.hooks.postgres import PostgresHook
5
from airflow.providers.amazon.aws.hooks.s3 import S3Hook
6

7
def extract_data_from_postgres():
8
    pg_hook = PostgresHook(postgres_conn_id='my_postgres')
9
    records = pg_hook.get_pandas_df("SELECT * FROM users")
10
    records.to_csv('/tmp/users_data.csv', index=False)
11

12
def upload_to_s3():
13
    s3_hook = S3Hook(aws_conn_id='my_s3')
14
    s3_hook.load_file('/tmp/users_data.csv', key='users_data.csv', bucket_name='my-bucket', replace=True)
15

16
default_args = {
17
    'owner': 'data_engineer',
18
    'start_date': datetime.datetime(2023, 9, 1),
19
}
20

21
with DAG(
22
    dag_id='extract_postgres_to_s3',
23
    default_args=default_args,
24
    schedule_interval='@daily',
25
    catchup=False,
26
) as dag:
27

28
    extract_task = PythonOperator(
29
        task_id='extract_postgres',
30
        python_callable=extract_data_from_postgres
31
    )
32

33
    upload_task = PythonOperator(
34
        task_id='upload_to_s3',
35
        python_callable=upload_to_s3
36
    )
37

38
    extract_task >> upload_task

Here, we have two Python functions: one to fetch data into a CSV, and another to upload that CSV to S3. Airflow’s PostgresHook sums up hours of boilerplate you’d otherwise need for connecting to databases.

Example Transformation#

Let’s illustrate a simple Python transformation for normalizing some data columns:

1
import datetime
2
from airflow import DAG
3
from airflow.operators.python_operator import PythonOperator
4

5
def transform_data():
6
    import pandas as pd
7

8
    df = pd.read_csv('/tmp/users_data.csv')
9
    # Normalize columns
10
    df['email'] = df['email'].str.lower()
11
    # Convert sign-up date from string to datetime
12
    df['sign_up_date'] = pd.to_datetime(df['sign_up_date'], format='%Y-%m-%d')
13
    # Save transformed data
14
    df.to_csv('/tmp/users_data_transformed.csv', index=False)
15

16
default_args = {
17
    'owner': 'data_engineer',
18
    'start_date': datetime.datetime(2023, 9, 1),
19
}
20

21
with DAG(
22
    dag_id='transform_data_dag',
23
    default_args=default_args,
24
    schedule_interval=None,
25
    catchup=False,
26
) as dag:
27

28
    transform_task = PythonOperator(
29
        task_id='transform',
30
        python_callable=transform_data
31
    )

• This DAG has a single task that reads a CSV, applies transformations, and writes the result to a new CSV.
• In more advanced contexts, you could read from a database or an object store, process large batches with a distributed framework, or chain multiple transformations together in a single DAG.

Common Targets for Load#

• Data Warehouse: Redshift, BigQuery, Snowflake.
• Analytical DB: Postgres, MySQL.
• NoSQL Stores: MongoDB, Cassandra.
• Data Lakes: S3 (Parquet or ORC files), HDFS.

Loading data often entails using specialized operators or writing custom scripts to push data to the destination. Below is an example using the BigQuery operator to load data from GCS:

1
import datetime
2
from airflow import DAG
3
from airflow.providers.google.cloud.operators.bigquery import BigQueryCreateExternalTableOperator
4

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

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

17
    create_external_table = BigQueryCreateExternalTableOperator(
18
        task_id='create_external_table',
19
        bucket='my-data',
20
        source_objects=['users_data_transformed.csv'],
21
        destination_project_dataset_table='my_project.my_dataset.users_data',
22
        source_format='CSV',
23
        skip_leading_rows=1,
24
        field_delimiter=',',
25
        schema_fields=[
26
            {'name': 'id', 'type': 'INTEGER', 'mode': 'REQUIRED'},
27
            {'name': 'name', 'type': 'STRING', 'mode': 'NULLABLE'},
28
            {'name': 'email', 'type': 'STRING', 'mode': 'NULLABLE'},
29
            {'name': 'sign_up_date', 'type': 'TIMESTAMP', 'mode': 'NULLABLE'}
30
        ],
31
        google_cloud_storage_conn_id='my_gcp_conn',
32
        bigquery_conn_id='my_bigquery_conn'
33
    )

This operator creates an external table in BigQuery reading from your file in GCS. The same concept applies to other warehouses like Snowflake or Redshift, where corresponding operators exist (e.g., SnowflakeOperator, RedshiftSQLOperator).

XComs (Cross-Communication)#

XCom is Airflow’s built-in mechanism for sharing data between tasks. Rather than persisting data in files, you can push and pull small amounts of data (like IDs, metadata, or simple objects) across tasks.

1
def push_data(**context):
2
    context['ti'].xcom_push(key='my_data', value='hello_world')
3

4
def pull_data(**context):
5
    result = context['ti'].xcom_pull(key='my_data', task_ids='push_task')
6
    print("Received data:", result)

Branching#

Branching allows you to conditionally choose a path in your DAG. For instance, you can decide which tasks to run based on whether your data meets some criteria.

1
from airflow.operators.python import BranchPythonOperator
2

3
def choose_branch(**context):
4
    # Pseudo condition
5
    if context['execution_date'].weekday() < 5:
6
        return 'week_day_task'
7
    else:
8
        return 'weekend_task'
9

10
branch_op = BranchPythonOperator(
11
    task_id='branch_op',
12
    python_callable=choose_branch
13
)

SubDAGs#

SubDAGs are a way to nest a DAG inside another DAG for modularity. However, SubDAGs are largely replaced by Task Groups in modern Airflow versions. Task Groups provide a simpler mechanism for grouping tasks visually and logically without introducing a brand new DAG.

Task Groups#

1
from airflow.utils.task_group import TaskGroup
2
with DAG(...) as dag:
3
    with TaskGroup("extraction_group") as extraction_group:
4
        # define tasks for extraction
5
        ...
6
    with TaskGroup("transformation_group") as transformation_group:
7
        # define tasks for transformation
8
        ...

Sensors#

Sensors are simply specialized operators that wait for a certain event. For example, an S3KeySensor waits for a file to land in an S3 bucket before allowing downstream tasks to proceed.

Airflow UI#

• DAGs View: A list of all the DAGs in your environment.
• Graph View: A graphical representation of tasks and their dependencies.
• Tree View: Historical runs of your DAG over time.
• Task Instance Logs: Logs from each task run.

Email and Slack Notifications#

You can configure email or Slack alerts for task failures or retries. Update your airflow.cfg or use environment variables to set email or Slack hooks. For Slack, you’d typically use the SlackWebhookOperator or custom on-failure callbacks:

1
from airflow.providers.slack.operators.slack_webhook import SlackWebhookOperator
2

3
def slack_fail_alert(context):
4
    slack_msg = f"DAG {context['dag'].dag_id} failed on task {context['task_instance'].task_id}"
5
    alert = SlackWebhookOperator(
6
        task_id='slack_fail',
7
        http_conn_id='slack_conn',
8
        webhook_token='YOUR_WEBHOOK_TOKEN',
9
        message=slack_msg,
10
        username='airflow'
11
    )
12
    return alert.execute(context=context)

And then in your task definition:

1
PythonOperator(
2
    task_id='some_task',
3
    python_callable=my_callable,
4
    on_failure_callback=slack_fail_alert
5
)

Executors#

Airflow comes with different executors to handle how tasks are run:

• SequentialExecutor: Runs one task at a time (default for local testing).
• LocalExecutor: Runs parallel tasks on a single machine.
• CeleryExecutor: Distributes tasks across multiple worker nodes.
• KubernetesExecutor: Schedules each task in a separate Kubernetes pod.

Database Optimizations#

• Ensure you use a production-grade database for the Airflow metadata (e.g., Postgres or MySQL).
• Avoid SQLite in production scenarios as it can cause concurrency issues.
• Clean up old logs and XCom entries regularly to prevent the metadata DB from bloating.

Scaling with Celery or Kubernetes#

• CeleryExecutor: Requires a message broker like RabbitMQ or Redis. You set up multiple worker nodes to process tasks in parallel.
• KubernetesExecutor: Each Airflow task runs in its own Kubernetes pod. This is highly dynamic and can scale well if you already have a Kubernetes cluster.

Parallelism and Concurrency#

• parallelism: The max number of task instances that can run across all DAGs.
• dag_concurrency: Max number of task instances for a single DAG at once.
• max_active_runs_per_dag: Limits how many DAG runs can be active simultaneously.

Tuning these parameters in airflow.cfg helps manage resource usage.