Super-charge Your SaaS & LLM Workflows

import requests

@task
def get_config():
    response = requests.get('https://api.config.com/settings')
    return response.json()

@task
def process_data(config):
    return config['processing_settings']
    
with DAG('your_dag', ...) as dag:
    config = get_config()
    process_data(config)

from airflow.sensors.base import BaseSensorOperator
import requests

# Replace with EcsTaskStateSensor or your own custom sensor
# implementation depending on what process was launched upstream
class ProcessingJobSensor(BaseSensorOperator):
    def __init__(
        self,
        *,
        job_id: str,
        api_url: str,
        **kwargs
    ) -> None:
        super().__init__(**kwargs)
        self.job_id = job_id
        self.api_url = api_url

    def poke(self, context) -> bool:
        response = requests.get(f"{self.api_url}/jobs/{self.job_id}")
        status = response.json()["status"]
        
        if status == "completed":
            return True
        if status == "failed":
            raise AirflowException("Processing job failed")
        return False

@task
def start_processing():
    # Start a pretend processing job via API.
    # This also pairs perfectly with containerized jobs launched 
    # using EcsRunTaskOperator and KubernetesPodOperator
    response = requests.post(
        "https://processing-service.example.com/jobs",
        json={
            "input_path": "s3://bucket/huge_file.csv",
            "output_path": "s3://bucket/output.csv"
        }
    )
    return response.json()["job_id"]

with DAG('process_data_dag',
    start_date=days_ago(1),
    schedule_interval='@daily',
) as dag:
    job = start_processing()
    wait_for_processing = ProcessingJobSensor(
        task_id='wait_for_processing',
        job_id="{{ task_instance.xcom_pull(task_ids='start_processing') }}",
        api_url="https://processing-service.example.com",
        mode="reschedule" # reschedule will cause the task to be retried without blocking the worker
    )
    job >> wait_for_processing

@task
def process_sales_data(date):
    # Deterministic filename based on logical date
    output_path = f's3://bucket/sales_{date:%Y/%m/%d}.csv'
    
    # Idempotent database operation
    with postgres.connect() as conn:
        conn.execute("""
            INSERT INTO sales_summary (date, total)
            SELECT %s, SUM(amount)
            FROM daily_sales
            WHERE date = %s
            ON CONFLICT (date) 
            DO UPDATE SET total = EXCLUDED.total
        """, [date, date])
    
    return output_path

with DAG(
    'sales_processing',
    start_date=days_ago(1),
    schedule_interval='@daily'
) as dag:
    process_sales_data('{{ ds }}')

@task
def fetch_data():
    response = requests.get('https://api.example.com/data')
    return response.json()

@task
def process_data(raw_data):
    df = pd.DataFrame(raw_data)
    df['processed'] = df['value'].apply(complex_calculation)
    return df.to_dict()

@task
def load_to_database(processed_data):
    df = pd.DataFrame(processed_data)
    df.to_sql('processed_data', engine)
    return "Database load complete"

@task
def save_to_s3(processed_data):
    df = pd.DataFrame(processed_data)
    path = 's3://bucket/processed_data.parquet'
    df.to_parquet(path)
    return path

@task
def notify_completion(db_status, s3_path):
    send_slack_message(
        f"Processing complete!\nDB Status: {db_status}\nS3 Path: {s3_path}"
    )

with DAG('atomic_processing_dag', ...) as dag:
    data = fetch_data()
    processed = process_data(data)
    db_status = load_to_database(processed)
    s3_path = save_to_s3(processed)
    notify_completion(db_status, s3_path)

@task
def process_data():
    # Token fetched just before it's needed
    token = get_oauth_token()
    headers = {"Authorization": f"Bearer {token}"}
    response = requests.get("https://api.example.com/data", headers=headers)
    return process(response.json())

with DAG('token_dag_good', ...) as dag:
    process_data()

@task
def get_config():
    # Config fetched during execution
    return fetch_config()['items']

@task
def process_item(item):
    return process_data(item)

@task
def summarize_results(results):
    return {
        'processed_count': len(results),
        'success_count': len([r for r in results if r['status'] == 'success'])
    }

with DAG('dynamic_dag_good', ...) as dag:
    config = get_config()
    results = process_item.expand(item=config)
    summarize_results(results)

@task(retries=3, retry_delay=timedelta(minutes=5))
def load_legacy_data():
    return pd.read_csv('old_data.csv').to_dict()

@task
def transform_data(data):
    df = pd.DataFrame(data)
    df['processed'] = df['value'] * 1.1
    return df.to_dict()

@task
def save_to_new_table(data):
    hook = PostgresHook(postgres_conn_id='your_postgres')
    df = pd.DataFrame(data)
    df.to_sql('new_table', hook.get_sqlalchemy_engine())
    return {'rows_processed': len(df)}

@task
def notify_completion(stats):
    send_slack_message(
        f"Data migration completed. {stats['rows_processed']} rows processed."
    )

with DAG(
    'one_time_migration_2024_03',
    schedule_interval=None,  # Manual triggers only
    start_date=days_ago(1),
    catchup=False,
    tags=['one-time', 'migration'],
) as dag:
    data = load_legacy_data()
    processed = transform_data(data)
    stats = save_to_new_table(processed)
    notify_completion(stats)
                            

@task
def generate_report(execution_date):
    df = process_data()
    # Using S3 with date partitioning
    path = f's3://your-bucket/reports/{execution_date:%Y/%m/%d}/report.csv'
    df.to_csv(path)
    return path

@task
def send_report(file_path):
    s3_hook = S3Hook()
    with s3_hook.get_key(file_path).get()['Body'] as f:
        send_email(
            to='team@company.com',
            subject='Daily Report',
            attachment=f
        )
    # Optionally move to archive location
    archive_path = file_path.replace('/reports/', '/archive/reports/')
    s3_hook.copy_object(file_path, archive_path)

with DAG('report_dag_good', ...) as dag:
    file_path = generate_report('{{ ds }}')
    send_report(file_path)

from airflow.exceptions import AirflowException
from requests.exceptions import RequestException
import tenacity

class DataQualityError(AirflowException):
    """Raised when data doesn't meet quality standards."""
    pass

@task(
    retries=3,
    retry_delay=timedelta(minutes=5),
    retry_exponential_backoff=True,
    max_retry_delay=timedelta(minutes=30),
    # Retry only on specific exceptions
    retry_on_exception=lambda e: isinstance(e, RequestException)
)
def process_api_data():
    try:
        response = requests.get(
            'https://api.example.com/data',
            timeout=10
        )
        response.raise_for_status()
        data = response.json()
        
        # Validate data quality
        if not validate_data(data):
            raise DataQualityError("Data validation failed")
            
        process_data(data)
        
    except RequestException as e:
        # Network/API errors - will be retried
        send_alert("API Error", f"Retrying due to: {str(e)}")
        raise
        
    except DataQualityError as e:
        # Data quality issues - no retry
        send_alert("Data Quality Error", str(e), severity="high")
        raise
        
    except Exception as e:
        # Unexpected errors
        send_alert(
            "Critical Error",
            f"Unexpected error: {str(e)}",
            severity="critical"
        )
        raise

with DAG('robust_error_dag', ...) as dag:
    process_api_data()