@dag(schedule_interval=timedelta(hours=1), start_date=datetime(2024, 1, 1), catchup=False)
def taskflow_example():
    @task
    def extract_data():
        return [{"id": 1, "name": "Alice"}]  # Auto-pushed to XCom

    @task
    def transform_data(raw_data: list) -> list:  # Auto-pulled from XCom
        return [r.upper() for r in raw_data]

    raw = extract_data()
    transformed = transform_data(raw)  # Dependency auto-inferred

@task
def extract_metrics() -> Tuple[List[Dict], Dict]:
    return records, metadata  # Each pushed as separate XCom key

@task
def process_records(records: list, metadata: dict):  # Pull by index
    print(f"Processing {metadata['total_records']} records")

@task
def get_regions() -> list:
    return ["us-east-1", "us-west-2", "eu-west-1"]

@task
def process_region(region: str) -> dict:
    return {"region": region, "status": "success"}

@task
def aggregate_results(region_results: list) -> dict:
    return {"total_regions": len(region_results)}

# Usage
regions = get_regions()
results = process_region.map(regions)  # Dynamic mapping
aggregate_results(results)

from airflow.decorators import task, dag
from datetime import datetime, timedelta
from typing import Optional, List, Dict
import json

@dag(
    schedule_interval="0 6 * * *",
    start_date=datetime(2024, 1, 1),
    catchup=False,
    tags=['taskflow', 'advanced'],
    doc_md="""
    ## Advanced TaskFlow Patterns
    Demonstrates TaskFlow API with:
    - Multiple return values
    - Dynamic task mapping
    - Error handling
    - Custom XCom keys
    """,
)
def advanced_taskflow_dag():
    
    @task(retries=3, retry_delay=timedelta(minutes=1))
    def extract_from_source(source: str) -> List[Dict]:
        """Extract data with retry logic."""
        import random
        if random.random() < 0.1:
            raise ConnectionError(f"Failed to connect to {source}")
        
        return [
            {"id": i, "source": source, "value": random.randint(1, 100)}
            for i in range(10)
        ]
    
    @task
    def validate_data(records: List[Dict]) -> tuple:
        """Validate and separate valid/invalid records."""
        valid = [r for r in records if 0 <= r["value"] <= 100]
        invalid = [r for r in records if r["value"] < 0 or r["value"] > 100]
        return valid, invalid
    
    @task
    def transform_record(record: Dict) -> Dict:
        """Transform a single record (used with .map())."""
        return {
            **record,
            "value_normalized": record["value"] / 100.0,
            "transformed": True,
        }
    
    @task
    def load_batch(records: List[Dict], destination: str) -> int:
        """Load a batch of records."""
        print(f"Loading {len(records)} records to {destination}")
        return len(records)
    
    @task
    def generate_report(loaded_counts: list, invalid_records: list) -> str:
        """Generate summary report."""
        total_loaded = sum(loaded_counts)
        total_invalid = len(invalid_records)
        
        report = {
            "total_loaded": total_loaded,
            "total_invalid": total_invalid,
            "success_rate": total_loaded / (total_loaded + total_invalid) * 100,
        }
        return json.dumps(report, indent=2)
    
    # Define sources
    sources = ["postgres", "mysql", "mongodb"]
    
    # Extract from multiple sources (dynamic mapping)
    raw_data = extract_from_source.map(sources)
    
    # Validate each source's data
    valid_data, invalid_data = validate_data.expand(raw_data)
    
    # Transform valid records (dynamic mapping)
    transformed = transform_record.map(valid_data)
    
    # Load to destination
    loaded_count = load_batch(transformed, "data_warehouse")
    
    # Generate report
    generate_report(loaded_count, invalid_data)

advanced_taskflow_dag()

from airflow.decorators import task, dag
from airflow.models.xcom import BaseXCom
from datetime import datetime, timedelta
import json
import pickle

class CustomXComBackend(BaseXCom):
    """Custom XCom backend for large payloads."""
    
    @staticmethod
    def serialize_value(value):
        """Serialize value for storage."""
        if isinstance(value, (dict, list)):
            return json.dumps(value)
        return str(value)
    
    @staticmethod
    def deserialize_value(result):
        """Deserialize value from storage."""
        try:
            return json.loads(result)
        except (json.JSONDecodeError, TypeError):
            return result

@dag(
    schedule_interval="@hourly",
    start_date=datetime(2024, 1, 1),
    catchup=False,
    tags=['taskflow', 'xcom-backend'],
    render_template_as_native_obj=True,
)
def xcom_backend_dag():
    
    @task(
        xcom_backend='path.to.CustomXComBackend',
        pool='data_processing',
    )
    def process_large_dataset(records: list) -> dict:
        """Process dataset using custom XCom backend."""
        import pandas as pd
        
        df = pd.DataFrame(records)
        
        result = {
            "mean": float(df["value"].mean()),
            "std": float(df["value"].std()),
            "count": int(len(df)),
            "percentiles": {
                "25": float(df["value"].quantile(0.25)),
                "50": float(df["value"].quantile(0.50)),
                "75": float(df["value"].quantile(0.75)),
            },
        }
        return result
    
    @task
    def summarize(stats: dict) -> str:
        """Create summary string."""
        return f"Processed {stats['count']} records, mean={stats['mean']:.2f}"
    
    sample_data = [{"value": i * 1.5} for i in range(100)]
    stats = process_large_dataset(sample_data)
    summarize(stats)

xcom_backend_dag()

from airflow.decorators import task, dag
from datetime import datetime, timedelta

@dag(
    schedule_interval="@daily",
    start_date=datetime(2024, 1, 1),
    catchup=False,
    tags=['taskflow', 'taskmap', 'cross-dependencies'],
)
def cross_task_dependency_dag():
    
    @task
    def get_config() -> dict:
        """Return configuration for downstream tasks."""
        return {
            "environments": ["dev", "staging", "prod"],
            "parallel_workers": 4,
            "timeout_seconds": 300,
        }
    
    @task
    def setup_environment(env: str, config: dict) -> dict:
        """Setup each environment using config."""
        return {
            "environment": env,
            "status": "ready",
            "workers": config["parallel_workers"],
        }
    
    @task
    def run_tests(env_setup: dict) -> dict:
        """Run tests in each environment."""
        return {
            "environment": env_setup["environment"],
            "tests_passed": 42,
            "tests_failed": 0,
        }
    
    @task
    def cleanup_environments(results: list) -> str:
        """Cleanup all environments after tests."""
        environments = [r["environment"] for r in results]
        return f"Cleaned up: {', '.join(environments)}"
    
    # Get configuration
    config = get_config()
    
    # Setup environments dynamically using config
    env_setups = setup_environment.map(
        env=config["environments"],
        config=config,  # Same config passed to all mapped tasks
    )
    
    # Run tests in each environment
    test_results = run_tests.expand(env_setups)
    
    # Cleanup
    cleanup_environments(test_results)

cross_task_dependency_dag()