Workflow Orchestration

← Previous: ETL/ELT Pipelines | Back to Index | Next: Data Quality & Testing →

Introduction

Apache Airflow has become my primary orchestration tool. It schedules, monitors, and manages complex data pipelines. This article covers production Airflow patterns I use daily.

Airflow Fundamentals

DAG (Directed Acyclic Graph)

# Python 3.12 - Basic DAG
from airflow import DAG
from airflow.operators.python import PythonOperator
from datetime import datetime, timedelta

default_args = {
    'owner': 'data-eng',
    'depends_on_past': False,
    'email': ['[email protected]'],
    'email_on_failure': True,
    'email_on_retry': False,
    'retries': 3,
    'retry_delay': timedelta(minutes=5),
}

with DAG(
    'user_metrics_pipeline',
    default_args=default_args,
    description='Calculate daily user metrics',
    schedule='0 2 * * *',  # Run at 2 AM daily
    start_date=datetime(2026, 1, 1),
    catchup=False,
    tags=['metrics', 'users'],
) as dag:
    
    def extract_data(**context):
        """Extract user data."""
        execution_date = context['ds']
        logging.info(f"Extracting data for {execution_date}")
        # Extraction logic
        return {'records': 1000}
    
    def transform_data(**context):
        """Transform extracted data."""
        ti = context['ti']
        extract_result = ti.xcom_pull(task_ids='extract')
        logging.info(f"Transforming {extract_result['records']} records")
        # Transformation logic
        return {'records': 950}
    
    def load_data(**context):
        """Load to warehouse."""
        ti = context['ti']
        transform_result = ti.xcom_pull(task_ids='transform')
        logging.info(f"Loading {transform_result['records']} records")
        # Load logic
    
    # Define tasks
    extract = PythonOperator(
        task_id='extract',
        python_callable=extract_data,
    )
    
    transform = PythonOperator(
        task_id='transform',
        python_callable=transform_data,
    )
    
    load = PythonOperator(
        task_id='load',
        python_callable=load_data,
    )
    
    # Define dependencies
    extract >> transform >> load

Production DAG Patterns

Dynamic Task Generation

# Python 3.12 - Dynamic tasks
from airflow.operators.python import PythonOperator

def process_table(table_name: str, **context):
    """Process a single table."""
    logging.info(f"Processing {table_name}")
    # Processing logic

with DAG('dynamic_pipeline', ...) as dag:
    
    # Generate tasks dynamically
    tables = ['users', 'transactions', 'products', 'orders']
    
    start = PythonOperator(task_id='start', python_callable=lambda: None)
    end = PythonOperator(task_id='end', python_callable=lambda: None)
    
    for table in tables:
        task = PythonOperator(
            task_id=f'process_{table}',
            python_callable=process_table,
            op_kwargs={'table_name': table},
        )
        
        start >> task >> end

Task Groups

# Python 3.12 - Task groups for organization
from airflow.utils.task_group import TaskGroup

with DAG('etl_with_groups', ...) as dag:
    
    with TaskGroup('extraction') as extract_group:
        extract_users = PythonOperator(task_id='extract_users', ...)
        extract_orders = PythonOperator(task_id='extract_orders', ...)
    
    with TaskGroup('transformation') as transform_group:
        transform_users = PythonOperator(task_id='transform_users', ...)
        transform_orders = PythonOperator(task_id='transform_orders', ...)
    
    with TaskGroup('loading') as load_group:
        load_users = PythonOperator(task_id='load_users', ...)
        load_orders = PythonOperator(task_id='load_orders', ...)
    
    extract_group >> transform_group >> load_group

Sensors for External Dependencies

# Python 3.12 - Wait for external events
from airflow.sensors.filesystem import FileSensor
from airflow.sensors.python import PythonSensor
from airflow.providers.postgres.sensors.postgres import PostgresSqlSensor

with DAG('sensor_example', ...) as dag:
    
    # Wait for file to appear
    wait_for_file = FileSensor(
        task_id='wait_for_file',
        filepath='/data/input/file.csv',
        poke_interval=60,  # Check every 60 seconds
        timeout=3600,  # Timeout after 1 hour
    )
    
    # Wait for database record
    wait_for_data = PostgresSqlSensor(
        task_id='wait_for_data',
        conn_id='postgres_default',
        sql="SELECT COUNT(*) FROM uploads WHERE status='complete' AND date='{{ ds }}'",
        poke_interval=300,
    )
    
    # Custom sensor
    def check_api_ready():
        import requests
        response = requests.get('https://api.example.com/health')
        return response.status_code == 200
    
    wait_for_api = PythonSensor(
        task_id='wait_for_api',
        python_callable=check_api_ready,
        poke_interval=60,
        timeout=1800,
    )
    
    [wait_for_file, wait_for_data, wait_for_api] >> process_data

XCom for Task Communication

# Python 3.12 - Share data between tasks
def extract_data(**context):
    """Extract and return data via XCom."""
    data = {'user_count': 1000, 'total_amount': 50000}
    return data  # Automatically pushed to XCom

def validate_data(**context):
    """Pull data from XCom and validate."""
    ti = context['ti']
    data = ti.xcom_pull(task_ids='extract_data')
    
    if data['user_count'] < 100:
        raise ValueError("Insufficient data")
    
    return {'status': 'valid'}

def load_data(**context):
    """Pull multiple XComs."""
    ti = context['ti']
    extract_result = ti.xcom_pull(task_ids='extract_data')
    validate_result = ti.xcom_pull(task_ids='validate_data')
    
    if validate_result['status'] == 'valid':
        # Load data
        logging.info(f"Loading {extract_result['user_count']} users")

Branching Logic

# Python 3.12 - Conditional execution
from airflow.operators.python import BranchPythonOperator

def determine_branch(**context):
    """Decide which path to take."""
    execution_date = context['ds']
    day_of_week = datetime.strptime(execution_date, '%Y-%m-%d').weekday()
    
    if day_of_week == 6:  # Sunday
        return 'full_refresh'
    else:
        return 'incremental_load'

with DAG('branching_pipeline', ...) as dag:
    
    branch = BranchPythonOperator(
        task_id='branch',
        python_callable=determine_branch,
    )
    
    full_refresh = PythonOperator(
        task_id='full_refresh',
        python_callable=lambda: logging.info("Running full refresh"),
    )
    
    incremental_load = PythonOperator(
        task_id='incremental_load',
        python_callable=lambda: logging.info("Running incremental load"),
    )
    
    complete = PythonOperator(
        task_id='complete',
        python_callable=lambda: logging.info("Pipeline complete"),
        trigger_rule='none_failed_min_one_success',
    )
    
    branch >> [full_refresh, incremental_load] >> complete

Error Handling and Retries

# Python 3.12 - Custom retry logic
from airflow.exceptions import AirflowException

def task_with_custom_retry(**context):
    """Task with smart retry logic."""
    try:
        # Attempt operation
        result = risky_operation()
        return result
    
    except TemporaryError as e:
        # Retry for temporary errors
        raise AirflowException(f"Temporary error, will retry: {e}")
    
    except PermanentError as e:
        # Don't retry for permanent errors
        logging.error(f"Permanent error, skipping retries: {e}")
        # Mark task as failed without retry
        raise AirflowException(f"Permanent failure: {e}")

task = PythonOperator(
    task_id='custom_retry',
    python_callable=task_with_custom_retry,
    retries=5,
    retry_delay=timedelta(minutes=10),
    retry_exponential_backoff=True,
    max_retry_delay=timedelta(hours=1),
)

Database Operators

# Python 3.12 - Database operations
from airflow.providers.postgres.operators.postgres import PostgresOperator
from airflow.providers.postgres.hooks.postgres import PostgresHook

with DAG('database_pipeline', ...) as dag:
    
    # SQL operator
    create_table = PostgresOperator(
        task_id='create_table',
        postgres_conn_id='postgres_warehouse',
        sql="""
            CREATE TABLE IF NOT EXISTS user_metrics (
                user_id VARCHAR(50),
                metric_date DATE,
                total_amount DECIMAL(10,2),
                transaction_count INTEGER,
                PRIMARY KEY (user_id, metric_date)
            )
        """,
    )
    
    # Custom SQL with parameters
    def run_custom_query(**context):
        """Run parameterized query."""
        hook = PostgresHook(postgres_conn_id='postgres_warehouse')
        
        query = """
            INSERT INTO user_metrics
            SELECT user_id, DATE('{{ ds }}'), SUM(amount), COUNT(*)
            FROM transactions
            WHERE DATE(created_at) = DATE('{{ ds }}')
            GROUP BY user_id
        """
        
        hook.run(query)
    
    load_metrics = PythonOperator(
        task_id='load_metrics',
        python_callable=run_custom_query,
    )
    
    create_table >> load_metrics

Monitoring and Alerting

# Python 3.12 - Custom alerts
from airflow.operators.email import EmailOperator

def check_data_quality(**context):
    """Check data quality and alert if issues."""
    ti = context['ti']
    row_count = ti.xcom_pull(task_ids='load_data')['row_count']
    
    if row_count < 1000:
        # Send alert
        send_slack_alert(f"Warning: Only {row_count} rows loaded")
        raise AirflowException("Data quality check failed")
    
    return row_count

quality_check = PythonOperator(
    task_id='quality_check',
    python_callable=check_data_quality,
)

# Email on failure
alert_email = EmailOperator(
    task_id='alert_email',
    to='[email protected]',
    subject='Pipeline {{ dag.dag_id }} failed on {{ ds }}',
    html_content='Check Airflow UI for details',
    trigger_rule='one_failed',
)

Best Practices

1. Idempotent Tasks

def idempotent_load(**context):
    """
    Delete and insert pattern.
    Safe to re-run without duplicates.
    """
    execution_date = context['ds']
    
    # Delete existing data for this date
    delete_query = f"DELETE FROM metrics WHERE date = '{execution_date}'"
    hook.run(delete_query)
    
    # Insert new data
    insert_query = f"INSERT INTO metrics SELECT * FROM staging WHERE date = '{execution_date}'"
    hook.run(insert_query)

2. Resource Management

with DAG('resource_aware', ...) as dag:
    
    # Limit concurrent tasks
    heavy_task = PythonOperator(
        task_id='heavy_processing',
        python_callable=process_large_dataset,
        pool='heavy_tasks',  # Use pool to limit concurrency
        priority_weight=10,  # Higher priority
    )

3. Testing DAGs

# Python 3.12 - Test DAG
import pytest
from airflow.models import DagBag

def test_dag_loaded():
    """Test DAG loads without errors."""
    dagbag = DagBag(dag_folder='dags/', include_examples=False)
    assert len(dagbag.import_errors) == 0
    assert 'user_metrics_pipeline' in dagbag.dags

def test_dag_structure():
    """Test DAG structure."""
    dagbag = DagBag(dag_folder='dags/')
    dag = dagbag.get_dag('user_metrics_pipeline')
    
    assert len(dag.tasks) == 3
    assert dag.has_task('extract')
    assert dag.has_task('transform')
    assert dag.has_task('load')

Conclusion

Airflow orchestrates complex data workflows reliably. Use DAGs for scheduling, sensors for dependencies, and XCom for task communication.

Key takeaways:

• Design DAGs to be idempotent

• Use sensors for external dependencies

• Implement proper error handling and retries

• Monitor pipeline health

• Test DAGs before production

---

Navigation:

• ← Previous: ETL/ELT Pipelines

• Next: Data Quality & Testing →

• Back to Overview