Data Pipeline Orchestration — Airflow, Prefect, and Dagster Guide

# quality_etl_dag.py
# Airflow DAG with data quality checks
from datetime import datetime, timedelta
from airflow import DAG
from airflow.operators.python import PythonOperator, BranchPythonOperator
from airflow.operators.bash import BashOperator
from airflow.operators.dummy import DummyOperator

default_args = {
    'owner': 'dodatech',
    'depends_on_past': False,
    'start_date': datetime(2026, 1, 1),
    'retries': 2,
    'retry_delay': timedelta(minutes=5),
}

dag = DAG(
    'quality_etl_pipeline',
    default_args=default_args,
    description='ETL with data quality checks and branching',
    schedule_interval='0 3 * * *',
    catchup=False,
    tags=['quality', 'etl'],
)

def extract_orders(**context):
    import json
    # Simulated extraction
    data = {
        "orders": [
            {"id": 1, "amount": 100.0, "date": "2026-06-01"},
            {"id": 2, "amount": 250.0, "date": "2026-06-01"},
            {"id": 3, "amount": -50.0, "date": "2026-06-02"},  # invalid
        ]
    }
    context['ti'].xcom_push(key='orders', value=data)
    print(f"Extracted {len(data['orders'])} orders")

def check_data_quality(**context):
    ti = context['ti']
    data = ti.xcom_pull(key='orders')
    errors = []
    
    for order in data['orders']:
        if order['amount'] <= 0:
            errors.append(f"Order {order['id']}: negative amount {order['amount']}")
    
    if errors:
        ti.xcom_push(key='quality_errors', value=errors)
        ti.xcom_push(key='quality_status', value='failed')
        print(f"Quality check FAILED: {errors}")
        return 'alert_team'
    else:
        ti.xcom_push(key='quality_status', value='passed')
        print("Quality check PASSED")
        return 'transform_orders'

def alert_team(**context):
    ti = context['ti']
    errors = ti.xcom_pull(key='quality_errors')
    print(f"ALERT: Sending notification for {len(errors)} errors")
    for err in errors:
        print(f"  - {err}")

def transform_orders(**context):
    ti = context['ti']
    data = ti.xcom_pull(key='orders')
    transformed = [
        {**o, "category": "high" if o["amount"] > 200 else "standard"}
        for o in data["orders"] if o["amount"] > 0
    ]
    ti.xcom_push(key='transformed', value=transformed)
    print(f"Transformed {len(transformed)} valid orders")

# Task definitions
start = DummyOperator(task_id='start', dag=dag)
extract = PythonOperator(
    task_id='extract_orders',
    python_callable=extract_orders,
    dag=dag,
)
quality_check = BranchPythonOperator(
    task_id='check_data_quality',
    python_callable=check_data_quality,
    dag=dag,
)
alert = PythonOperator(
    task_id='alert_team',
    python_callable=alert_team,
    dag=dag,
)
transform = PythonOperator(
    task_id='transform_orders',
    python_callable=transform_orders,
    dag=dag,
)
load = BashOperator(
    task_id='load_to_warehouse',
    bash_command='echo "Loading {{ ti.xcom_pull(key=\'transformed\') | length }} records"',
    dag=dag,
)
end = DummyOperator(task_id='end', dag=dag)

# Dependencies with branching
start >> extract >> quality_check
quality_check >> alert >> end
quality_check >> transform >> load >> end

# sensor_example.py
# Airflow sensor implementation
from airflow.sensors.base import BaseSensorOperator
from airflow.utils.decorators import apply_defaults
import os

class FileExistenceSensor(BaseSensorOperator):
    @apply_defaults
    def __init__(self, filepath, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.filepath = filepath

    def poke(self, context):
        exists = os.path.exists(self.filepath)
        if exists:
            self.log.info(f"File {self.filepath} exists!")
        else:
            self.log.info(f"File {self.filepath} not found, waiting...")
        return exists

# Usage in DAG:
# wait_for_file = FileExistenceSensor(
#     task_id='wait_for_data_file',
#     filepath='/data/input/orders_20260601.csv',
#     poke_interval=60,
#     timeout=3600,
#     mode='reschedule',
#     dag=dag,
# )

# prefect_flow.py
# Prefect flow with retries and caching
from prefect import flow, task
from prefect.task_runners import ConcurrentTaskRunner
import httpx
import time

@task(retries=3, retry_delay_seconds=30, cache_key_fn=None)
def fetch_weather_data(city: str):
    """Fetch weather data from an API with retries."""
    response = httpx.get(f"https://api.weather.com/v1/{city}", timeout=10)
    response.raise_for_status()
    data = response.json()
    print(f"Fetched {len(data)} records for {city}")
    return data

@task
def process_weather_data(data: dict):
    """Transform raw weather data."""
    processed = {
        "city": data.get("name"),
        "temp_c": data.get("main", {}).get("temp"),
        "humidity": data.get("main", {}).get("humidity"),
        "timestamp": time.time(),
    }
    print(f"Processed: {processed['city']} - {processed['temp_c']}C")
    return processed

@task
def validate_temperature(processed: dict):
    """Data quality: check temperature range."""
    if processed["temp_c"] is None:
        raise ValueError("Temperature is missing")
    if processed["temp_c"] < -50 or processed["temp_c"] > 60:
        raise ValueError(f"Temperature {processed['temp_c']}C out of range")
    print(f"Validation passed for {processed['city']}")
    return processed

@task
def store_in_database(processed: dict):
    """Simulate database storage."""
    print(f"INSERT INTO weather VALUES ('{processed['city']}', {processed['temp_c']}, {processed['humidity']})")
    return {"status": "stored", "city": processed["city"]}

@flow(name="weather_pipeline", task_runner=ConcurrentTaskRunner())
def weather_pipeline(cities: list):
    """Orchestrate the full weather data pipeline."""
    results = []
    for city in cities:
        raw = fetch_weather_data(city)
        processed = process_weather_data(raw)
        validated = validate_temperature(processed)
        stored = store_in_database(validated)
        results.append(stored)
    return results

# Run the flow
if __name__ == "__main__":
    cities = ["London", "Tokyo", "New York"]
    results = weather_pipeline(cities)
    print(f"\nPipeline complete. {len(results)} cities processed.")
    for r in results:
        print(f"  {r['city']}: {r['status']}")

# dagster_pipeline.py
# Dagster software-defined assets pipeline
from dagster import asset, Definitions, Out, AssetIn
import pandas as pd

@asset
def raw_orders() -> pd.DataFrame:
    """Extract raw orders data."""
    data = pd.DataFrame([
        {"order_id": 1, "amount": 100.0, "status": "completed", "date": "2026-06-01"},
        {"order_id": 2, "amount": 250.0, "status": "completed", "date": "2026-06-01"},
        {"order_id": 3, "amount": -50.0, "status": "cancelled", "date": "2026-06-02"},
    ])
    print(f"Extracted {len(data)} orders")
    return data

@asset
def cleaned_orders(raw_orders: pd.DataFrame) -> pd.DataFrame:
    """Remove invalid orders with negative amounts."""
    before = len(raw_orders)
    cleaned = raw_orders[raw_orders["amount"] > 0].copy()
    removed = before - len(cleaned)
    print(f"Cleaned: removed {removed} invalid orders")
    return cleaned

@asset
def order_summary(cleaned_orders: pd.DataFrame) -> pd.DataFrame:
    """Compute daily order summary."""
    cleaned_orders["order_date"] = pd.to_datetime(cleaned_orders["date"])
    summary = cleaned_orders.groupby("order_date").agg(
        total_orders=("order_id", "count"),
        total_revenue=("amount", "sum"),
    ).reset_index()
    print("Order summary computed:")
    print(summary.to_string())
    return summary

# Definitions for Dagster UI
defs = Definitions(assets=[raw_orders, cleaned_orders, order_summary])

# .github/workflows/data-pipeline-ci.yml
name: Data Pipeline CI/CD
on:
  push:
    branches: [main]
    paths:
      - 'dags/**'
      - 'flows/**'
      - 'assets/**'

jobs:
  test:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4
      
      - name: Setup Python
        uses: actions/setup-python@v5
        with:
          python-version: '3.11'
      
      - name: Install dependencies
        run: |
          pip install -r requirements.txt
          pip install pytest pytest-cov
      
      - name: Run tests
        run: |
          python -m pytest tests/ --cov=dags --cov-report=xml
      
      - name: Validate DAGs
        run: |
          python -c "from airflow.models import DAG; exec(open('dags/quality_etl_dag.py').read()); print('DAG validation passed')"
      
      - name: Lint
        run: |
          pip install ruff
          ruff check dags/ --max-line-length=100

  deploy:
    needs: test
    runs-on: ubuntu-latest
    if: github.ref == 'refs/heads/main'
    steps:
      - uses: actions/checkout@v4
      
      - name: Deploy DAGs to Airflow
        run: |
          rsync -avz dags/ airflow-server:/opt/airflow/dags/

# monitoring.py
# Airflow callback for Slack alerts
from airflow.providers.slack.hooks.slack_webhook import SlackWebhookHook

def task_failure_alert(context):
    """Send Slack alert on task failure."""
    dag_id = context['dag'].dag_id
    task_id = context['task_instance'].task_id
    execution_date = context['execution_date']
    log_url = context['task_instance'].log_url
    
    message = f"""
    🔴 Pipeline Failed
    DAG: {dag_id}
    Task: {task_id}
    Execution: {execution_date}
    Logs: {log_url}
    """
    
    hook = SlackWebhookHook(
        webhook_token='T_TOKEN/B_TOKEN/xxxxx'
    )
    hook.send(text=message)
    print(f"Alert sent for {dag_id}.{task_id}")

# Add to DAG default_args:
# 'on_failure_callback': task_failure_alert,

# custom_block.py
# Custom Prefect block for API polling
from prefect.blocks.core import Block
from pydantic import Field
import httpx
import time

class ApiPoller(Block):
    """Block that polls an API endpoint until data is ready."""
    
    endpoint: str = Field(description="API endpoint URL")
    poll_interval: int = Field(default=30, description="Seconds between polls")
    timeout: int = Field(default=600, description="Max seconds to wait")
    
    def block_initialization(self):
        self.start_time = time.time()
    
    def poll(self) -> dict:
        """Poll the API until data is available or timeout."""
        while time.time() - self.start_time < self.timeout:
            response = httpx.get(self.endpoint, timeout=10)
            if response.status_code == 200 and response.json().get("ready"):
                print(f"Data ready after {time.time() - self.start_time:.1f}s")
                return response.json()
            
            elapsed = time.time() - self.start_time
            print(f"Waiting... ({elapsed:.0f}s elapsed)")
            time.sleep(self.poll_interval)
        
        raise TimeoutError(f"API not ready after {self.timeout}s")

# Usage:
# poller = ApiPoller(endpoint="https://api.example.com/status")
# result = poller.poll()

if __name__ == "__main__":
    # Simulate usage
    from unittest.mock import patch
    
    poller = ApiPoller(endpoint="https://api.example.com/status")
    
    with patch.object(httpx, 'get') as mock_get:
        mock_response = mock_get.return_value
        mock_response.status_code = 200
        mock_response.json.side_effect = [
            {"ready": False, "data": None},
            {"ready": False, "data": None},
            {"ready": True, "data": {"records": 100}},
        ]
        
        result = poller.poll()
        print(f"Polling result: {result}")