Upload graders.py

server/graders.py

@@ -1250,3 +1250,148 @@ def grade_dedup_bug(
         },
         partial_scores=partial_scores.copy(),
     )
+
+
+def grade_fk_integrity(
+    action: SQLAction,
+    ground_truth: Dict[str, Any],
+    conn: sqlite3.Connection,
+) -> SQLReward:
+    """Single-step fk_integrity scorer — validates grader exists."""
+    fix_sql = (action.fix_query or action.sql or "").strip()
+    target_table = ground_truth.get("table", "invoices")
+    parent_table = ground_truth.get("parent_table", "customers")
+    parent_key = ground_truth.get("parent_key", "id")
+    fk_col = ground_truth.get("column", "customer_id")
+    snapshot = _snapshot_table(conn, target_table)
+
+    if not fix_sql or fix_sql.upper().startswith("SELECT"):
+        return SQLReward(
+            score=_clamp(0.0),
+            step_reward=_clamp(0.0),
+            feedback="No fix query provided",
+            dimensions={"legacy_mode": True},
+            partial_scores={},
+        )
+
+    test_conn = _fresh_test_db(conn)
+    try:
+        test_conn.executescript(fix_sql)
+        test_conn.commit()
+        orphan_count = test_conn.execute(
+            f"SELECT COUNT(*) FROM {quote_identifier(target_table)} t "
+            f"LEFT JOIN {quote_identifier(parent_table)} p ON t.{quote_identifier(fk_col)} = p.{quote_identifier(parent_key)} "
+            f"WHERE p.{quote_identifier(parent_key)} IS NULL"
+        ).fetchone()[0]
+        issue_resolved = orphan_count == 0
+        preservation = _measure_data_preservation(
+            test_conn, snapshot, target_table
+        )
+        score = 0.0
+        if issue_resolved:
+            score += 0.6
+        if preservation >= 0.95:
+            score += 0.4
+        return SQLReward(
+            score=_clamp(score),
+            step_reward=_clamp(score),
+            feedback=(
+                "FK integrity resolved" if issue_resolved
+                else f"Orphans remaining: {orphan_count}"
+            ),
+            dimensions={"legacy_mode": True, "composite": _clamp(score)},
+            partial_scores={"issue_resolved": float(issue_resolved),
+                           "preservation": preservation},
+        )
+    except sqlite3.Error as exc:
+        return SQLReward(
+            score=_clamp(0.0),
+            step_reward=_clamp(0.0),
+            feedback=f"SQL error: {exc}",
+            dimensions={"legacy_mode": True},
+            partial_scores={},
+        )
+    finally:
+        test_conn.close()
+
+
+def grade_schema_drift(
+    action: SQLAction,
+    ground_truth: Dict[str, Any],
+    conn: sqlite3.Connection,
+) -> SQLReward:
+    """Single-step schema_drift scorer — validates grader exists."""
+    fix_sql = (action.fix_query or action.sql or "").strip()
+    canonical_table = ground_truth.get("table", "canonical_events")
+    expected_cols = ground_truth.get(
+        "expected_columns",
+        ["event_id", "user_id", "event_type", "amount", "event_time"]
+    )
+    expected_rows = ground_truth.get("expected_row_count", 1200)
+
+    if not fix_sql:
+        return SQLReward(
+            score=_clamp(0.0),
+            step_reward=_clamp(0.0),
+            feedback="No fix query provided",
+            dimensions={"legacy_mode": True},
+            partial_scores={},
+        )
+
+    test_conn = _fresh_test_db(conn)
+    try:
+        test_conn.executescript(fix_sql)
+        test_conn.commit()
+        tables = {
+            r[0] for r in test_conn.execute(
+                "SELECT name FROM sqlite_master WHERE type='table'"
+            ).fetchall()
+        }
+        if canonical_table not in tables:
+            return SQLReward(
+                score=_clamp(0.0),
+                step_reward=_clamp(0.0),
+                feedback=f"Table {canonical_table} not created",
+                dimensions={"legacy_mode": True},
+                partial_scores={},
+            )
+        cols = {
+            r[1] for r in test_conn.execute(
+                f"PRAGMA table_info({quote_identifier(canonical_table)})"
+            )
+        }
+        row_count = test_conn.execute(
+            f"SELECT COUNT(*) FROM {quote_identifier(canonical_table)}"
+        ).fetchone()[0]
+        col_overlap = len(set(expected_cols) & cols) / max(1, len(expected_cols))
+        row_align = 1.0 - min(
+            1.0, abs(row_count - expected_rows) / max(1, expected_rows)
+        )
+        score = col_overlap * 0.5 + row_align * 0.5
+        return SQLReward(
+            score=_clamp(score),
+            step_reward=_clamp(score),
+            feedback=f"Schema drift: col_overlap={col_overlap:.2f} row_align={row_align:.2f}",
+            dimensions={"legacy_mode": True, "composite": _clamp(score)},
+            partial_scores={"col_overlap": col_overlap, "row_align": row_align},
+        )
+    except sqlite3.Error as exc:
+        return SQLReward(
+            score=_clamp(0.0),
+            step_reward=_clamp(0.0),
+            feedback=f"SQL error: {exc}",
+            dimensions={"legacy_mode": True},
+            partial_scores={},
+        )
+    finally:
+        test_conn.close()
+
+
+__all__ = [
+    "grade_null_audit",
+    "grade_type_repair",
+    "grade_dedup_bug",
+    "grade_fk_integrity",
+    "grade_schema_drift",
+    "compute_episode_score",
+]