Practice — Implement Medallion Lakehouse on S3

# exercise_medallion_lakehouse.py — Step 1: Bronze, Lambda, and Iceberg metadata
import duckdb
import pandas as pd
import numpy as np
import json
import uuid
from datetime import datetime, timezone, timedelta, date
from collections import defaultdict

np.random.seed(42)

# ── Simulated S3 object store ──────────────────────────────────────────────────
S3_STORE: dict[str, pd.DataFrame] = {}  # s3_path → DataFrame
PIPELINE_RUNS: dict[str, str]     = {}  # source_file → run_id (idempotency)

def s3_put(path: str, df: pd.DataFrame) -> None:
    """Simulate S3 object write — immutable after first write."""
    if path in S3_STORE:
        raise ValueError(f"Immutability violation: {path} already exists")
    S3_STORE[path] = df.copy()

def s3_get(path: str) -> pd.DataFrame:
    return S3_STORE[path].copy()

# ── Lambda-style trigger ──────────────────────────────────────────────────────
def lambda_on_s3_event(event: dict) -> dict:
    """
    Triggered when a new scorecard file lands in S3.
    Idempotent: only starts one pipeline run per source file.
    """
    s3_key = event["Records"][0]["s3"]["object"]["key"]

    if not s3_key.endswith(".csv") or "scorecards" not in s3_key:
        return {"action": "skipped", "reason": "not a scorecard CSV"}

    if s3_key in PIPELINE_RUNS:
        return {"action": "skipped", "reason": "already_started",
                "run_id": PIPELINE_RUNS[s3_key]}

    run_id = str(uuid.uuid4())[:8]
    PIPELINE_RUNS[s3_key] = run_id
    return {"action": "started", "run_id": run_id, "source_key": s3_key}

# ── Simulate scorecard arrival ─────────────────────────────────────────────────
SCORECARD_S3_KEY = "raw/scorecards/2024-04-20/scorecard.csv"
event = {"Records": [{"s3": {"object": {"key": SCORECARD_S3_KEY}}}]}

result1 = lambda_on_s3_event(event)
assert result1["action"] == "started"
result2 = lambda_on_s3_event(event)  # duplicate trigger
assert result2["action"] == "skipped"
assert len(PIPELINE_RUNS) == 1
print(f"  Lambda trigger: {result1['action']} (run={result1['run_id']}) ✓")
print(f"  Duplicate trigger: {result2['action']} (idempotency ✓)")

# ── Bronze landing: raw data + metadata columns ───────────────────────────────
raw_deliveries = pd.DataFrame({
    "delivery_id":  [str(i) for i in range(1, 241)] + ["1", "2"],  # 2 duplicates
    "match_id":     ["10001"]*120 + ["10002"]*120 + ["10001"]*2,
    "over":         [str((i//6)+1) for i in range(242)],
    "runs_scored":  [str(int(np.random.choice([0,1,2,4,6]))) for _ in range(242)],
    "is_wicket":    [str(np.random.random() < 0.05) for _ in range(242)],
    "bowler":       list(np.random.choice(["Bumrah","Shami","Hardik",None], 242, p=[0.3,0.3,0.3,0.1])),
    "batter":       list(np.random.choice(["Rohit","Kohli","Gill"], 242)),
})
raw_deliveries["_ingested_at"] = datetime.now(timezone.utc).isoformat()
raw_deliveries["_source_file"] = SCORECARD_S3_KEY

bronze_path = f"bronze/deliveries/year=2024/month=04/day=20/deliveries.parquet"
s3_put(bronze_path, raw_deliveries)

bronze_df = s3_get(bronze_path)
assert "_ingested_at" in bronze_df.columns
assert "_source_file" in bronze_df.columns
assert len(bronze_df) == 242  # includes duplicates — Bronze never filters

# Iceberg snapshot simulation
ibronze_snapshot = {
    "snapshot_id":    1,
    "operation":      "append",
    "table_path":     f"s3://ipl-data-lake/{bronze_path}",
    "summary":        {"added_records": len(bronze_df), "total_records": len(bronze_df)},
    "schema_id":      0,
    "sequence_number":1,
}
print(f"  Bronze: {len(bronze_df)} raw rows (including duplicates) ✓")
print(f"  Iceberg snapshot: {ibronze_snapshot['summary']}")

# Immutability test: Bronze write should fail
try:
    s3_put(bronze_path, bronze_df)
    assert False, "Should have raised immutability error"
except ValueError as e:
    print(f"  Immutability: write blocked as expected ✓")
print("Step 1 ✓: Bronze landing, Lambda trigger, and Iceberg metadata complete")

# exercise_medallion_lakehouse.py — Step 2: Silver, Gold, schema evolution, cost tags
import duckdb
import pandas as pd
import numpy as np

np.random.seed(42)

con = duckdb.connect(":memory:")
con.register("bronze_raw", bronze_df)

# ── Silver transformation ─────────────────────────────────────────────────────
silver_df = con.execute("""
WITH dedup AS (
    SELECT *, ROW_NUMBER() OVER (PARTITION BY delivery_id ORDER BY _ingested_at) AS rn
    FROM bronze_raw
)
SELECT
    CAST(delivery_id AS BIGINT)     AS delivery_id,
    CAST(match_id    AS INTEGER)    AS match_id,
    CAST(over        AS INTEGER)    AS over,
    CAST(runs_scored AS INTEGER)    AS runs_scored,
    CAST(is_wicket   AS BOOLEAN)    AS is_wicket,
    bowler,
    batter,
    CAST(_ingested_at AS TIMESTAMPTZ) AS ingested_at,
    CASE WHEN CAST(over AS INTEGER) <= 6  THEN 'powerplay'
         WHEN CAST(over AS INTEGER) <= 15 THEN 'middle'
         ELSE 'death' END            AS phase
FROM dedup
WHERE rn = 1
  AND bowler IS NOT NULL
  AND CAST(runs_scored AS INTEGER) BETWEEN 0 AND 6
""").df()

silver_path = "silver/deliveries/year=2024/month=04/day=20/deliveries.parquet"
S3_STORE[silver_path] = silver_df  # Silver is mutable (MERGE-based)

assert len(silver_df) < len(bronze_df), "Silver must have fewer rows after dedup+quality"
assert silver_df["bowler"].notna().all()
assert silver_df["runs_scored"].between(0, 6).all()
print(f"  Silver: {len(silver_df)} rows (bronze={len(bronze_df)}, rejected={len(bronze_df)-len(silver_df)}) ✓")

# ── Schema evolution: add is_boundary column ─────────────────────────────────
silver_df["is_boundary"] = silver_df["runs_scored"] >= 4
# Old Bronze rows get default False (simulating Iceberg schema evolution)
bronze_df["is_boundary"] = False  # new column with default for existing data
assert "is_boundary" in silver_df.columns
assert "is_boundary" in bronze_df.columns
print(f"  Schema evolution: 'is_boundary' added, old rows default=False ✓")

# ── Gold mart ─────────────────────────────────────────────────────────────────
con.register("silver_deliveries", silver_df)
gold_df = con.execute("""
SELECT
    bowler,
    COUNT(*)                                        AS deliveries,
    ROUND(COUNT(*)/6.0, 1)                          AS overs,
    SUM(runs_scored)                                AS runs_conceded,
    SUM(CASE WHEN is_wicket THEN 1 ELSE 0 END)      AS wickets,
    ROUND(SUM(runs_scored)/NULLIF(COUNT(*)/6.0,0),2) AS economy
FROM silver_deliveries
GROUP BY bowler
ORDER BY economy
""").df()

assert gold_df["economy"].notna().all()
assert int(gold_df["deliveries"].sum()) == len(silver_df)
print(f"  Gold: {len(gold_df)} bowler records, conservation verified ✓")

# ── Resource tagging validation ───────────────────────────────────────────────
RESOURCE_TAGS = {
    "Team": "data-engineering", "Project": "ipl-platform",
    "Environment": "prod", "Pipeline": "medallion-lakehouse",
    "Owner": "[email protected]", "CostCenter": "DE-001", "ManagedBy": "Terraform",
}
REQUIRED_KEYS = {"Team","Project","Environment","Pipeline","Owner","CostCenter","ManagedBy"}
missing_tags = REQUIRED_KEYS - set(RESOURCE_TAGS.keys())
assert missing_tags == set(), f"Missing required tags: {missing_tags}"
print(f"  Tag validation: all {len(REQUIRED_KEYS)} required tags present ✓")

# ── Cost estimation ────────────────────────────────────────────────────────────
ON_DEMAND_COST_HR  = 0.192   # m5.xlarge per hour
SPOT_DISCOUNT       = 0.70
JOB_DURATION_HR     = 0.5    # 30-minute Glue/Spark job
WORKER_COUNT        = 4

cost_on_demand = ON_DEMAND_COST_HR * WORKER_COUNT * JOB_DURATION_HR
cost_spot      = cost_on_demand * (1 - SPOT_DISCOUNT)
print(f"  Cost estimate: on-demand=${cost_on_demand:.3f}, spot=${cost_spot:.3f} per run")
print(f"  Monthly saving (daily): ${(cost_on_demand-cost_spot)*30:.2f}")

print("\n=== Medallion Lakehouse Exercise Complete ===")
print(f"  Bronze rows:  {len(bronze_df)} (raw, immutable)")
print(f"  Silver rows:  {len(silver_df)} (cleaned, typed)")
print(f"  Gold records: {len(gold_df)} bowlers")
print("  All assertions passed ✓")