Prefect: Modern Workflow Orchestration

from prefect import flow, task, get_run_logger
from prefect.tasks import task_input_hash
from prefect.cache_policies import NO_CACHE, INPUTS
from prefect.deployments import Deployment
from datetime import timedelta, datetime
from typing import List, Dict, Optional
import pandas as pd
import hashlib
import time
import logging

logger = logging.getLogger(__name__)


# ------------------------------------------------------
# TASKS
# ------------------------------------------------------
@task(
    name="extract-orders",                    # Task name for UI display
    retries=3,                                # Retry up to 3 times on failure
    retry_delay_seconds=10,                   # Wait 10 seconds between retries
    cache_key_fn=task_input_hash,             # Cache based on function inputs
    cache_expiration=timedelta(hours=1),      # Cache expires after 1 hour
    log_prints=True,                          # Capture print() statements in logs
)
def extract_orders(execution_date: str) -> pd.DataFrame:
    """
    Extract orders from source database for the given date.
    
    Parameters:
        execution_date (str): Date string in 'YYYY-MM-DD' format
        
    Returns:
        pd.DataFrame: Order records for the specified date
        
    Caching behavior:
        If same execution_date is passed within 1 hour, cached result is returned
        without hitting the database.
    """
    import sqlalchemy

    engine = sqlalchemy.create_engine("postgresql://user:pass@host/db")
    query = f"""
        SELECT order_id, customer_id, amount, status, created_at
        FROM orders
        WHERE DATE(created_at) = '{execution_date}'
    """
    df = pd.read_sql(query, engine)
    logger.info(f"Extracted {len(df)} orders for {execution_date}")
    return df


@task(
    name="validate-raw-data",
    retries=2,
    retry_delay_seconds=5,
)
def validate_raw_data(df: pd.DataFrame) -> pd.DataFrame:
    """Validate raw data quality before transformation."""
    initial_count = len(df)

    # Remove nulls in critical columns
    df = df.dropna(subset=["order_id", "customer_id", "amount"])

    # Remove negative amounts
    df = df[df["amount"] > 0]

    # Remove duplicates
    df = df.drop_duplicates(subset=["order_id"])

    validation_stats = {
        "initial_count": initial_count,
        "final_count": len(df),
        "removed_count": initial_count - len(df),
        "removal_rate": (initial_count - len(df)) / initial_count if initial_count > 0 else 0,
    }

    if validation_stats["removal_rate"] > 0.1:
        raise ValueError(
            f"Data quality issue: {validation_stats['removal_rate']:.1%} "
            f"of records removed during validation"
        )

    logger.info(f"Validation passed: {validation_stats}")
    return df


@task(
    name="transform-orders",
    retries=2,
    retry_delay_seconds=30,
)
def transform_orders(df: pd.DataFrame) -> pd.DataFrame:
    """Transform orders with enrichment and aggregation."""
    # Enrich with customer data
    import sqlalchemy
    engine = sqlalchemy.create_engine("postgresql://user:pass@host/db")
    customer_df = pd.read_sql("SELECT customer_id, tier, segment FROM customers", engine)

    enriched = df.merge(customer_df, on="customer_id", how="left")
    enriched["tier"] = enriched["tier"].fillna("unknown")

    # Calculate derived fields
    enriched["order_date"] = pd.to_datetime(enriched["created_at"]).dt.date
    enriched["amount_bucket"] = pd.cut(
        enriched["amount"],
        bins=[0, 10, 100, 1000, float("inf")],
        labels=["micro", "small", "medium", "large"],
    )

    logger.info(f"Transformed {len(enriched)} orders")
    return enriched


@task(
    name="load-to-warehouse",
    retries=3,                    # More retries for load phase (critical)
    retry_delay_seconds=60,       # Longer delay for database recovery
)
def load_to_warehouse(df: pd.DataFrame, target_table: str) -> Dict:
    """
    Load transformed data to the warehouse using idempotent upsert.
    
    Parameters:
        df (pd.DataFrame): Transformed data to load
        target_table (str): Target table name in the warehouse
        
    Returns:
        Dict: Load statistics with record count and target table
        
    Idempotency:
        Uses ON CONFLICT DO UPDATE (UPSERT) to ensure repeated runs
        produce the same result without duplicates.
    """
    import sqlalchemy

    engine = sqlalchemy.create_engine("postgresql://user:pass@host/db")

    # Use temporary table for staging the upsert
    temp_table = f"{target_table}_temp"
    df.to_sql(temp_table, engine, if_exists="replace", index=False)

    # Upsert via MERGE (PostgreSQL: INSERT ... ON CONFLICT)
    merge_query = f"""
        INSERT INTO {target_table}
        SELECT * FROM {temp_table}
        ON CONFLICT (order_id)
        DO UPDATE SET
            amount = EXCLUDED.amount,
            status = EXCLUDED.status,
            tier = EXCLUDED.tier,
            segment = EXCLUDED.segment
    """
    engine.execute(merge_query)
    engine.execute(f"DROP TABLE {temp_table}")

    result = {"loaded_count": len(df), "target_table": target_table}
    logger.info(f"Loaded {len(df)} records to {target_table}")
    return result


@task(name="send-notification")
def send_notification(message: str, severity: str = "info"):
    """Send notification about pipeline completion."""
    logger.info(f"Notification [{severity}]: {message}")
    # In production, integrate with Slack, PagerDuty, etc.


# ------------------------------------------------------
# FLOW
# ------------------------------------------------------
@flow(
    name="daily-order-etl",
    description="Daily order extraction  transformation  and warehouse load",
    log_prints=True,
    retry_delay_seconds=300,
    retries=1,
)
def daily_order_etl(execution_date: str = None):
    """Main ETL flow orchestrating extraction, transformation, and loading."""
    if execution_date is None:
        execution_date = datetime.now().strftime("%Y-%m-%d")

    logger.info(f"Starting daily order ETL for {execution_date}")

    # Step 1: Extract
    raw_df = extract_orders(execution_date)

    # Step 2: Validate
    validated_df = validate_raw_data(raw_df)

    # Step 3: Transform
    transformed_df = transform_orders(validated_df)

    # Step 4: Load
    load_result = load_to_warehouse(transformed_df, "orders_fact")

    # Step 5: Notify
    send_notification(
        message=f"ETL completed for {execution_date}: {load_result['loaded_count']} records loaded",
        severity="info",
    )

    return load_result


# ------------------------------------------------------
# DEPLOYMENT
# ------------------------------------------------------
if __name__ == "__main__":
    # Run locally for testing
    result = daily_order_etl(execution_date="2024-01-15")
    print(result)

    # Deploy to Prefect Cloud
    deployment = Deployment.build_from_flow(
        daily_order_etl,
        name="daily-schedule",
        schedule={
            "cron": "0 6 * * *",
            "timezone": "UTC",
        },
        work_pool_name="default-agent-pool",
        parameters={"execution_date": "{{ data_interval_end }}"},
    )
    deployment.apply()

from prefect import flow, task, get_run_logger
from typing import List
import pandas as pd


@task(name="discover-tables")
def discover_tables(connection_string: str) -> List[str]:
    """Dynamically discover tables to process."""
    import sqlalchemy
    engine = sqlalchemy.create_engine(connection_string)
    with engine.connect() as conn:
        result = conn.execute(
            "SELECT table_name FROM information_schema.tables "
            "WHERE table_schema = 'public' AND table_name LIKE 'events_%'"
        )
        tables = [row[0] for row in result]
    logger = get_run_logger()
    logger.info(f"Discovered {len(tables)} tables: {tables}")
    return tables


@task(name="extract-table", retries=2, retry_delay_seconds=10)
def extract_table(connection_string: str, table_name: str) -> str:
    """Extract a single table to a parquet file."""
    import sqlalchemy
    engine = sqlalchemy.create_engine(connection_string)
    df = pd.read_sql(f"SELECT * FROM {table_name}", engine)
    output_path = f"/tmp/raw/{table_name}.parquet"
    df.to_parquet(output_path, index=False)
    logger = get_run_logger()
    logger.info(f"Extracted {len(df)} rows from {table_name}")
    return output_path


@task(name="validate-table")
def validate_table(table_name: str, file_path: str) -> dict:
    """Validate extracted data quality."""
    df = pd.read_parquet(file_path)
    null_rate = df.isnull().mean().mean()
    duplicate_rate = 1 - (len(df.drop_duplicates()) / len(df))

    result = {
        "table": table_name,
        "row_count": len(df),
        "null_rate": float(null_rate),
        "duplicate_rate": float(duplicate_rate),
        "is_valid": null_rate < 0.1 and duplicate_rate < 0.05,
    }

    logger = get_run_logger()
    if not result["is_valid"]:
        logger.warning(f"Validation FAILED for {table_name}: {result}")
    else:
        logger.info(f"Validation passed for {table_name}")

    return result


@task(name="aggregate-results")
def aggregate_results(validations: List[dict]) -> dict:
    """Aggregate validation results across all tables."""
    total = len(validations)
    valid = sum(1 for v in validations if v["is_valid"])
    return {
        "total_tables": total,
        "valid_tables": valid,
        "validity_rate": valid / total if total > 0 else 0,
        "details": validations,
    }


@flow(name="dynamic-table-pipeline")
def dynamic_table_pipeline(connection_string: str):
    """Dynamic pipeline that processes discovered tables in parallel."""
    # Discover tables dynamically
    tables = discover_tables(connection_string)

    # Map extraction across tables (parallel execution)
    raw_files = extract_table.map(
        connection_string=connection_string,
        table_name=tables,
    )

    # Map validation across tables
    validations = validate_table.map(
        table_name=tables,
        file_path=raw_files,
    )

    # Aggregate results
    summary = aggregate_results(validations)

    logger = get_run_logger()
    logger.info(f"Pipeline summary: {summary}")

    if summary["validity_rate"] < 0.9:
        raise ValueError(
            f"Data quality below threshold: {summary['validity_rate']:.1%} "
            f"of tables passed validation"
        )

    return summary


# Run the dynamic pipeline
if __name__ == "__main__":
    result = dynamic_table_pipeline("postgresql://user:pass@host/db")
    print(result)