scripts/generate_benchmark.py

#!/usr/bin/env python3
"""Generate a small synthetic 'Squirrel-mini' benchmark.

This is a toy proxy for the Squirrel/SQLBench construction pipeline (Section 3 of
the paper). The real seed corpus (1,000+ proprietary enterprise SQL scripts) is not
public, so this script synthesizes long, multi-CTE ETL-style SQL scripts across
several business domains, then injects syntax and semantic bugs following the
paper's bug taxonomy (Table 3/4), producing (issue_sql, reference_sql, error/query)
triples we can use to probe LLM debugging ability and to measure script-complexity
statistics (lines, tokens, functions, AST width/depth) for comparison with Table 1.
"""
from __future__ import annotations

import json
import re
from pathlib import Path

DOMAINS = [
    ("ecommerce", "order_events", "customer_profile", "product_catalog"),
    ("finance", "transaction_ledger", "account_profile", "instrument_catalog"),
    ("healthcare", "visit_events", "patient_profile", "procedure_catalog"),
    ("logistics", "shipment_events", "carrier_profile", "route_catalog"),
    ("education", "enrollment_events", "student_profile", "course_catalog"),
    ("telecom", "usage_events", "subscriber_profile", "plan_catalog"),
    ("real_estate", "listing_events", "agent_profile", "property_catalog"),
    ("entertainment", "stream_events", "viewer_profile", "title_catalog"),
    ("travel", "booking_events", "traveler_profile", "fare_catalog"),
    ("energy", "meter_events", "site_profile", "tariff_catalog"),
]

N_METRIC_COLS = 14
N_DIM_COLS = 10


def build_reference_sql(domain: str, fact: str, dim1: str, dim2: str) -> str:
    """Build a long multi-CTE enterprise ETL SQL script (Hive/Spark-style)."""
    metric_cols = [f"metric_{i}_amount" for i in range(N_METRIC_COLS)]
    dim_cols = [f"dim_{i}_code" for i in range(N_DIM_COLS)]

    lines = []
    lines.append(f"-- ETL pipeline for {domain} daily aggregation")
    lines.append("WITH base_events AS (")
    lines.append("    SELECT")
    lines.append(f"        {fact}.entity_id,")
    lines.append(f"        {fact}.event_timestamp,")
    for c in dim_cols[:5]:
        lines.append(f"        {fact}.{c},")
    for c in metric_cols[:7]:
        lines.append(f"        {fact}.{c},")
    lines.append(f"        {fact}.status_flag")
    lines.append(f"    FROM {domain}_db.{fact} {fact}")
    lines.append(f"    WHERE {fact}.event_date BETWEEN '${{start_date}}' AND '${{end_date}}'")
    lines.append("      AND status_flag = 'ACTIVE'")
    lines.append("),")

    lines.append("enriched_events AS (")
    lines.append("    SELECT")
    lines.append("        base_events.entity_id,")
    lines.append("        base_events.event_timestamp,")
    for c in dim_cols[:5]:
        lines.append(f"        base_events.{c},")
    for c in metric_cols[:7]:
        lines.append(
            f"        COALESCE(base_events.{c}, 0) AS {c},"
        )
    lines.append(f"        {dim1}.segment_code,")
    lines.append(f"        {dim1}.region_code,")
    lines.append(f"        {dim1}.tier_level,")
    lines.append(f"        {dim2}.category_code,")
    lines.append(f"        {dim2}.subcategory_code")
    lines.append("    FROM base_events")
    lines.append(f"    JOIN {domain}_db.{dim1} {dim1}")
    lines.append(f"        ON base_events.entity_id = {dim1}.entity_id")
    lines.append(f"    LEFT JOIN {domain}_db.{dim2} {dim2}")
    lines.append(f"        ON base_events.dim_0_code = {dim2}.dim_0_code")
    lines.append("),")

    lines.append("windowed_metrics AS (")
    lines.append("    SELECT")
    lines.append("        entity_id,")
    lines.append("        segment_code,")
    lines.append("        region_code,")
    lines.append("        category_code,")
    for c in metric_cols[:7]:
        lines.append(f"        {c},")
    lines.append(
        "        ROW_NUMBER() OVER (PARTITION BY entity_id ORDER BY event_timestamp DESC) AS recency_rank,"
    )
    lines.append(
        "        SUM(metric_0_amount) OVER (PARTITION BY entity_id, segment_code) AS segment_total,"
    )
    lines.append(
        "        AVG(metric_1_amount) OVER (PARTITION BY region_code ORDER BY event_timestamp "
        "ROWS BETWEEN 6 PRECEDING AND CURRENT ROW) AS rolling_avg_metric_1"
    )
    lines.append("    FROM enriched_events")
    lines.append("),")

    lines.append("aggregated AS (")
    lines.append("    SELECT")
    lines.append("        segment_code,")
    lines.append("        region_code,")
    lines.append("        category_code,")
    lines.append("        COUNT(DISTINCT entity_id) AS entity_count,")
    for i, c in enumerate(metric_cols[:7]):
        lines.append(f"        SUM({c}) AS total_{c},")
        lines.append(f"        MAX(CASE WHEN recency_rank = 1 THEN {c} ELSE NULL END) AS latest_{c},")
    lines.append("        MAX(segment_total) AS segment_total,")
    lines.append("        AVG(rolling_avg_metric_1) AS avg_rolling_metric_1")
    lines.append("    FROM windowed_metrics")
    lines.append("    GROUP BY")
    lines.append("        segment_code,")
    lines.append("        region_code,")
    lines.append("        category_code")
    lines.append("),")

    lines.append("classified AS (")
    lines.append("    SELECT")
    lines.append("        segment_code,")
    lines.append("        region_code,")
    lines.append("        category_code,")
    lines.append("        entity_count,")
    lines.append("        total_metric_0_amount,")
    lines.append("        total_metric_1_amount,")
    lines.append("        segment_total,")
    lines.append("        avg_rolling_metric_1,")
    lines.append("        CASE")
    lines.append("            WHEN total_metric_0_amount > 100000 THEN 'HIGH'")
    lines.append("            WHEN total_metric_0_amount > 10000 THEN 'MEDIUM'")
    lines.append("            ELSE 'LOW'")
    lines.append("        END AS value_tier,")
    lines.append(
        "        CONCAT_WS('_', segment_code, region_code) AS segment_region_key,"
    )
    lines.append(
        "        ROUND(total_metric_1_amount / NULLIF(entity_count, 0), 2) AS per_entity_metric_1"
    )
    lines.append("    FROM aggregated")
    lines.append("),")

    lines.append("legacy_snapshot AS (")
    lines.append("    SELECT")
    lines.append("        segment_code,")
    lines.append("        region_code,")
    lines.append("        category_code,")
    lines.append("        entity_count AS entity_count,")
    lines.append("        total_metric_0_amount AS total_metric_0_amount,")
    lines.append("        'legacy' AS data_source")
    lines.append(f"    FROM {domain}_db.legacy_{fact}_snapshot")
    lines.append("    WHERE snapshot_date = '${end_date}'")
    lines.append(")")

    lines.append("INSERT OVERWRITE TABLE fake_base_test.{}_daily_summary".format(domain))
    lines.append("PARTITION (processing_date = '${end_date}')")
    lines.append("SELECT")
    lines.append("    classified.segment_code,")
    lines.append("    classified.region_code,")
    lines.append("    classified.category_code,")
    lines.append("    classified.entity_count,")
    lines.append("    classified.total_metric_0_amount,")
    lines.append("    classified.total_metric_1_amount,")
    lines.append("    classified.segment_total,")
    lines.append("    classified.avg_rolling_metric_1,")
    lines.append("    classified.value_tier,")
    lines.append("    classified.segment_region_key,")
    lines.append("    classified.per_entity_metric_1,")
    lines.append("    'current' AS data_source")
    lines.append("FROM classified")
    lines.append("UNION ALL")
    lines.append("SELECT")
    lines.append("    legacy_snapshot.segment_code,")
    lines.append("    legacy_snapshot.region_code,")
    lines.append("    legacy_snapshot.category_code,")
    lines.append("    legacy_snapshot.entity_count,")
    lines.append("    legacy_snapshot.total_metric_0_amount,")
    lines.append("    0 AS total_metric_1_amount,")
    lines.append("    0 AS segment_total,")
    lines.append("    0 AS avg_rolling_metric_1,")
    lines.append("    'LEGACY' AS value_tier,")
    lines.append(
        "    CONCAT_WS('_', legacy_snapshot.segment_code, legacy_snapshot.region_code) AS segment_region_key,"
    )
    lines.append("    0 AS per_entity_metric_1,")
    lines.append("    legacy_snapshot.data_source")
    lines.append("FROM legacy_snapshot;")

    return "\n".join(lines)


def build_ddl(domain: str, fact: str, dim1: str, dim2: str) -> str:
    return (
        f"CREATE TABLE IF NOT EXISTS {domain}_db.{fact} (entity_id STRING, event_timestamp TIMESTAMP, "
        + ", ".join(f"dim_{i}_code STRING" for i in range(N_DIM_COLS))
        + ", "
        + ", ".join(f"metric_{i}_amount DOUBLE" for i in range(N_METRIC_COLS))
        + ", status_flag STRING, event_date STRING);\n"
        f"CREATE TABLE IF NOT EXISTS {domain}_db.{dim1} (entity_id STRING, segment_code STRING, "
        "region_code STRING, tier_level STRING);\n"
        f"CREATE TABLE IF NOT EXISTS {domain}_db.{dim2} (dim_0_code STRING, category_code STRING, "
        "subcategory_code STRING);\n"
        f"CREATE TABLE IF NOT EXISTS {domain}_db.legacy_{fact}_snapshot (segment_code STRING, "
        "region_code STRING, category_code STRING, entity_count BIGINT, total_metric_0_amount DOUBLE, "
        "snapshot_date STRING);"
    )


# ---- Bug injectors (mirroring Table 3 / Table 4 taxonomy) ----

def inject_syntax_bug(sql: str):
    """Return (issue_sql, level1, level2, level3, error_message) or None."""
    # CASE Expression: Missing END (most frequent syntax bug type in Table 4, count=72)
    m = list(re.finditer(r"CASE\n(.*?)\n(\s*)END AS (\w+)", sql, re.S))
    if m:
        m = m[0]
        issue = sql[: m.start()] + f"CASE\n{m.group(1)}\n{m.group(2)}{m.group(3)}" + sql[m.end():]
        err = (
            f"ParseException: mismatched input 'AS' expecting 'END' near '{m.group(3)}'"
        )
        return issue, "Grammar & Structure", "CASE Expression", "Missing END in CASE WHEN", err
    return None


def inject_syntax_bug_group_by(sql: str):
    # Grammar & Structure / Clause Structure: trailing comma before FROM in GROUP BY
    m = re.search(r"(GROUP BY\n\s+segment_code,\n\s+region_code,\n\s+category_code)\n(\))", sql)
    if m:
        issue = sql[: m.start()] + m.group(1) + ",\n" + m.group(2) + sql[m.end():]
        err = "ParseException: extraneous input ')' expecting {identifier} near GROUP BY clause"
        return issue, "Grammar & Structure", "Clause Structure", "Trailing comma in GROUP BY", err
    return None


def inject_semantic_bug_join_type(sql: str):
    # Joins & Grouping / JOIN Type Selection: LEFT JOIN -> INNER JOIN drops unmatched rows
    if "LEFT JOIN" in sql:
        issue = sql.replace("LEFT JOIN", "JOIN", 1)
        query = (
            "Some entities that should appear in the daily summary (e.g. those without a "
            "matching category record) are missing from the output entirely. Expected: all "
            "entities from the base events should be retained even if category enrichment is "
            "unavailable. Please fix the bug."
        )
        return issue, "Joins & Grouping", "JOIN Type Selection", "Using INNER JOIN when LEFT JOIN is needed", query
    return None


def inject_semantic_bug_union(sql: str):
    # Semantics & Logic / Set Operations: UNION ALL -> UNION causes silent row loss via dedup
    if "UNION ALL" in sql:
        issue = sql.replace("UNION ALL", "UNION", 1)
        query = (
            "The daily summary table sometimes has fewer rows than expected when the current "
            "period and legacy snapshot happen to produce identical rows for a segment; those "
            "duplicate-looking rows are actually distinct facts and should both be kept. Please fix the bug."
        )
        return issue, "Semantics & Logic", "Set Operations", "UNION vs UNION ALL misuse", query
    return None


def inject_semantic_bug_groupby_missing(sql: str):
    # Joins & Grouping / GROUP BY Logic: aggregate SELECT includes non-grouped info via wrong join key
    m = re.search(r"ON base_events\.entity_id = " + r"(\w+)\.entity_id", sql)
    if m:
        dim1 = m.group(1)
        issue = sql.replace(
            f"ON base_events.entity_id = {dim1}.entity_id",
            f"ON base_events.dim_0_code = {dim1}.entity_id",
            1,
        )
        query = (
            "The join between base events and the profile dimension looks wrong: segment_code "
            "and region_code values in the output don't correspond to the correct entity anymore, "
            "producing mismatched enrichment. Please fix the bug so profile attributes are joined "
            "on the correct key."
        )
        return issue, "Joins & Grouping", "JOIN Logic", "Wrong join key used", query
    return None


def _parses_ok(sql: str) -> bool:
    import sqlglot

    try:
        sqlglot.parse(sql, read="hive")
        return True
    except Exception:
        return False


SYNTAX_INJECTORS = [inject_syntax_bug, inject_syntax_bug_group_by]
SEMANTIC_INJECTORS = [
    inject_semantic_bug_join_type,
    inject_semantic_bug_union,
    inject_semantic_bug_groupby_missing,
]


def main():
    tasks = []
    seeds = []
    for i, (domain, fact, dim1, dim2) in enumerate(DOMAINS):
        ref_sql = build_reference_sql(domain, fact, dim1, dim2)
        ddl = build_ddl(domain, fact, dim1, dim2)
        seeds.append({"domain": domain, "reference_sql": ref_sql, "ddl": ddl})

        result = None
        for syn_fn in SYNTAX_INJECTORS[i % 2 :] + SYNTAX_INJECTORS[: i % 2]:
            candidate = syn_fn(ref_sql)
            if candidate and not _parses_ok(candidate[0]):
                result = candidate
                break
        if result is None:
            result = inject_syntax_bug(ref_sql)
        if result:
            issue_sql, l1, l2, l3, err = result
            tasks.append(
                {
                    "task_id": f"{domain}-syntax-{i}",
                    "type": "syntax",
                    "domain": domain,
                    "ddl": ddl,
                    "issue_sql": issue_sql,
                    "reference_sql": ref_sql,
                    "error_message": err,
                    "level1_error_type": l1,
                    "level2_error_type": l2,
                    "level3_error_type": l3,
                }
            )

        sem_fn = SEMANTIC_INJECTORS[i % len(SEMANTIC_INJECTORS)]
        result = sem_fn(ref_sql)
        if result:
            issue_sql, l1, l2, l3, query = result
            tasks.append(
                {
                    "task_id": f"{domain}-semantic-{i}",
                    "type": "semantic",
                    "domain": domain,
                    "ddl": ddl,
                    "issue_sql": issue_sql,
                    "reference_sql": ref_sql,
                    "user_query": query,
                    "level1_error_type": l1,
                    "level2_error_type": l2,
                    "level3_error_type": l3,
                }
            )

    out_dir = Path(__file__).resolve().parent.parent / "data"
    out_dir.mkdir(exist_ok=True)
    (out_dir / "squirrel_mini_seeds.json").write_text(json.dumps(seeds, indent=2))
    (out_dir / "squirrel_mini_tasks.json").write_text(json.dumps(tasks, indent=2))
    print(f"Generated {len(seeds)} seed scripts and {len(tasks)} tasks")
    print(f"  syntax tasks: {sum(1 for t in tasks if t['type'] == 'syntax')}")
    print(f"  semantic tasks: {sum(1 for t in tasks if t['type'] == 'semantic')}")


if __name__ == "__main__":
    main()