Kalletlamadhav/sql-optimization-env

1#!/usr/bin/env python32"""3data/fixtures/generate_fixtures.py4-----------------------------------5Thin wrapper around data/seed_database.py that:6  1. Generates the canonical benchmark_seed42.db (SEED=42, 100 000 rows).7  2. Prints a row-count summary so you can confirm the DB is correct.8  3. Optionally generates a larger "scale" database for curriculum Level 4/5.9 10Usage:11    # Standard benchmark DB (used by inference.py and the environment):12    python data/fixtures/generate_fixtures.py13 14    # Large-scale DB for Level 4/5 curriculum (5 M rows):15    python data/fixtures/generate_fixtures.py --scale16 17    # Custom row count:18    python data/fixtures/generate_fixtures.py --rows 50000019"""20 21import argparse22import sqlite323import subprocess24import sys25from pathlib import Path26 27# ── Paths ─────────────────────────────────────────────────────────────────────28FIXTURES_DIR = Path(__file__).parent29STANDARD_DB = FIXTURES_DIR / 'benchmark_seed42.db'30SCALE_DB    = FIXTURES_DIR / 'benchmark_scale_seed42.db'31SEEDER      = Path(__file__).parent.parent / 'seed_database.py'32 33# ── Expected row counts (approximate) after --rows 100000 ────────────────────34EXPECTED_COUNTS = {35    'gst_invoice_records':      100_000,36    'gst_invoice_items':        300_000,   # ~3 items/invoice average37    'ration_card_beneficiaries': 20_000,   # n // 538    'pds_allotments':            80_000,   # ~4 months × 20 k cards39    'railway_trains':                 5,   # Fixed set40    'railway_pnr_bookings':     100_000,41    'mgnrega_workers':           33_333,   # n // 342    'mgnrega_attendance':     2_500_000,   # ~75 days/worker average43    'mgnrega_payments':         400_000,   # 12 months × 33 k workers44}45 46 47def run_seeder(db_path: Path, rows: int) -> None:48    """Invoke seed_database.py via subprocess."""49    cmd = [50        sys.executable,51        str(SEEDER),52        '--rows', str(rows),53        '--db',   str(db_path),54    ]55    print(f'Running: {" ".join(cmd)}')56    result = subprocess.run(cmd, check=True)57    if result.returncode != 0:58        sys.exit(1)59 60 61def print_summary(db_path: Path) -> None:62    """Print row counts for all tables and flag large deviations."""63    conn = sqlite3.connect(db_path)64    tables = [65        row[0] for row in66        conn.execute("SELECT name FROM sqlite_master WHERE type='table' ORDER BY name").fetchall()67    ]68    print(f'\n{"Table":<35} {"Rows":>12}  {"Expected":>12}  {"Status"}')69    print('-' * 72)70    all_ok = True71    for table in tables:72        count = conn.execute(f'SELECT COUNT(*) FROM {table}').fetchone()[0]73        expected = EXPECTED_COUNTS.get(table, '?')74        if isinstance(expected, int):75            ratio  = count / expected if expected else 076            status = '✅' if 0.5 <= ratio <= 2.0 else '⚠️  CHECK'77            if ratio < 0.5 or ratio > 2.0:78                all_ok = False79        else:80            status = '—'81        print(f'{table:<35} {count:>12,}  {str(expected):>12}  {status}')82 83    indexes = conn.execute(84        "SELECT name, tbl_name FROM sqlite_master WHERE type='index' ORDER BY tbl_name, name"85    ).fetchall()86    print(f'\nIndexes present ({len(indexes)}):')87    for idx_name, tbl_name in indexes:88        print(f'  {tbl_name}.{idx_name}')89 90    conn.close()91    print()92    if all_ok:93        print('✅  All row counts within expected range.')94    else:95        print('⚠️   Some row counts are outside expected range — re-run seeder or check SEED.')96 97 98def main() -> None:99    parser = argparse.ArgumentParser(description='Generate benchmark fixture databases.')100    parser.add_argument(101        '--rows', type=int, default=100_000,102        help='Number of GST invoices (base scale). Default: 100 000.'103    )104    parser.add_argument(105        '--scale', action='store_true',106        help='Generate the large-scale DB (5 M rows) for curriculum Level 4/5.'107    )108    parser.add_argument(109        '--skip-generate', action='store_true',110        help='Skip generation — only print summary of existing DB.'111    )112    args = parser.parse_args()113 114    if args.scale:115        db_path = SCALE_DB116        rows    = 5_000_000117        print(f'Generating LARGE-SCALE fixture: {db_path}  ({rows:,} rows) …')118    else:119        db_path = STANDARD_DB120        rows    = args.rows121        print(f'Generating STANDARD fixture:    {db_path}  ({rows:,} rows) …')122 123    if not args.skip_generate:124        run_seeder(db_path, rows)125 126    if db_path.exists():127        print_summary(db_path)128    else:129        print(f'ERROR: DB not found at {db_path}', file=sys.stderr)130        sys.exit(1)131 132 133if __name__ == '__main__':134    main()