DE Pipeline
Shipment
Analytics

@retry(stop=stop_after_attempt(3),
       wait=wait_exponential(multiplier=1, min=2, max=10))
def fetch_shipments():
    resp = requests.get(
        f"{API_BASE}/shipments",
        timeout=30
    )
    resp.raise_for_status()
    return resp.json()

def extract_to_staging(conn):
    data = fetch_shipments()
    cur = conn.cursor()
    cur.execute("TRUNCATE staging.shipments")
    cur.executemany(
        """INSERT INTO staging.shipments
           (shipment_id, customer_id, cost, status, ship_date)
           VALUES (%s,%s,%s,%s,%s)""",
        [(r['id'],r['customer'],r['cost'],
          r['status'],r['date']) for r in data]
    )
    conn.commit()

def extract_tiers_to_staging(conn):
    df = pd.read_csv(
        DATA_DIR / "customer_tiers.csv",
        dtype={"customer_id": str, "tier": str}
    )
    # Validate expected columns exist
    required = {"customer_id", "tier"}
    assert required.issubset(df.columns), \
        f"Missing columns: {required - set(df.columns)}"

    cur = conn.cursor()
    cur.execute("TRUNCATE staging.customer_tiers")
    for _, row in df.iterrows():
        cur.execute(
            """INSERT INTO staging.customer_tiers
               (customer_id, tier) VALUES (%s, %s)""",
            (row["customer_id"], row["tier"])
        )
    conn.commit()

def transform(conn):
    cur = conn.cursor()
    cur.execute("TRUNCATE staging.shipments_with_tiers")
    cur.execute("""
        INSERT INTO staging.shipments_with_tiers
        SELECT
            s.shipment_id,
            s.customer_id,
            s.cost,
            s.ship_date,
            COALESCE(t.tier, 'Unknown') AS tier
        FROM (
            -- Deduplicate: keep latest per shipment_id
            SELECT DISTINCT ON (shipment_id) *
            FROM staging.shipments
            ORDER BY shipment_id, created_at DESC
        ) s
        LEFT JOIN staging.customer_tiers t
            ON s.customer_id = t.customer_id
        WHERE
            s.customer_id IS NOT NULL    -- Gate 1: nulls
            AND s.cost > 0               -- Gate 2: negatives
            AND s.status != 'CANCELLED'  -- Gate 3: cancelled
    """)
    conn.commit()

def load_analytics(conn):
    """Idempotent: TRUNCATE + INSERT every run."""
    cur = conn.cursor()
    cur.execute(
        "TRUNCATE analytics.shipping_spend_by_tier"
    )
    cur.execute("""
        INSERT INTO analytics.shipping_spend_by_tier
            (tier, month, total_spend, shipment_count)
        SELECT
            tier,
            DATE_TRUNC('month', ship_date) AS month,
            SUM(cost)                      AS total_spend,
            COUNT(*)                       AS shipment_count
        FROM staging.shipments_with_tiers
        GROUP BY tier, DATE_TRUNC('month', ship_date)
        ORDER BY month, tier
    """)
    conn.commit()
    cur.execute(
        "SELECT COUNT(*) FROM "
        "analytics.shipping_spend_by_tier"
    )
    n = cur.fetchone()[0]
    print(f"Loaded {n} rows into analytics")