Apache Airflow at Scale: DAG Patterns, Best Practices, and Monitoring

from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.operators.bash import BashOperator
from airflow.sensors.external_task import ExternalTaskSensor
from datetime import datetime, timedelta

# ✅ Good: Reusable, testable DAG
default_args = {
    'owner': 'data-team',
    'depends_on_past': False,
    'email': ['alerts@example.com'],
    'email_on_failure': True,
    'email_on_retry': False,
    'retries': 3,
    'retry_delay': timedelta(minutes=5),
    'retry_exponential_backoff': True,
    'max_retry_delay': timedelta(hours=1)
}

with DAG(
    'etl_user_analytics',
    default_args=default_args,
    description='Daily user analytics ETL',
    schedule='0 2 * * *',  # 2 AM daily
    start_date=datetime(2024, 1, 1),
    catchup=False,  # Don't backfill on deploy
    max_active_runs=1,  # Prevent overlapping runs
    tags=['analytics', 'daily']
) as dag:
    
    # 1. Extract phase
    extract_users = PythonOperator(
        task_id='extract_users',
        python_callable=extract_users_from_db,
        op_kwargs={'date': '{{ ds }}'}  # Templated date
    )
    
    extract_events = PythonOperator(
        task_id='extract_events',
        python_callable=extract_events_from_s3,
        op_kwargs={'date': '{{ ds }}'}
    )
    
    # 2. Transform phase (runs after extracts)
    transform_data = PythonOperator(
        task_id='transform_data',
        python_callable=transform_user_events,
        op_kwargs={
            'users_path': '{{ ti.xcom_pull(task_ids="extract_users") }}',
            'events_path': '{{ ti.xcom_pull(task_ids="extract_events") }}'
        }
    )
    
    # 3. Load phase
    load_to_warehouse = PythonOperator(
        task_id='load_to_warehouse',
        python_callable=load_to_redshift,
        op_kwargs={'data_path': '{{ ti.xcom_pull(task_ids="transform_data") }}'}
    )
    
    # 4. Data quality check
    validate_data = PythonOperator(
        task_id='validate_data',
        python_callable=run_data_quality_checks
    )
    
    # Define dependencies
    [extract_users, extract_events] >> transform_data >> load_to_warehouse >> validate_data

from airflow.exceptions import AirflowException, AirflowSkipException
from airflow.providers.postgres.hooks.postgres import PostgresHook
from airflow.utils.trigger_rule import TriggerRule

def extract_with_error_handling(**context):
    """Extract with proper error handling"""
    try:
        # Attempt extraction
        data = fetch_data_from_api()
        
        # Data quality check
        if not data or len(data) == 0:
            raise AirflowException("No data returned from API")
        
        # Save to XCom
        return save_to_s3(data)
        
    except requests.Timeout:
        # Retriable error - will retry based on DAG retries
        raise
    except requests.HTTPError as e:
        if e.response.status_code == 404:
            # Skip this task - data doesn't exist yet
            raise AirflowSkipException(f"Data not found: {e}")
        elif e.response.status_code >= 500:
            # Server error - retry
            raise
        else:
            # Client error - don't retry
            raise AirflowException(f"API error: {e}")

# Cleanup task that runs even if pipeline fails
cleanup_task = PythonOperator(
    task_id='cleanup_temp_files',
    python_callable=cleanup_temp_files,
    trigger_rule=TriggerRule.ALL_DONE  # Run regardless of upstream status
)

# Notification task for failures
def send_failure_notification(context):
    """Send Slack notification on failure"""
    dag_id = context['dag'].dag_id
    task_id = context['task'].task_id
    execution_date = context['execution_date']
    log_url = context['task_instance'].log_url
    
    send_slack_message(
        f"❌ Task {task_id} in DAG {dag_id} failed\n"
        f"Execution: {execution_date}\n"
        f"Logs: {log_url}"
    )

failure_alert = PythonOperator(
    task_id='failure_alert',
    python_callable=send_failure_notification,
    trigger_rule=TriggerRule.ONE_FAILED
)

# Task dependency
[extract_users, extract_events] >> transform_data >> [load_to_warehouse, cleanup_task]
load_to_warehouse >> [validate_data, failure_alert]

from airflow.decorators import task, task_group
from airflow.models import Variable

# ✅ Dynamic DAGs based on configuration
TABLES = Variable.get("tables_to_sync", deserialize_json=True)
# ["users", "orders", "products", "reviews"]

with DAG('dynamic_table_sync', ...) as dag:
    
    @task_group(group_id="extract_tables")
    def extract_all_tables():
        for table in TABLES:
            @task(task_id=f"extract_{table}")
            def extract_table(table_name: str):
                return extract_from_postgres(table_name)
            
            extract_table(table)
    
    @task_group(group_id="transform_tables")
    def transform_all_tables():
        for table in TABLES:
            @task(task_id=f"transform_{table}")
            def transform_table(table_name: str):
                return transform_data(table_name)
            
            transform_table(table)
    
    extract_all_tables() >> transform_all_tables()

# Using TaskFlow API (cleaner syntax)
from airflow.decorators import dag

@dag(
    schedule='@daily',
    start_date=datetime(2024, 1, 1),
    catchup=False,
    tags=['taskflow']
)
def modern_etl():
    
    @task
    def extract():
        """Extract from source"""
        return fetch_data()
    
    @task
    def transform(data: list):
        """Transform data"""
        return process_data(data)
    
    @task
    def load(data: list):
        """Load to warehouse"""
        save_to_warehouse(data)
    
    # Automatic dependency management
    load(transform(extract()))

dag_instance = modern_etl()

# 1. Data quality checks
from great_expectations.core import ExpectationSuite
from airflow.providers.great_expectations.operators.great_expectations import \
    GreatExpectationsOperator

validate_data = GreatExpectationsOperator(
    task_id="validate_user_data",
    expectation_suite_name="user_data_suite",
    data_context_root_dir="/path/to/great_expectations",
    fail_task_on_validation_failure=True
)

# 2. SLA monitoring
sla_miss_callback = lambda context: send_sla_alert(context)

with DAG(
    'time_sensitive_pipeline',
    default_args={
        'sla': timedelta(hours=2),  # Must complete within 2 hours
        'sla_miss_callback': sla_miss_callback
    },
    ...
) as dag:
    pass

# 3. Custom metrics
from airflow.providers.statsd.hooks.statsd import StatsHook

def track_metrics(**context):
    statsd = StatsHook()
    
    # Record execution time
    duration = context['task_instance'].duration
    statsd.timing('airflow.task.duration', duration)
    
    # Record row counts
    row_count = get_row_count()
    statsd.gauge('airflow.pipeline.rows_processed', row_count)

# 4. Logs to external systems
import logging
from pythonjsonlogger import jsonlogger

logger = logging.getLogger(__name__)
handler = logging.StreamHandler()
formatter = jsonlogger.JsonFormatter()
handler.setFormatter(formatter)
logger.addHandler(handler)

def log_pipeline_metrics(**context):
    logger.info('Pipeline completed', extra={
        'dag_id': context['dag'].dag_id,
        'execution_date': str(context['execution_date']),
        'rows_processed': get_row_count(),
        'duration_seconds': context['task_instance'].duration
    })