Data Pipeline Orchestration: Airflow vs Prefect vs Dagster

# airflow_dag.py
from airflow import DAG
from airflow.operators.python import PythonOperator, BranchPythonOperator
from airflow.operators.bash import BashOperator
from airflow.providers.postgres.operators.postgres import PostgresOperator
from airflow.sensors.external_task import ExternalTaskSensor
from datetime import datetime, timedelta

default_args = {
    'owner': 'data_team',
    'depends_on_past': False,
    'retries': 3,
    'retry_delay': timedelta(minutes=5),
}

with DAG(
    'etl_pipeline',
    default_args=default_args,
    description='Daily ETL pipeline for analytics',
    schedule_interval='0 2 * * *',  # 2 AM daily
    start_date=datetime(2024, 1, 1),
    catchup=False,
    tags=['etl', 'analytics'],
) as dag:

    def extract_data(**context):
        """Extract data from source."""
        import requests
        
        # Pull data from API
        response = requests.get(
            'https://api.example.com/data',
            params={'date': context['ds']}
        )
        
        # Store in XCom for downstream tasks
        context['task_instance'].xcom_push(
            key='raw_data',
            value=response.json()
        )
        
        return 'Data extracted successfully'

    def transform_data(**context):
        """Transform extracted data."""
        import json
        
        # Pull from XCom
        ti = context['task_instance']
        raw_data = ti.xcom_pull(key='raw_data', task_ids='extract_data')
        
        # Transform
        transformed = [item for item in raw_data if item['status'] == 'active']
        
        # Push transformed data
        ti.xcom_push(key='transformed_data', value=transformed)
        
        return f'Transformed {len(transformed)} records'

    def load_to_warehouse(**context):
        """Load data to data warehouse."""
        # Pull transformed data
        ti = context['task_instance']
        data = ti.xcom_pull(key='transformed_data', task_ids='transform_data')
        
        # Load to warehouse (pseudocode)
        # warehouse.insert('table', data)
        
        print(f'Loaded {len(data)} records to warehouse')

    def check_data_quality(**context):
        """Check data quality rules."""
        ti = context['task_instance']
        data = ti.xcom_pull(key='transformed_data', task_ids='transform_data')
        
        if len(data) < 100:
            raise ValueError(f'Data quality check failed: only {len(data)} records')
        
        return 'Data quality checks passed'

    # Task definitions
    extract = PythonOperator(
        task_id='extract_data',
        python_callable=extract_data,
    )

    transform = PythonOperator(
        task_id='transform_data',
        python_callable=transform_data,
    )

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

    load = PythonOperator(
        task_id='load_to_warehouse',
        python_callable=load_to_warehouse,
    )

    # Dependencies
    extract >> transform >> check >> load

from airflow.utils.task_group import TaskGroup

with DAG('complex_etl', ...) as dag:
    
    with TaskGroup('processing') as processing_group:
        task_a = PythonOperator(task_id='task_a', ...)
        task_b = PythonOperator(task_id='task_b', ...)
        task_a >> task_b
    
    with TaskGroup('loading') as loading_group:
        task_c = PythonOperator(task_id='task_c', ...)
        task_d = PythonOperator(task_id='task_d', ...)
        task_c >> task_d
    
    processing_group >> loading_group

from airflow.sensors.base import BaseSensorOperator

class MyCustomSensor(BaseSensorOperator):
    def poke(self, context):
        import time
        # Check if data is ready
        return time.time() % 10 == 0  # Simplified

# Using sensors
wait_for_data = MyCustomSensor(
    task_id='wait_for_data',
    poke_interval=60,
    timeout=3600,
)

wait_for_external = ExternalTaskSensor(
    task_id='wait_for_upstream',
    external_dag_id='upstream_dag',
    external_task_id='finish',
    timeout=7200,
)

# prefect_flow.py
from prefect import flow, task, get_run_logger
from typing import List
import httpx

@task(name='Extract Data', retries=3, retry_delay_seconds=60)
def extract_data(date: str) -> dict:
    """Extract data from API."""
    logger = get_run_logger()
    
    response = httpx.get(
        'https://api.example.com/data',
        params={'date': date}
    )
    response.raise_for_status()
    
    logger.info(f'Extracted {len(response.json())} records')
    return response.json()

@task(name='Transform Data', cache_key_fn=lambda *args, **kwargs: kwargs.get('date'))
def transform_data(raw_data: List[dict], date: str) -> List[dict]:
    """Transform extracted data."""
    logger = get_run_logger()
    
    transformed = [
        item for item in raw_data 
        if item.get('status') == 'active'
    ]
    
    logger.info(f'Transformed {len(transformed)} records')
    return transformed

@task(name='Load to Warehouse', result_storage='s3://bucket/results/')
def load_to_warehouse(data: List[dict], table: str) -> int:
    """Load data to warehouse."""
    logger = get_run_logger()
    
    # Load to warehouse
    # warehouse.insert(table, data)
    
    logger.info(f'Loaded {len(data)} records to {table}')
    return len(data)

@task(name='Check Data Quality')
def check_data_quality(data: List[dict], min_records: int = 100) -> bool:
    """Validate data meets quality requirements."""
    if len(data) < min_records:
        raise ValueError(f'Only {len(data)} records, need at least {min_records}')
    return True

@flow(
    name='ETL Pipeline',
    description='Daily data pipeline',
    schedule={'cron': '0 2 * * *'},
    log_prints=True,
)
def etl_pipeline(date: str):
    """Main ETL flow."""
    
    # Extract
    raw_data = extract_data(date)
    
    # Transform
    transformed = transform_data(raw_data, date=date)
    
    # Validate
    check_data_quality(transformed)
    
    # Load
    records_loaded = load_to_warehouse(transformed, table='analytics.data')
    
    return {'records_extracted': len(raw_data), 'records_loaded': records_loaded}


# Run locally
if __name__ == '__main__':
    result = etl_pipeline('2024-01-15')
    print(result)

# prefect.yaml (deployment configuration)
name: etl-pipeline
version: 1.0.0

schedule:
  cron: 0 2 * * *

parameters:
  date: "{{ today }}"

infrastructure:
  type: kubernetes-job
  job:
    image: my-registry/etl-pipeline:latest
    env:
      - DATABASE_URL: "postgresql://..."

storage:
  type: s3
  bucket: my-bucket
  path: /prefect-flows/

actions:
  - type: run
    triggers:
      - type: completion
        match:
          state: COMPLETED

# Using blocks for configuration
from prefect import flow
from prefect.blocks.system import Secret

@flow
def sensitive_flow():
    # Access secrets securely
    api_key = Secret.load('api-key').get()
    # Use api_key...

# Using Sub-Flows
@flow
def parent_flow():
    result = child_flow()  # Wait for sub-flow
    
@flow
def child_flow():
    return "child result"

# Parallel execution with task mapping
@task
def process_item(item):
    return item * 2

@flow
def parallel_flow(items: List[int]):
    # Map task across items (runs in parallel)
    results = process_item.map(items)
    return results

# dagster_pipeline.py
from dagster import (
    asset, AssetExecutionContext, AssetKey,
    job, op, In, Out, get_dagster_logger
)
import pandas as pd

logger = get_dagster_logger()

# Define assets (software-defined)
@asset(
    key='raw_data',
    description='Raw data from API',
    group_name='etl',
    metadata={'source': 'api'}
)
def extract_data(context: AssetExecutionContext) -> pd.DataFrame:
    """Extract raw data from API."""
    import httpx
    
    context.log_info('Extracting data from API...')
    
    response = httpx.get('https://api.example.com/data')
    data = response.json()
    
    df = pd.DataFrame(data)
    context.log_info(f'Extracted {len(df)} records')
    
    return df

@asset(
    key='transformed_data',
    description='Transformed data',
    group_name='etl',
    deps=['raw_data']
)
def transform_data(context: AssetExecutionContext, raw_data: pd.DataFrame) -> pd.DataFrame:
    """Transform the raw data."""
    context.log_info(f'Transforming {len(raw_data)} records...')
    
    # Transformation logic
    transformed = raw_data[raw_data['status'] == 'active'].copy()
    transformed['processed_at'] = pd.Timestamp.now()
    
    context.log_info(f'Created {len(transformed)} transformed records')
    return transformed

@asset(
    key='analytics_table',
    description='Final analytics table',
    group_name='etl',
    deps=['transformed_data']
)
def load_to_warehouse(
    context: AssetExecutionContext, 
    transformed_data: pd.DataFrame
) -> None:
    """Load data to warehouse."""
    context.log_info(f'Loading {len(transformed_data)} records...')
    
    # Load to warehouse
    # warehouse.insert('analytics', transformed_data)
    
    context.log_info('Data loaded successfully')

# Alternative: Using ops and jobs for more control
@op
def my_op(context, input_data):
    """Single operation."""
    context.log.info(f'Processing {len(input_data)} items')
    return input_data

@job
def my_job():
    """Job that combines ops."""
    data = my_op([1, 2, 3])
    my_op(data)

from dagster import schedule, sensor, RunRequest

# Time-based schedule
@schedule(
    cron_schedule='0 2 * * *',
    job_name='etl_pipeline',
    execution_timezone='UTC'
)
def daily_schedule(context):
    """Daily ETL schedule."""
    return RunRequest(
        run_key=None,
        tags={'date': context.scheduled_execution_time.strftime('%Y-%m-%d')}
    )

# Event-based sensor
@sensor(job_name='etl_pipeline')
def api_sensor():
    """Poll API for new data."""
    import httpx
    
    response = httpx.get('https://api.example.com/last-update')
    last_update = response.json().get('timestamp')
    
    # Check if new data available
    if last_update != get_stored_timestamp():
        return RunRequest(
            run_key=f'sensor_{last_update}',
            tags={'last_update': last_update}
        )
    
    return None

# Asset sensors
from dagster import AssetMaterialization

@sensor(job_name='etl_pipeline')
def upstream_sensor():
    """Trigger when upstream asset is materialized."""
    # Implementation
    pass

import pytest
from dagster import build_op_context, asset_definitions_test

# Test individual ops
def test_my_op():
    context = build_op_context()
    result = my_op(context, [1, 2, 3])
    assert len(result) == 3

# Test assets
def test_extract_data_asset():
    result = extract_data(build_op_context())
    assert isinstance(result, pd.DataFrame)

# Test full job
def test_my_job():
    from dagster import execute_job
    
    result = execute_job(
        job=my_job,
        run_config={'ops': {'my_op': {'inputs': {'input_data': [1, 2, 3]}}}}
    )
    
    assert result.success

# Asset lineage testing
def test_asset_dependencies():
    # Verify asset dependencies are correct
    assets = get_assets_defs()
    
    raw_data = next(a for a in assets if a.key.path == ['raw_data'])
    transformed = next(a for a in assets if a.key.path == ['transformed_data'])
    
    assert transformed.key in raw_data.required_downsstream_assets

# Using Prefect with Kubernetes
from prefect import flow
from prefect.infrastructure.kubernetes import KubernetesJob

@flow(
    name='Production ETL',
    infrastructure=KubernetesJob(
        image='my-registry/etl:v1',
        namespace='production',
        service_account_name='etl-runner',
        env={'DATABASE_URL': {'env': 'DATABASE_URL'}},
    )
)
def production_flow():
    extract()
    transform()
    load()

# Local development with same code
@flow(name='Local ETL')
def local_flow():
    extract()
    transform()
    load()

# Airflow code
with DAG('my_dag', ...) as dag:
    task1 = PythonOperator(task_id='task1', python_callable=func1)
    task2 = PythonOperator(task_id='task2', python_callable=func2)
    task1 >> task2

# Equivalent Prefect code
@flow
def my_flow():
    result1 = func1()
    func2(result1)

# Or with explicit task definitions
@task
def task1():
    return func1()

@task
def task2(input):
    return func2(input)

@flow
def my_flow():
    result = task1()
    task2(result)