from datetime import datetime, timedelta
from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator
from airflow.sensors.s3_key_sensor import S3KeySensor

default_args = {
    "owner": "data-engineering",
    "depends_on_past": False,
    "email_on_failure": True,
    "email": ["data-team@example.com"],
    "retries": 2,
    "retry_delay": timedelta(minutes=5),
}

with DAG(
    dag_id="daily_orders_pipeline",
    default_args=default_args,
    description="Ingest, transform, and validate daily orders",
    schedule="0 6 * * *",  # 06:00 UTC daily
    start_date=datetime(2025, 1, 1),
    catchup=False,
    tags=["orders", "daily"],
) as dag:

    # 1. Wait for source file
    wait_for_file = S3KeySensor(
        task_id="wait_for_orders_file",
        bucket_name="raw-data",
        bucket_key="orders/{{ ds }}/orders.parquet",
        poke_interval=300,  # check every 5 minutes
        timeout=3600,       # fail after 1 hour
        mode="poke",
    )

    # 2. Extract and load raw data
    def extract_orders(**context):
        execution_date = context["ds"]
        # Read from S3, apply schema, write to staging
        # ... (Spark submit or direct S3 read)
        return f"Extracted orders for {execution_date}"

    extract = PythonOperator(
        task_id="extract_orders",
        python_callable=extract_orders,
    )

    # 3. Transform with SQL
    transform = SQLExecuteQueryOperator(
        task_id="transform_orders",
        conn_id="warehouse_conn",
        sql="sql/transform_orders.sql",
        params={"run_date": "{{ ds }}"},
    )

    # 4. Data quality check
    def validate_orders(**context):
        execution_date = context["ds"]
        # Run validation queries, raise on failure
        row_count = run_query(f"SELECT COUNT(*) FROM fact_orders WHERE order_date = '{execution_date}'")
        if row_count == 0:
            raise ValueError(f"No rows loaded for {execution_date}")

    validate = PythonOperator(
        task_id="validate_orders",
        python_callable=validate_orders,
    )

    # Dependencies
    wait_for_file >> extract >> transform >> validate

# Backfill orders pipeline for the first two weeks of June
airflow dags backfill daily_orders_pipeline \
    --start-date 2025-06-01 \
    --end-date 2025-06-15

# Default: all upstream tasks must succeed
task_a >> task_b  # task_b runs only if task_a succeeds

# Trigger rules for non-default behaviour
from airflow.utils.trigger_rule import TriggerRule

cleanup = PythonOperator(
    task_id="cleanup",
    python_callable=cleanup_fn,
    trigger_rule=TriggerRule.ALL_DONE,  # runs regardless of upstream success/failure
)

from prefect import flow, task
from prefect.tasks import task_input_hash
from datetime import timedelta

@task(retries=3, retry_delay_seconds=60, cache_key_fn=task_input_hash, cache_expiration=timedelta(hours=1))
def extract_orders(date: str) -> dict:
    # Extract logic here
    return {"rows": 1500, "date": date}

@task
def transform_orders(raw_data: dict) -> dict:
    # Transform logic here
    return {"rows": raw_data["rows"], "transformed": True}

@task
def validate_orders(data: dict):
    assert data["rows"] > 0, "No rows in output"

@flow(name="daily-orders-pipeline", log_prints=True)
def daily_orders(date: str):
    raw = extract_orders(date)
    transformed = transform_orders(raw)
    validate_orders(transformed)
    print(f"Pipeline complete for {date}: {transformed['rows']} rows")

# Run
daily_orders("2025-06-15")

from dagster import asset, Definitions, ScheduleDefinition

@asset
def raw_orders():
    """Extract orders from source system."""
    return extract_from_source("orders")

@asset
def clean_orders(raw_orders):
    """Clean and validate order data."""
    df = raw_orders
    df = df.dropna(subset=["order_id"])
    df = df[df["amount"] > 0]
    return df

@asset
def order_summary(clean_orders):
    """Aggregate orders by customer."""
    return clean_orders.groupby("customer_id").agg(
        total_spent=("amount", "sum"),
        order_count=("order_id", "count")
    )

daily_schedule = ScheduleDefinition(
    job=define_asset_job("daily_orders_job", selection=[raw_orders, clean_orders, order_summary]),
    cron_schedule="0 6 * * *",
)

defs = Definitions(
    assets=[raw_orders, clean_orders, order_summary],
    schedules=[daily_schedule],
)