wenshu/scripts/generate_iceberg_demo.py

#!/usr/bin/env python3
from __future__ import annotations

import csv
import json
import random
from dataclasses import dataclass
from datetime import date, datetime, timedelta
from pathlib import Path
from typing import Iterable


ROOT = Path(__file__).resolve().parents[1]
OUT_DIR = ROOT / "backend" / "src" / "main" / "resources" / "db" / "iceberg_demo"
RANDOM_SEED = 20260327


@dataclass(frozen=True)
class ProductRow:
    product_id: int
    sku: str
    product_name: str
    category: str
    price: float
    stock_qty: int
    launch_date: str
    active_flag: str


def ensure_output_dir() -> None:
    OUT_DIR.mkdir(parents=True, exist_ok=True)


def date_range(start: date, end: date) -> Iterable[date]:
    current = start
    while current <= end:
        yield current
        current += timedelta(days=1)


def pick_region(rng: random.Random) -> str:
    return rng.choice(["North", "South", "East", "West"])


def noisy_text(rng: random.Random, value: str, probability: float = 0.35) -> str:
    if rng.random() > probability:
        return value
    choice = rng.randint(0, 4)
    if choice == 0:
        return f" {value} "
    if choice == 1:
        return value.upper()
    if choice == 2:
        return value.lower()
    if choice == 3:
        return value.title()
    return f"  {value}  "


def build_customers(rng: random.Random) -> list[dict[str, object]]:
    first_names = [
        "Ava", "Ben", "Cody", "Dora", "Evan", "Fiona", "Gina", "Hank",
        "Iris", "Jack", "Kira", "Liam", "Mia", "Nora", "Owen", "Pia",
        "Quinn", "Rita", "Sean", "Tina", "Uma", "Vera", "Wade", "Xena",
        "Yara", "Zane",
    ]
    last_names = [
        "Adams", "Brown", "Carter", "Diaz", "Evans", "Foster", "Grant", "Hayes",
        "Irwin", "Jones", "King", "Lopez", "Moore", "Nguyen", "Owens", "Price",
        "Quinn", "Reed", "Stone", "Turner", "Usher", "Vance", "Walker", "Xu",
        "Young", "Zimmer",
    ]
    cities = [
        "Shanghai", "Beijing", "Shenzhen", "Hangzhou", "Nanjing", "Chengdu",
        "Wuhan", "Suzhou", "Guangzhou", "Xian", "Tianjin", "Chongqing",
    ]
    tiers = ["bronze", "silver", "gold", "platinum"]
    rows = []
    base_date = date(2024, 1, 1)
    for customer_id in range(1, 1001):
        first = first_names[(customer_id - 1) % len(first_names)]
        last = last_names[((customer_id - 1) * 3) % len(last_names)]
        full_name = noisy_text(rng, f"{first} {last}", probability=0.45)
        email = f"{first}.{last}{customer_id}@example.com".lower()
        if customer_id % 37 == 0:
            email = f" {email.upper()} "
        elif customer_id % 41 == 0:
            email = f"{first.lower()}{customer_id}"
        city = noisy_text(rng, cities[(customer_id - 1) % len(cities)], probability=0.25)
        region = noisy_text(rng, pick_region(rng), probability=0.25)
        signup_date = base_date + timedelta(days=(customer_id * 17) % 780)
        loyalty_tier = noisy_text(rng, tiers[(customer_id - 1) % len(tiers)], probability=0.40)
        status = "ACTIVE" if customer_id % 9 else "inactive"
        if customer_id % 53 == 0:
            status = " suspended "
        age = 18 + (customer_id * 7) % 45
        rows.append({
            "customer_id": customer_id,
            "full_name": full_name,
            "email": email,
            "city": city,
            "region": region,
            "signup_date": signup_date.isoformat(),
            "loyalty_tier": loyalty_tier,
            "status": status,
            "age": age,
        })
    return rows


def build_products(rng: random.Random) -> tuple[list[dict[str, object]], dict[int, ProductRow]]:
    categories = [
        "electronics", "home", "fashion", "sports", "beauty", "books",
    ]
    rows = []
    valid_lookup: dict[int, ProductRow] = {}
    base_date = date(2023, 6, 1)
    for product_id in range(1, 121):
        category = categories[(product_id - 1) % len(categories)]
        category_value = noisy_text(rng, category, probability=0.55)
        price = round(8.0 + (product_id * 3.37) % 180 + rng.uniform(0.25, 19.75), 2)
        stock_qty = 20 + (product_id * 29) % 900
        if product_id % 29 == 0:
            price = -abs(price)
        elif product_id % 37 == 0:
            stock_qty = -stock_qty
        active_flag = "Y" if product_id % 5 else "n"
        if product_id % 11 == 0:
            active_flag = " true "
        sku = noisy_text(rng, f"SKU-{product_id:04d}", probability=0.15)
        product_name = f"{category.title()} Item {product_id:03d}"
        launch_date = base_date + timedelta(days=(product_id * 13) % 560)
        row = {
            "product_id": product_id,
            "sku": sku,
            "product_name": product_name,
            "category": category_value,
            "price": price,
            "stock_qty": stock_qty,
            "launch_date": launch_date.isoformat(),
            "active_flag": active_flag,
        }
        rows.append(row)
        if price > 0 and stock_qty >= 0:
            valid_lookup[product_id] = ProductRow(
                product_id=product_id,
                sku=f"SKU-{product_id:04d}",
                product_name=product_name,
                category=category.title(),
                price=price,
                stock_qty=stock_qty,
                launch_date=launch_date.isoformat(),
                active_flag=active_flag.strip().upper(),
            )
    return rows, valid_lookup


def build_orders(
    rng: random.Random,
    customers: list[dict[str, object]],
    products: dict[int, ProductRow],
) -> list[dict[str, object]]:
    valid_customer_ids = [row["customer_id"] for row in customers if row["customer_id"] % 41 != 0]
    valid_product_ids = list(products.keys())
    statuses = ["PAID", " shipped ", "returned", "CANCELLED", "pending"]
    channels = ["web", "app", "store", "partner"]
    notes = ["", "priority", "gift", "repeat_customer", "bulk_order"]
    start = date(2025, 1, 1)
    end = date(2026, 3, 20)
    total_days = (end - start).days
    rows = []
    for offset in range(1, 10001):
        customer_id = rng.choice(valid_customer_ids)
        product_id = rng.choice(valid_product_ids)
        product = products[product_id]
        order_date = start + timedelta(days=(offset * 19 + product_id) % total_days)
        quantity = 1 + (offset * 7) % 8
        if offset % 97 == 0:
            quantity = 0
        unit_price = round(product.price * (0.92 + ((offset % 13) * 0.01)), 2)
        if offset % 151 == 0:
            unit_price = -unit_price
        discount_rate = round(((offset % 9) * 0.04), 2)
        if offset % 53 == 0:
            discount_rate = ""
        status = statuses[offset % len(statuses)]
        channel = channels[(offset + product_id) % len(channels)]
        region = noisy_text(rng, pick_region(rng), probability=0.22)
        note = notes[offset % len(notes)]
        rows.append({
            "order_id": 7000000 + offset,
            "customer_id": customer_id,
            "product_id": product_id,
            "order_date": order_date.isoformat(),
            "quantity": quantity,
            "unit_price": unit_price,
            "discount_rate": discount_rate,
            "status": status,
            "channel": channel,
            "region": region,
            "note": note,
        })
    return rows


def write_csv(path: Path, rows: list[dict[str, object]], header: list[str]) -> None:
    with path.open("w", newline="", encoding="utf-8") as fp:
        writer = csv.DictWriter(fp, fieldnames=header)
        writer.writeheader()
        writer.writerows(rows)


def sql_escape(value: str) -> str:
    return value.replace("'", "''")


def build_workflow_sql() -> str:
    dag_json = {
        "nodes": [
            {"taskId": 964011},
            {"taskId": 964012},
            {"taskId": 964013},
            {"taskId": 964014},
            {"taskId": 964015},
        ],
        "edges": [
            {"fromTaskId": 964011, "toTaskId": 964014},
            {"fromTaskId": 964012, "toTaskId": 964014},
            {"fromTaskId": 964013, "toTaskId": 964014},
            {"fromTaskId": 964014, "toTaskId": 964015},
        ],
    }
    dag_json_text = json.dumps(dag_json, ensure_ascii=False, separators=(",", ":"))

    customer_clean_sql = (
        "CREATE OR REPLACE TABLE polaris.demo_iceberg.customer_clean USING iceberg AS "
        "SELECT customer_id, "
        "trim(regexp_replace(full_name, '\\\\s+', ' ')) AS full_name, "
        "lower(trim(email)) AS email, "
        "initcap(trim(city)) AS city, "
        "upper(trim(region)) AS region, "
        "to_date(signup_date) AS signup_date, "
        "upper(trim(loyalty_tier)) AS loyalty_tier, "
        "CASE WHEN lower(trim(status)) IN ('active','1','y') THEN 'ACTIVE' ELSE 'INACTIVE' END AS status, "
        "CAST(age AS INT) AS age "
        "FROM polaris.demo_iceberg.demo_customers "
        "WHERE customer_id IS NOT NULL AND trim(email) <> '' AND email LIKE '%@%';"
    )

    product_clean_sql = (
        "CREATE OR REPLACE TABLE polaris.demo_iceberg.product_clean USING iceberg AS "
        "SELECT product_id, "
        "upper(trim(sku)) AS sku, "
        "trim(product_name) AS product_name, "
        "initcap(lower(trim(category))) AS category, "
        "CAST(price AS DECIMAL(12,2)) AS price, "
        "CAST(stock_qty AS INT) AS stock_qty, "
        "to_date(launch_date) AS launch_date, "
        "CASE WHEN lower(trim(active_flag)) IN ('y','1','true') THEN true ELSE false END AS is_active "
        "FROM polaris.demo_iceberg.demo_products "
        "WHERE product_id IS NOT NULL AND price > 0 AND stock_qty >= 0;"
    )

    order_clean_sql = (
        "CREATE OR REPLACE TABLE polaris.demo_iceberg.order_clean USING iceberg AS "
        "SELECT order_id, "
        "customer_id, "
        "product_id, "
        "to_date(order_date) AS order_date, "
        "CAST(quantity AS INT) AS quantity, "
        "CAST(unit_price AS DECIMAL(12,2)) AS unit_price, "
        "CAST(COALESCE(discount_rate, 0) AS DECIMAL(5,4)) AS discount_rate, "
        "upper(trim(status)) AS status, "
        "upper(trim(channel)) AS channel, "
        "upper(trim(region)) AS region, "
        "NULLIF(trim(note), '') AS note "
        "FROM polaris.demo_iceberg.demo_orders "
        "WHERE order_id IS NOT NULL AND customer_id IS NOT NULL AND product_id IS NOT NULL "
        "AND quantity > 0 AND unit_price > 0;"
    )

    enriched_sql = (
        "CREATE OR REPLACE TABLE polaris.demo_iceberg.order_enriched USING iceberg AS "
        "SELECT o.order_id, "
        "o.order_date, "
        "o.customer_id, "
        "c.full_name, "
        "c.city, "
        "c.region AS customer_region, "
        "p.category, "
        "p.sku, "
        "o.channel, "
        "o.status, "
        "o.quantity, "
        "o.unit_price, "
        "o.discount_rate, "
        "CAST(o.quantity * o.unit_price AS DECIMAL(14,2)) AS gross_amount, "
        "CAST(o.quantity * o.unit_price * (1 - COALESCE(o.discount_rate, 0)) AS DECIMAL(14,2)) AS net_amount "
        "FROM polaris.demo_iceberg.order_clean o "
        "JOIN polaris.demo_iceberg.customer_clean c ON o.customer_id = c.customer_id "
        "JOIN polaris.demo_iceberg.product_clean p ON o.product_id = p.product_id;"
    )

    summary_sql = (
        "CREATE OR REPLACE TABLE polaris.demo_iceberg.sales_kpi_daily USING iceberg AS "
        "SELECT order_date, "
        "customer_region, "
        "category, "
        "COUNT(*) AS order_count, "
        "COUNT(DISTINCT customer_id) AS customer_count, "
        "SUM(quantity) AS total_quantity, "
        "ROUND(SUM(net_amount), 2) AS net_revenue, "
        "ROUND(AVG(net_amount), 2) AS avg_net_order_value, "
        "ROUND(MAX(net_amount), 2) AS max_net_order_value "
        "FROM polaris.demo_iceberg.order_enriched "
        "GROUP BY order_date, customer_region, category;"
    )

    workflow_name = "Iceberg_Sales_Clean_Analysis"
    description = "Import three related Iceberg tables, clean dirty records, and generate sales KPIs"
    now = "2026-03-27 10:00:00"

    dag_json_sql = sql_escape(dag_json_text)
    return f"""USE wenshu_platform;

INSERT INTO workflow_definition (
  workflow_id,
  workflow_name,
  description,
  dag_json,
  timeout_seconds,
  create_time,
  failure_strategy
) VALUES
  (
    964001,
    '{sql_escape(workflow_name)}',
    '{sql_escape(description)}',
    '{dag_json_sql}',
    9000,
    '{now}',
    'STOP'
  )
ON DUPLICATE KEY UPDATE
  workflow_name = VALUES(workflow_name),
  description = VALUES(description),
  dag_json = VALUES(dag_json),
  timeout_seconds = VALUES(timeout_seconds),
  create_time = VALUES(create_time),
  failure_strategy = VALUES(failure_strategy);

INSERT INTO task_definition (
  task_id,
  task_name,
  workflow_id,
  task_type,
  task_content,
  exector_id,
  timeout_seconds,
  retry_times
) VALUES
  (964011, '清洗客户表', 964001, 'SPARK_SQL', '{sql_escape(customer_clean_sql)}', 3001, 1200, 1),
  (964012, '清洗产品表', 964001, 'SPARK_SQL', '{sql_escape(product_clean_sql)}', 3001, 1200, 1),
  (964013, '清洗订单表', 964001, 'SPARK_SQL', '{sql_escape(order_clean_sql)}', 3001, 1500, 1),
  (964014, '构建明细宽表', 964001, 'SPARK_SQL', '{sql_escape(enriched_sql)}', 3001, 1800, 1),
  (964015, '生成日粒度指标', 964001, 'SPARK_SQL', '{sql_escape(summary_sql)}', 3001, 1200, 1)
ON DUPLICATE KEY UPDATE
  task_name = VALUES(task_name),
  workflow_id = VALUES(workflow_id),
  task_type = VALUES(task_type),
  task_content = VALUES(task_content),
  exector_id = VALUES(exector_id),
  timeout_seconds = VALUES(timeout_seconds),
  retry_times = VALUES(retry_times);

INSERT INTO workflow_instance (
  workflow_instance_id,
  workflow_id,
  workflow_name,
  start_time,
  end_time,
  state,
  error_message
) VALUES
  (965001, 964001, 'Iceberg_Sales_Clean_Analysis', '2026-03-27 10:05:00', '2026-03-27 10:12:48', 'SUCCESS', NULL)
ON DUPLICATE KEY UPDATE
  workflow_id = VALUES(workflow_id),
  workflow_name = VALUES(workflow_name),
  start_time = VALUES(start_time),
  end_time = VALUES(end_time),
  state = VALUES(state),
  error_message = VALUES(error_message);

INSERT INTO task_instance (
  task_instance_id,
  workflow_instance_id,
  task_id,
  task_name,
  execution_script,
  state,
  start_time,
  end_time,
  executor_id,
  retry_count,
  engine_task_id
) VALUES
  (965011, 965001, 964011, '清洗客户表', '{sql_escape(customer_clean_sql)}', 'SUCCESS', '2026-03-27 10:05:00', '2026-03-27 10:06:05', 3001, 0, 'spark-demo-clean-customers'),
  (965012, 965001, 964012, '清洗产品表', '{sql_escape(product_clean_sql)}', 'SUCCESS', '2026-03-27 10:05:00', '2026-03-27 10:06:10', 3001, 0, 'spark-demo-clean-products'),
  (965013, 965001, 964013, '清洗订单表', '{sql_escape(order_clean_sql)}', 'SUCCESS', '2026-03-27 10:05:00', '2026-03-27 10:07:35', 3001, 0, 'spark-demo-clean-orders'),
  (965014, 965001, 964014, '构建明细宽表', '{sql_escape(enriched_sql)}', 'SUCCESS', '2026-03-27 10:07:35', '2026-03-27 10:09:40', 3001, 0, 'spark-demo-build-fact'),
  (965015, 965001, 964015, '生成日粒度指标', '{sql_escape(summary_sql)}', 'SUCCESS', '2026-03-27 10:09:40', '2026-03-27 10:12:48', 3001, 0, 'spark-demo-build-kpi')
ON DUPLICATE KEY UPDATE
  workflow_instance_id = VALUES(workflow_instance_id),
  task_id = VALUES(task_id),
  task_name = VALUES(task_name),
  execution_script = VALUES(execution_script),
  state = VALUES(state),
  start_time = VALUES(start_time),
  end_time = VALUES(end_time),
  executor_id = VALUES(executor_id),
  retry_count = VALUES(retry_count),
  engine_task_id = VALUES(engine_task_id);
"""


def main() -> None:
    rng = random.Random(RANDOM_SEED)
    ensure_output_dir()

    customers = build_customers(rng)
    products, valid_products = build_products(rng)
    orders = build_orders(rng, customers, valid_products)

    write_csv(
        OUT_DIR / "demo_customers.csv",
        customers,
        [
            "customer_id",
            "full_name",
            "email",
            "city",
            "region",
            "signup_date",
            "loyalty_tier",
            "status",
            "age",
        ],
    )
    write_csv(
        OUT_DIR / "demo_products.csv",
        products,
        [
            "product_id",
            "sku",
            "product_name",
            "category",
            "price",
            "stock_qty",
            "launch_date",
            "active_flag",
        ],
    )
    write_csv(
        OUT_DIR / "demo_orders.csv",
        orders,
        [
            "order_id",
            "customer_id",
            "product_id",
            "order_date",
            "quantity",
            "unit_price",
            "discount_rate",
            "status",
            "channel",
            "region",
            "note",
        ],
    )
    (OUT_DIR / "test-iceberg-demo-data.sql").write_text(build_workflow_sql(), encoding="utf-8")

    print("Generated demo Iceberg data:")
    print(f"  {OUT_DIR / 'demo_customers.csv'}")
    print(f"  {OUT_DIR / 'demo_products.csv'}")
    print(f"  {OUT_DIR / 'demo_orders.csv'}")
    print(f"  {OUT_DIR / 'test-iceberg-demo-data.sql'}")
    print(f"Rows: customers={len(customers)}, products={len(products)}, orders={len(orders)}")


if __name__ == "__main__":
    main()