Kalletlamadhav/sql-optimization-env

1# tasks/hard/gst_multi_join.py2from pathlib import Path3from tasks.base_task import BaseTask4 5# FIX: Added missing import + schema_ddl + hint + reference_fix (all absent in PDF)6_schema_file = Path('data/schemas/gst_schema.sql')7_schema_ddl = _schema_file.read_text() if _schema_file.exists() else (8    "CREATE TABLE gst_invoice_records (invoice_id TEXT PRIMARY KEY, "9    "gstin_supplier TEXT NOT NULL, gstin_buyer TEXT NOT NULL, "10    "invoice_date DATE NOT NULL, invoice_type TEXT DEFAULT 'B2B', "11    "taxable_value DECIMAL(15,2), cgst_amount DECIMAL(10,2), "12    "sgst_amount DECIMAL(10,2), igst_amount DECIMAL(10,2), "13    "cess_amount DECIMAL(10,2) DEFAULT 0.0, state_code CHAR(2) NOT NULL, "14    "hsn_code TEXT, filing_status TEXT DEFAULT 'FILED', "15    "created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP);"16    "CREATE TABLE gst_invoice_items ("17    "item_id INTEGER PRIMARY KEY AUTOINCREMENT, "18    "invoice_id TEXT NOT NULL REFERENCES gst_invoice_records(invoice_id), "19    "item_description TEXT, hsn_code TEXT, quantity DECIMAL(10,3), "20    "unit_value DECIMAL(12,2), taxable_value DECIMAL(15,2), "21    "gst_rate DECIMAL(5,2));"22)23 24TASK = BaseTask(25    task_id='gst_multi_join',26    goal=(27        'A tax auditor needs a reconciliation report: for each supplier in Maharashtra, '28        'show total invoices, total taxable value, total GST paid, and count of line items. '29        'The query uses a correlated subquery with a nested IN clause across 2M rows — '30        'a classic N+1 combined with a missing index. '31        'Optimize for <2 second execution. Multiple anti-patterns are present.'32    ),33    slow_query="""34        SELECT35            r.gstin_supplier,36            COUNT(DISTINCT r.invoice_id) AS invoice_count,37            SUM(r.taxable_value) AS total_taxable,38            SUM(r.cgst_amount + r.sgst_amount + r.igst_amount) AS total_gst,39            (SELECT COUNT(*) FROM gst_invoice_items40             WHERE invoice_id IN (41                 SELECT invoice_id FROM gst_invoice_records42                 WHERE gstin_supplier = r.gstin_supplier43             )) AS item_count44        FROM gst_invoice_records r45        WHERE r.state_code = '27'46        GROUP BY r.gstin_supplier47    """,48    expected_pattern='N_PLUS_ONE',   # Primary — also has MISSING_INDEX49    tables=['gst_invoice_records', 'gst_invoice_items'],50    schema_ddl=_schema_ddl,51    difficulty='hard',52    curriculum_level=4,53    max_steps=5,54    hint='',   # No hint at level 455    reference_fix="""56        CREATE INDEX idx_gst_state_supplier57        ON gst_invoice_records(state_code, gstin_supplier);58 59        CREATE INDEX idx_items_invoice60        ON gst_invoice_items(invoice_id);61 62        SELECT63            r.gstin_supplier,64            COUNT(DISTINCT r.invoice_id) AS invoice_count,65            SUM(r.taxable_value) AS total_taxable,66            SUM(r.cgst_amount + r.sgst_amount + r.igst_amount) AS total_gst,67            COUNT(it.item_id) AS item_count68        FROM gst_invoice_records r69        LEFT JOIN gst_invoice_items it ON it.invoice_id = r.invoice_id70        WHERE r.state_code = '27'71        GROUP BY r.gstin_supplier72    """73)