Update app.py

app.py

@@ -1,13 +1,9 @@
 # Adaptive SQL Trainer — Domain Randomized with OpenAI (Gradio + SQLite)
-# - Randomizes a relational domain via OpenAI (bookstore, retail sales, wholesaler,
-#   sales tax, oil & gas wells, marketing) OR falls back to a built-in dataset.
-# - Builds 3–4 related tables (schema + seed rows) in SQLite.
-# - Generates 8–12 randomized SQL questions with varied phrasings.
-# - Validates answers by executing canonical SQL and comparing result sets.
-# - Provides tailored feedback (SQLite dialect, cartesian products, aggregates, aliases).
-# - Always shows data results at the bottom pane.
-#
-# Hugging Face Spaces: set OPENAI_API_KEY in secrets to enable randomization.
+# - OpenAI randomizes a domain and questions (fallback dataset if unavailable).
+# - 3–4 related tables with seed rows installed in SQLite.
+# - Students practice SELECT, WHERE, JOINs (INNER/LEFT), aliases, views, CTAS/SELECT INTO.
+# - Validator now enforces columns only when the prompt asks for them; otherwise it focuses on rows.
+# - ERD shows all FK edges in light gray and dynamically HIGHLIGHTS edges implied by the student’s JOINs.
 
 import os
 import re
@@ -18,7 +14,7 @@ import sqlite3
 import threading
 from dataclasses import dataclass
 from datetime import datetime, timezone
-from typing import List, Dict, Any, Tuple, Optional
+from typing import List, Dict, Any, Tuple, Optional, Set
 
 import gradio as gr
 import pandas as pd
@@ -43,7 +39,169 @@ def _candidate_models():
     seen = set()
     return [m for m in base if m and (m not in seen and not seen.add(m))]
 
-# -------------------- Global settings --------------------
+# -------------------- ERD drawing (headless) --------------------
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+from matplotlib.patches import Rectangle
+from io import BytesIO
+from PIL import Image
+
+PLOT_FIGSIZE = (7.6, 3.8)
+PLOT_DPI = 120
+PLOT_HEIGHT = 300
+
+def _fig_to_pil(fig) -> Image.Image:
+    buf = BytesIO()
+    fig.tight_layout()
+    fig.savefig(buf, format="png", dpi=PLOT_DPI, bbox_inches="tight")
+    plt.close(fig)
+    buf.seek(0)
+    return Image.open(buf)
+
+def draw_dynamic_erd(
+    schema: Dict[str, Any],
+    highlight_tables: Optional[Set[str]] = None,
+    highlight_edges: Optional[Set[Tuple[str, str]]] = None,
+) -> Image.Image:
+    """
+    Draw tables + FK edges. If highlight_* provided, overlay those tables/edges in bold.
+    highlight_edges uses (src_table, dst_table) with dst_table = referenced table.
+    """
+    highlight_tables = set(highlight_tables or [])
+    # normalize edges so (A,B) & (B,A) match regardless of direction
+    def _norm_edge(a, b): return tuple(sorted([a, b]))
+    H = set(_norm_edge(*e) for e in (highlight_edges or set()))
+
+    tables = schema.get("tables", [])
+    if not tables:
+        fig, ax = plt.subplots(figsize=PLOT_FIGSIZE); ax.axis("off")
+        ax.text(0.5, 0.5, "No tables to diagram.", ha="center", va="center")
+        return _fig_to_pil(fig)
+
+    # Layout tables horizontally
+    n = len(tables)
+    fig, ax = plt.subplots(figsize=PLOT_FIGSIZE); ax.axis("off")
+    margin = 0.03
+    width = (1 - margin * (n + 1)) / max(n, 1)
+    height = 0.70
+    y = 0.20
+
+    # Build quick FK lookup: [(src_table, dst_table)]
+    fk_edges = []
+    for t in tables:
+        for fk in t.get("fks", []) or []:
+            dst = fk.get("ref_table")
+            if dst:
+                fk_edges.append((t["name"], dst))
+
+    # Draw table boxes + columns
+    boxes: Dict[str, Tuple[float,float,float,float]] = {}
+    for i, t in enumerate(tables):
+        tx = margin + i * (width + margin)
+        boxes[t["name"]] = (tx, y, width, height)
+
+        # Highlight table border if used in current SQL
+        lw = 2.0 if t["name"] in highlight_tables else 1.2
+        ax.add_patch(Rectangle((tx, y), width, height, fill=False, lw=lw))
+        ax.text(tx + 0.01, y + height - 0.04, t["name"], fontsize=10, ha="left", va="top", weight="bold")
+
+        yy = y + height - 0.09
+        pkset = set(t.get("pk", []) or [])
+        # For FK annotation by column
+        fk_map: Dict[str, List[Tuple[str, str]]] = {}
+        for fk in t.get("fks", []) or []:
+            ref_tbl = fk.get("ref_table", "")
+            for c, rc in zip(fk.get("columns", []) or [], fk.get("ref_columns", []) or []):
+                fk_map.setdefault(c, []).append((ref_tbl, rc))
+
+        for col in t.get("columns", []):
+            nm = col.get("name", "")
+            tag = ""
+            if nm in pkset:
+                tag = " (PK)"
+            if nm in fk_map:
+                ref = fk_map[nm][0]
+                tag = f" (FK→{ref[0]}.{ref[1]})" if not tag else tag.replace(")", f", FK→{ref[0]}.{ref[1]})")
+            ax.text(tx + 0.016, yy, f"{nm}{tag}", fontsize=9, ha="left", va="top")
+            yy -= 0.055
+
+    # Draw FK edges: light gray
+    for (src, dst) in fk_edges:
+        if src not in boxes or dst not in boxes: 
+            continue
+        (x1, y1, w1, h1) = boxes[src]
+        (x2, y2, w2, h2) = boxes[dst]
+        ax.annotate("",
+                    xy=(x2 + w2/2.0, y2 + h2),
+                    xytext=(x1 + w1/2.0, y1),
+                    arrowprops=dict(arrowstyle="->", lw=1.0, color="#cccccc"))
+
+    # Overlay highlighted edges: bold, darker
+    for (src, dst) in fk_edges:
+        if _norm_edge(src, dst) in H:
+            (x1, y1, w1, h1) = boxes[src]
+            (x2, y2, w2, h2) = boxes[dst]
+            ax.annotate("",
+                        xy=(x2 + w2/2.0, y2 + h2),
+                        xytext=(x1 + w1/2.0, y1),
+                        arrowprops=dict(arrowstyle="->", lw=2.6, color="#2b6cb0"))
+
+    ax.text(0.5, 0.06, f"Domain: {schema.get('domain','unknown')}", fontsize=9, ha="center")
+    return _fig_to_pil(fig)
+
+# Parse JOINs from SQL and turn them into tables/edges to highlight on ERD
+JOIN_TBL_RE = re.compile(r"\b(?:from|join)\s+([a-z_]\w*)(?:\s+(?:as\s+)?([a-z_]\w*))?", re.IGNORECASE)
+EQ_ON_RE = re.compile(r"([a-z_]\w*)\.[a-z_]\w*\s*=\s*([a-z_]\w*)\.[a-z_]\w*", re.IGNORECASE)
+USING_RE = re.compile(r"\bjoin\s+([a-z_]\w*)(?:\s+(?:as\s+)?([a-z_]\w*))?\s+using\s*\(", re.IGNORECASE)
+
+def sql_highlights(sql: str, schema: Dict[str, Any]) -> Tuple[Set[str], Set[Tuple[str, str]]]:
+    """
+    Returns (highlight_tables, highlight_edges) based on the student's SQL.
+    - Tables: any table appearing after FROM or JOIN (by name or alias).
+    - Edges: pairs inferred from `a.col = b.col` or JOIN ... USING (...).
+    """
+    if not sql:
+        return set(), set()
+
+    low = " ".join(sql.strip().split())
+    # Alias map alias->table and list of tables in join order
+    alias_to_table: Dict[str, str] = {}
+    join_order: List[str] = []
+
+    for m in JOIN_TBL_RE.finditer(low):
+        table = m.group(1)
+        alias = m.group(2) or table
+        alias_to_table[alias] = table
+        join_order.append(alias)
+
+    # Edges from explicit equality ON clauses
+    edges: Set[Tuple[str, str]] = set()
+    for a1, a2 in EQ_ON_RE.findall(low):
+        t1 = alias_to_table.get(a1, a1)
+        t2 = alias_to_table.get(a2, a2)
+        if t1 != t2:
+            edges.add((t1, t2))
+
+    # Heuristic for USING(): connect the previous alias with the current JOIN alias
+    if USING_RE.search(low) and len(join_order) >= 2:
+        for i in range(1, len(join_order)):
+            t_left = alias_to_table.get(join_order[i-1], join_order[i-1])
+            t_right = alias_to_table.get(join_order[i], join_order[i])
+            if t_left != t_right:
+                edges.add((t_left, t_right))
+
+    # Highlight tables used (map aliases back to table names)
+    used_tables = {alias_to_table.get(a, a) for a in join_order}
+
+    # Normalize edges to actual table names present in schema
+    schema_tables = {t["name"] for t in schema.get("tables", [])}
+    edges = { (a, b) for (a, b) in edges if a in schema_tables and b in schema_tables }
+    used_tables = { t for t in used_tables if t in schema_tables }
+
+    return used_tables, edges
+
+# -------------------- SQLite + locking --------------------
 DB_DIR = "/data" if os.path.exists("/data") else "."
 DB_PATH = os.path.join(DB_DIR, "sql_trainer_dynamic.db")
 EXPORT_DIR = "."
@@ -51,15 +209,9 @@ RANDOM_SEED = int(os.getenv("RANDOM_SEED", "7"))
 random.seed(RANDOM_SEED)
 SYS_RAND = random.SystemRandom()
 
-# -------------------- SQLite connection + locking --------------------
 DB_LOCK = threading.RLock()
 
 def connect_db():
-    """
-    Single shared connection usable across threads.
-    All operations (reads + writes) serialized via DB_LOCK.
-    WAL mode enables concurrent reads.
-    """
     con = sqlite3.connect(DB_PATH, check_same_thread=False)
     con.execute("PRAGMA journal_mode=WAL;")
     con.execute("PRAGMA synchronous=NORMAL;")
@@ -104,7 +256,7 @@ def init_progress_tables(con: sqlite3.Connection):
 
 init_progress_tables(CONN)
 
-# -------------------- Fallback dataset (if no OpenAI) --------------------
+# -------------------- Fallback dataset & questions --------------------
 FALLBACK_SCHEMA = {
     "domain": "bookstore",
     "tables": [
@@ -217,10 +369,8 @@ FALLBACK_QUESTIONS = [
      "requires_aliases":False,"required_aliases":[]},
 ]
 
-# -------------------- OpenAI JSON request helpers --------------------
-DOMAIN_AND_QUESTIONS_SCHEMA = {
-    "required": ["domain", "tables", "questions"]
-}
+# --------------- OpenAI prompts + parsing helpers ---------------
+DOMAIN_AND_QUESTIONS_SCHEMA = {"required": ["domain", "tables", "questions"]}
 
 def _domain_prompt(prev_domain: Optional[str]) -> str:
     extra = f" Avoid using the previous domain '{prev_domain}' if possible." if prev_domain else ""
@@ -232,28 +382,22 @@ Rules:
 - One domain chosen from: bookstore, retail sales, wholesaler, sales tax, oil and gas wells, marketing.
 - Tables: SQLite-friendly. Use snake_case. Each table has: name, pk (list of column names),
   columns (list of {{name,type}}), fks (list of {{columns,ref_table,ref_columns}}), rows (8–15 small seed rows).
-- Questions: diverse natural language. Categories: "SELECT *", "SELECT columns", "WHERE", "Aliases",
+- Questions: categories among "SELECT *", "SELECT columns", "WHERE", "Aliases",
   "JOIN (INNER)", "JOIN (LEFT)", "Aggregation", "VIEW", "CTAS / SELECT INTO".
   Include at least one LEFT JOIN, one VIEW creation, one CTAS or SELECT INTO.
   Provide 1–3 'answer_sql' strings per question. Prefer SQLite-compatible SQL. Do NOT use RIGHT/FULL OUTER JOIN.
   For 1–2 questions, set requires_aliases=true and list required_aliases.
 
-Example top-level keys (do not include comments in output):
-{{
-  "domain": "retail sales",
-  "tables": [...],
-  "questions": [...]
-}}
+Example top-level keys:
+{{"domain":"retail sales","tables":[...],"questions":[...]}}
 """
 
 def _loose_json_parse(s: str) -> Optional[dict]:
-    """Extract the first JSON object from a possibly-wrapped string."""
     try:
         return json.loads(s)
     except Exception:
         pass
-    start = s.find("{")
-    end = s.rfind("}")
+    start = s.find("{"); end = s.rfind("}")
     if start != -1 and end != -1 and end > start:
         try:
             return json.loads(s[start:end+1])
@@ -261,94 +405,64 @@ def _loose_json_parse(s: str) -> Optional[dict]:
             return None
     return None
 
-# -------------------- Canonicalization & validation --------------------
+# Canonicalization of LLM output (questions & tables)
 _SQL_FENCE = re.compile(r"```sql(.*?)```", re.IGNORECASE | re.DOTALL)
 _CODE_FENCE = re.compile(r"```(.*?)```", re.DOTALL)
 
 def _strip_code_fences(txt: str) -> str:
-    if txt is None:
-        return ""
+    if txt is None: return ""
     m = _SQL_FENCE.findall(txt)
-    if m:
-        return "\n".join([x.strip() for x in m if x.strip()])
+    if m: return "\n".join([x.strip() for x in m if x.strip()])
     m2 = _CODE_FENCE.findall(txt)
-    if m2:
-        return "\n".join([x.strip() for x in m2 if x.strip()])
+    if m2: return "\n".join([x.strip() for x in m2 if x.strip()])
     return txt.strip()
 
 def _as_list_of_sql(val) -> List[str]:
-    if val is None:
-        return []
+    if val is None: return []
     if isinstance(val, str):
         s = _strip_code_fences(val)
         parts = [p.strip() for p in s.split("\n") if p.strip()]
-        # if it’s a single long line, keep as is
         return parts or ([s] if s else [])
     if isinstance(val, list):
         out = []
         for v in val:
             if isinstance(v, str):
                 s = _strip_code_fences(v)
-                if s:
-                    out.append(s)
+                if s: out.append(s)
         return out
     return []
 
 def _canon_question(q: Dict[str, Any]) -> Optional[Dict[str, Any]]:
-    """Normalize one question; return None if missing critical fields."""
-    if not isinstance(q, dict):
-        return None
-    # field mapping/synonyms
+    if not isinstance(q, dict): return None
     cat = q.get("category") or q.get("type") or q.get("topic")
     prompt = q.get("prompt_md") or q.get("prompt") or q.get("question") or q.get("text")
     answer_sql = q.get("answer_sql") or q.get("answers") or q.get("solutions") or q.get("sql")
     diff = q.get("difficulty") or 1
     req_alias = bool(q.get("requires_aliases", False))
     req_aliases = q.get("required_aliases") or []
-
     cat = str(cat).strip() if cat is not None else ""
     prompt = str(prompt).strip() if prompt is not None else ""
     answers = _as_list_of_sql(answer_sql)
-
-    if not cat or not prompt or not answers:
-        return None
-
-    # keep only known categories if provided; otherwise accept free text
+    if not cat or not prompt or not answers: return None
     known = {
         "SELECT *","SELECT columns","WHERE","Aliases",
         "JOIN (INNER)","JOIN (LEFT)","Aggregation","VIEW","CTAS / SELECT INTO"
     }
     if cat not in known:
-        # Try to map rough names to our set
         low = cat.lower()
-        if "select *" in low:
-            cat = "SELECT *"
-        elif "select col" in low or "columns" in low:
-            cat = "SELECT columns"
-        elif "where" in low or "filter" in low:
-            cat = "WHERE"
-        elif "alias" in low:
-            cat = "Aliases"
-        elif "left" in low:
-            cat = "JOIN (LEFT)"
-        elif "inner" in low or "join" in low:
-            cat = "JOIN (INNER)"
-        elif "agg" in low or "group" in low:
-            cat = "Aggregation"
-        elif "view" in low:
-            cat = "VIEW"
-        elif "into" in low or "ctas" in low or "create table" in low:
-            cat = "CTAS / SELECT INTO"
-        else:
-            # leave as-is; still usable for practice buckets
-            pass
-
-    # normalize aliases list
+        if "select *" in low: cat = "SELECT *"
+        elif "columns" in low: cat = "SELECT columns"
+        elif "where" in low or "filter" in low: cat = "WHERE"
+        elif "alias" in low: cat = "Aliases"
+        elif "left" in low: cat = "JOIN (LEFT)"
+        elif "inner" in low or "join" in low: cat = "JOIN (INNER)"
+        elif "agg" in low or "group" in low: cat = "Aggregation"
+        elif "view" in low: cat = "VIEW"
+        elif "into" in low or "ctas" in low: cat = "CTAS / SELECT INTO"
     if isinstance(req_aliases, str):
         req_aliases = [a.strip() for a in re.split(r"[,\s]+", req_aliases) if a.strip()]
     elif not isinstance(req_aliases, list):
         req_aliases = []
-
     return {
         "id": str(q.get("id") or f"LLM_{int(time.time()*1000)}_{random.randint(100,999)}"),
         "category": cat,
@@ -362,22 +476,17 @@ def _canon_question(q: Dict[str, Any]) -> Optional[Dict[str, Any]]:
 def _canon_tables(tables: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
     out = []
     for t in (tables or []):
-        if not isinstance(t, dict): 
-            continue
+        if not isinstance(t, dict): continue
         name = str(t.get("name","")).strip()
-        if not name: 
-            continue
+        if not name: continue
         cols = t.get("columns") or []
         good_cols = []
         for c in cols:
-            if not isinstance(c, dict): 
-                continue
+            if not isinstance(c, dict): continue
             cname = str(c.get("name","")).strip()
             ctype = str(c.get("type","TEXT")).strip() or "TEXT"
-            if cname:
-                good_cols.append({"name": cname, "type": ctype})
-        if not good_cols:
-            continue
+            if cname: good_cols.append({"name": cname, "type": ctype})
+        if not good_cols: continue
         pk = t.get("pk") or []
         if isinstance(pk, str): pk = [pk]
         fks = t.get("fks") or []
@@ -391,31 +500,22 @@ def _canon_tables(tables: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
         })
     return out
 
-# -------------------- LLM call --------------------
 def llm_generate_domain_and_questions(prev_domain: Optional[str]) -> Tuple[Optional[Dict[str,Any]], Optional[str], Optional[str], Dict[str,int]]:
-    """
-    Returns (obj, error_message, model_used, stats_dict).
-    stats_dict contains {"accepted_questions": n, "dropped_questions": m}
-    """
     if not OPENAI_AVAILABLE or not os.getenv("OPENAI_API_KEY"):
         return None, "OpenAI client not available or OPENAI_API_KEY missing.", None, {"accepted_questions":0,"dropped_questions":0}
-
     errors = []
     prompt = _domain_prompt(prev_domain)
-
     for model in _candidate_models():
         try:
-            # Try JSON mode first
             try:
                 chat = _client.chat.completions.create(
                     model=model,
                     messages=[{"role":"user","content": prompt}],
                     temperature=0.6,
-                    response_format={"type":"json_object"}  # newer SDKs
+                    response_format={"type":"json_object"}
                 )
                 data_text = chat.choices[0].message.content
             except TypeError:
-                # Older SDKs: no response_format ⇒ plain chat + strict instructions
                 chat = _client.chat.completions.create(
                     model=model,
                     messages=[{"role":"system","content":"Return ONLY a JSON object. No markdown."},
@@ -423,52 +523,35 @@ def llm_generate_domain_and_questions(prev_domain: Optional[str]) -> Tuple[Optio
                     temperature=0.6
                 )
                 data_text = chat.choices[0].message.content
-
             obj_raw = _loose_json_parse(data_text or "")
             if not obj_raw:
                 raise RuntimeError("Could not parse JSON from model output.")
-
-            # Minimal top-level validation
             for k in DOMAIN_AND_QUESTIONS_SCHEMA["required"]:
                 if k not in obj_raw:
                     raise RuntimeError(f"Missing key '{k}'")
-
-            # Canonicalize tables
             tables = _canon_tables(obj_raw.get("tables", []))
-            if not tables:
-                raise RuntimeError("No usable tables in LLM output.")
+            if not tables: raise RuntimeError("No usable tables in LLM output.")
             obj_raw["tables"] = tables
-
-            # Canonicalize questions
             dropped = 0
             clean_qs = []
             for q in obj_raw.get("questions", []):
                 cq = _canon_question(q)
-                if not cq:
-                    dropped += 1
-                    continue
-                # Strip RIGHT/FULL joins from answers
+                if not cq: dropped += 1; continue
                 answers = [a for a in cq["answer_sql"] if " right join " not in a.lower() and " full " not in a.lower()]
-                if not answers:
-                    dropped += 1
-                    continue
+                if not answers: dropped += 1; continue
                 cq["answer_sql"] = answers
                 clean_qs.append(cq)
-
             if not clean_qs:
                 raise RuntimeError("No usable questions after canonicalization.")
             stats = {"accepted_questions": len(clean_qs), "dropped_questions": dropped}
-
             obj_raw["questions"] = clean_qs
             return obj_raw, None, model, stats
-
         except Exception as e:
             errors.append(f"{model}: {e}")
             continue
-
     return None, "; ".join(errors) if errors else "Unknown LLM error.", None, {"accepted_questions":0,"dropped_questions":0}
 
-# -------------------- Schema install & question handling --------------------
+# -------------------- Install schema & prepare questions --------------------
 def drop_existing_domain_tables(con: sqlite3.Connection, keep_internal=True):
     with DB_LOCK:
         cur = con.cursor()
@@ -487,131 +570,30 @@ def install_schema(con: sqlite3.Connection, schema: Dict[str,Any]):
     drop_existing_domain_tables(con, keep_internal=True)
     with DB_LOCK:
         cur = con.cursor()
-        # Create tables
         for t in schema.get("tables", []):
             cols_sql = []
             pk = t.get("pk", [])
             for c in t.get("columns", []):
-                cname = c["name"]
-                ctype = c.get("type","TEXT")
-                cols_sql.append(f"{cname} {ctype}")
-            if pk:
-                cols_sql.append(f"PRIMARY KEY ({', '.join(pk)})")
+                cols_sql.append(f"{c['name']} {c.get('type','TEXT')}")
+            if pk: cols_sql.append(f"PRIMARY KEY ({', '.join(pk)})")
             create_sql = f"CREATE TABLE {t['name']} ({', '.join(cols_sql)})"
             cur.execute(create_sql)
-        # Insert rows
         for t in schema.get("tables", []):
-            if not t.get("rows"):
-                continue
+            if not t.get("rows"): continue
             cols = [c["name"] for c in t.get("columns", [])]
             qmarks = ",".join(["?"]*len(cols))
             insert_sql = f"INSERT INTO {t['name']} ({', '.join(cols)}) VALUES ({qmarks})"
             for r in t["rows"]:
-                if isinstance(r, dict):
-                    vals = [r.get(col, None) for col in cols]
+                if isinstance(r, dict): vals = [r.get(col, None) for col in cols]
                 elif isinstance(r, (list, tuple)):
-                    vals = list(r) + [None]*(len(cols)-len(r))
-                    vals = vals[:len(cols)]
-                else:
-                    continue
+                    vals = list(r) + [None]*(len(cols)-len(r)); vals = vals[:len(cols)]
+                else: continue
                 cur.execute(insert_sql, vals)
         con.commit()
-        # Persist schema JSON
         cur.execute("INSERT OR REPLACE INTO session_meta(id, domain, schema_json) VALUES (1, ?, ?)",
                     (schema.get("domain","unknown"), json.dumps(schema)))
         con.commit()
 
-def run_df(con: sqlite3.Connection, sql: str) -> pd.DataFrame:
-    with DB_LOCK:
-        return pd.read_sql_query(sql, con)
-
-def rewrite_select_into(sql: str) -> Tuple[str, Optional[str]]:
-    s = sql.strip().strip(";")
-    if re.search(r"\bselect\b.+\binto\b.+\bfrom\b", s, flags=re.IGNORECASE|re.DOTALL):
-        m = re.match(r"(?is)^\s*select\s+(.*?)\s+into\s+([A-Za-z_][A-Za-z0-9_]*)\s+from\s+(.*)$", s)
-        if m:
-            cols, tbl, rest = m.groups()
-            return f"CREATE TABLE {tbl} AS SELECT {cols} FROM {rest}", tbl
-    return sql, None
-
-def detect_unsupported_joins(sql: str) -> Optional[str]:
-    low = sql.lower()
-    if " right join " in low:
-        return "SQLite does not support RIGHT JOIN. Use LEFT JOIN in the opposite direction."
-    if " full join " in low or " full outer join " in low:
-        return "SQLite does not support FULL OUTER JOIN. Use LEFT JOIN plus UNION for the other side."
-    if " ilike " in low:
-        return "SQLite has no ILIKE. Use LOWER(col) LIKE LOWER('%pattern%')."
-    return None
-
-def detect_cartesian(con: sqlite3.Connection, sql: str, df_result: pd.DataFrame) -> Optional[str]:
-    low = sql.lower()
-    if " cross join " in low:
-        return "Query uses CROSS JOIN (cartesian product). Ensure this is intended."
-    comma_from = re.search(r"\bfrom\b\s+([a-z_]\w*)\s*,\s*([a-z_]\w*)", low)
-    missing_on = (" join " in low) and (" on " not in low) and (" using " not in low) and (" natural " not in low)
-    if comma_from or missing_on:
-        try:
-            with DB_LOCK:
-                cur = con.cursor()
-                if comma_from:
-                    t1, t2 = comma_from.groups()
-                else:
-                    m = re.search(r"\bfrom\b\s+([a-z_]\w*)", low)
-                    j = re.search(r"\bjoin\b\s+([a-z_]\w*)", low)
-                    if not m or not j:
-                        return "Possible cartesian product: no join condition detected."
-                    t1, t2 = m.group(1), j.group(1)
-                cur.execute(f"SELECT COUNT(*) FROM {t1}")
-                n1 = cur.fetchone()[0]
-                cur.execute(f"SELECT COUNT(*) FROM {t2}")
-                n2 = cur.fetchone()[0]
-            prod = n1 * n2
-            if len(df_result) == prod and prod > 0:
-                return f"Result row count equals {n1}×{n2}={prod}. Likely cartesian product (missing join)."
-        except Exception:
-            return "Possible cartesian product: no join condition detected."
-    return None
-
-def results_equal(df_a: pd.DataFrame, df_b: pd.DataFrame) -> bool:
-    if df_a.shape != df_b.shape:
-        return False
-    a = df_a.copy()
-    b = df_b.copy()
-    a.columns = [c.lower() for c in a.columns]
-    b.columns = [c.lower() for c in b.columns]
-    a = a.sort_values(list(a.columns)).reset_index(drop=True)
-    b = b.sort_values(list(b.columns)).reset_index(drop=True)
-    return a.equals(b)
-
-def aliases_present(sql: str, required_aliases: List[str]) -> bool:
-    low = re.sub(r"\s+", " ", sql.lower())
-    for al in required_aliases:
-        if f" {al}." not in low and f" as {al} " not in low:
-            return False
-    return True
-
-# -------------------- Question model helpers --------------------
-@dataclass
-class SQLQuestion:
-    id: str
-    category: str
-    difficulty: int
-    prompt_md: str
-    answer_sql: List[str]
-    requires_aliases: bool = False
-    required_aliases: List[str] = None
-
-def to_question_dict(q) -> Dict[str,Any]:
-    d = dict(q)
-    d.setdefault("requires_aliases", False)
-    d.setdefault("required_aliases", [])
-    return d
-
-def load_questions(obj_list: List[Dict[str,Any]]) -> List[Dict[str,Any]]:
-    return [to_question_dict(o) for o in obj_list]
-
-# -------------------- Domain bootstrap --------------------
 def bootstrap_domain_with_llm_or_fallback(prev_domain: Optional[str]):
     obj, err, model_used, stats = llm_generate_domain_and_questions(prev_domain)
     if obj is None:
@@ -621,13 +603,13 @@ def bootstrap_domain_with_llm_or_fallback(prev_domain: Optional[str]):
 def install_schema_and_prepare_questions(prev_domain: Optional[str]):
     schema, questions, info = bootstrap_domain_with_llm_or_fallback(prev_domain)
     install_schema(CONN, schema)
-    # Safety: if questions empty, fall back
     if not questions:
         questions = FALLBACK_QUESTIONS
-        info = {"source":"openai+fallback-questions","model":info.get("model"),"error":"LLM returned 0 usable questions; used fallback bank.","accepted":0,"dropped":0}
+        info = {"source":"openai+fallback-questions","model":info.get("model"),
+                "error":"LLM returned 0 usable questions; used fallback bank.","accepted":0,"dropped":0}
     return schema, questions, info
 
-# -------------------- Session state --------------------
+# -------------------- Session globals --------------------
 CURRENT_SCHEMA, CURRENT_QS, CURRENT_INFO = install_schema_and_prepare_questions(prev_domain=None)
 
 # -------------------- Progress + mastery --------------------
@@ -662,7 +644,6 @@ def fetch_attempts(con: sqlite3.Connection, user_id: str) -> pd.DataFrame:
         return pd.read_sql_query("SELECT * FROM attempts WHERE user_id=? ORDER BY id DESC", con, params=(user_id,))
 
 def pick_next_question(user_id: str) -> Dict[str,Any]:
-    # Defensive: ensure we always have a pool
     pool = CURRENT_QS if CURRENT_QS else FALLBACK_QUESTIONS
     df = fetch_attempts(CONN, user_id)
     stats = topic_stats(df)
@@ -671,21 +652,99 @@ def pick_next_question(user_id: str) -> Dict[str,Any]:
     cands = [q for q in pool if str(q.get("category","")).strip() == weakest] or pool
     return dict(random.choice(cands))
 
-# -------------------- Execution & feedback --------------------
+# -------------------- SQL execution & grading --------------------
+def run_df(con: sqlite3.Connection, sql: str) -> pd.DataFrame:
+    with DB_LOCK:
+        return pd.read_sql_query(sql, con)
+
+def rewrite_select_into(sql: str) -> Tuple[str, Optional[str]]:
+    s = sql.strip().strip(";")
+    if re.search(r"\bselect\b.+\binto\b.+\bfrom\b", s, flags=re.IGNORECASE|re.DOTALL):
+        m = re.match(r"(?is)^\s*select\s+(.*?)\s+into\s+([A-Za-z_][A-Za-z0-9_]*)\s+from\s+(.*)$", s)
+        if m:
+            cols, tbl, rest = m.groups()
+            return f"CREATE TABLE {tbl} AS SELECT {cols} FROM {rest}", tbl
+    return sql, None
+
+def detect_unsupported_joins(sql: str) -> Optional[str]:
+    low = sql.lower()
+    if " right join " in low:
+        return "SQLite does not support RIGHT JOIN. Use LEFT JOIN in the opposite direction."
+    if " full join " in low or " full outer join " in low:
+        return "SQLite does not support FULL OUTER JOIN. Use LEFT JOIN plus UNION."
+    if " ilike " in low:
+        return "SQLite has no ILIKE. Use LOWER(col) LIKE LOWER('%pattern%')."
+    return None
+
+def detect_cartesian(con: sqlite3.Connection, sql: str, df_result: pd.DataFrame) -> Optional[str]:
+    low = sql.lower()
+    if " cross join " in low: return "Query uses CROSS JOIN (cartesian product). Ensure this is intended."
+    comma_from = re.search(r"\bfrom\b\s+([a-z_]\w*)\s*,\s*([a-z_]\w*)", low)
+    missing_on = (" join " in low) and (" on " not in low) and (" using " not in low) and (" natural " not in low)
+    if comma_from or missing_on:
+        try:
+            with DB_LOCK:
+                cur = con.cursor()
+                if comma_from: t1, t2 = comma_from.groups()
+                else:
+                    m = re.search(r"\bfrom\b\s+([a-z_]\w*)", low); j = re.search(r"\bjoin\b\s+([a-z_]\w*)", low)
+                    if not m or not j: return "Possible cartesian product: no join condition detected."
+                    t1, t2 = m.group(1), j.group(1)
+                cur.execute(f"SELECT COUNT(*) FROM {t1}"); n1 = cur.fetchone()[0]
+                cur.execute(f"SELECT COUNT(*) FROM {t2}"); n2 = cur.fetchone()[0]
+            prod = n1 * n2
+            if len(df_result) == prod and prod > 0:
+                return f"Result row count equals {n1}×{n2}={prod}. Likely cartesian product (missing join)."
+        except Exception:
+            return "Possible cartesian product: no join condition detected."
+    return None
+
+# Column enforcement policy — only when the prompt asks for it
+def should_enforce_columns(q: Dict[str, Any]) -> bool:
+    cat = (q.get("category") or "").strip()
+    if cat in ("SELECT columns", "Aggregation", "VIEW", "CTAS / SELECT INTO"):
+        return True
+    prompt = (q.get("prompt_md") or "").lower()
+    # Signals that the projection is specified in the prompt
+    if re.search(r"`[^`]+`", q.get("prompt_md") or ""):  # backticked names
+        return True
+    if re.search(r"\((?:show|return|display)[^)]+\)", prompt):  # e.g., "(show title, price)"
+        return True
+    if re.search(r"\b(show|return|display|select)\b[^.]{0,100}\b(columns?|fields?|name|title|price)\b", prompt):
+        return True
+    return False
+
+def _normalize_columns(df: pd.DataFrame) -> pd.DataFrame:
+    out = df.copy()
+    out.columns = [str(c).strip().lower() for c in out.columns]
+    return out
+
+def results_equal_or_superset(df_student: pd.DataFrame, df_expected: pd.DataFrame) -> Tuple[bool, Optional[str]]:
+    a = _normalize_columns(df_student); b = _normalize_columns(df_expected)
+    if set(a.columns) == set(b.columns):
+        a2 = a[sorted(a.columns)].sort_values(sorted(a.columns)).reset_index(drop=True)
+        b2 = b[sorted(b.columns)].sort_values(sorted(b.columns)).reset_index(drop=True)
+        return (a2.equals(b2), None)
+    if set(b.columns).issubset(set(a.columns)):
+        a_proj = a[b.columns]
+        a2 = a_proj.sort_values(list(b.columns)).reset_index(drop=True)
+        b2 = b.sort_values(list(b.columns)).reset_index(drop=True)
+        if a2.equals(b2):
+            return True, "extra_columns"
+    return False, None
+
+def results_equal_rowcount_only(df_student: pd.DataFrame, df_expected: pd.DataFrame) -> bool:
+    # When projection isn't specified, match on row count only.
+    return df_student.shape[0] == df_expected.shape[0]
+
 def exec_student_sql(sql_text: str) -> Tuple[Optional[pd.DataFrame], Optional[str], Optional[str], Optional[str]]:
     if not sql_text or not sql_text.strip():
         return None, "Enter a SQL statement.", None, None
-
     sql_raw = sql_text.strip().rstrip(";")
     sql_rew, created_tbl = rewrite_select_into(sql_raw)
-    note = None
-    if sql_rew != sql_raw:
-        note = "Rewrote `SELECT ... INTO` to `CREATE TABLE ... AS SELECT ...` for SQLite."
-
+    note = "Rewrote `SELECT ... INTO` to `CREATE TABLE ... AS SELECT ...` for SQLite." if sql_rew != sql_raw else None
     unsup = detect_unsupported_joins(sql_rew)
-    if unsup:
-        return None, unsup, None, note
-
+    if unsup: return None, unsup, None, note
     try:
         low = sql_rew.lower()
         if low.startswith("select"):
@@ -695,72 +754,56 @@ def exec_student_sql(sql_text: str) -> Tuple[Optional[pd.DataFrame], Optional[st
         else:
             with DB_LOCK:
                 cur = CONN.cursor()
-                cur.execute(sql_rew)
-                CONN.commit()
-                # Preview newly created objects
+                cur.execute(sql_rew); CONN.commit()
                 if low.startswith("create view"):
                     m = re.match(r"(?is)^\s*create\s+view\s+(if\s+not\s+exists\s+)?([a-z_]\w*)\s+as\s+(select.*)$", low)
                     name = m.group(2) if m else None
                     if name:
-                        try:
-                            df = pd.read_sql_query(f"SELECT * FROM {name}", CONN)
-                            return df, None, None, note
-                        except Exception:
-                            return None, "View created but could not be queried.", None, note
+                        try: return pd.read_sql_query(f"SELECT * FROM {name}", CONN), None, None, note
+                        except Exception: return None, "View created but could not be queried.", None, note
                 if low.startswith("create table"):
                     tbl = created_tbl
                     if not tbl:
                         m = re.match(r"(?is)^\s*create\s+table\s+(if\s+not\s+exists\s+)?([a-z_]\w*)\s+as\s+select.*$", low)
                         tbl = m.group(2) if m else None
                     if tbl:
-                        try:
-                            df = pd.read_sql_query(f"SELECT * FROM {tbl}", CONN)
-                            return df, None, None, note
-                        except Exception:
-                            return None, "Table created but could not be queried.", None, note
+                        try: return pd.read_sql_query(f"SELECT * FROM {tbl}", CONN), None, None, note
+                        except Exception: return None, "Table created but could not be queried.", None, note
             return pd.DataFrame(), None, None, note
     except Exception as e:
         msg = str(e)
-        if "no such table" in msg.lower():
-            return None, f"{msg}. Check table names for this randomized domain.", None, note
-        if "no such column" in msg.lower():
-            return None, f"{msg}. Use correct column names or prefixes (alias.column).", None, note
-        if "ambiguous column name" in msg.lower():
-            return None, f"{msg}. Qualify the column with a table alias.", None, note
+        if "no such table" in msg.lower(): return None, f"{msg}. Check table names for this randomized domain.", None, note
+        if "no such column" in msg.lower(): return None, f"{msg}. Use correct column names or prefixes (alias.column).", None, note
+        if "ambiguous column name" in msg.lower(): return None, f"{msg}. Qualify the column with a table alias.", None, note
         if "misuse of aggregate" in msg.lower() or "aggregate functions are not allowed in" in msg.lower():
             return None, f"{msg}. You might need a GROUP BY for non-aggregated columns.", None, note
         if "near \"into\"" in msg.lower() and "syntax error" in msg.lower():
             return None, "SQLite doesn’t support `SELECT ... INTO`. I can rewrite it automatically—try again.", None, note
         if "syntax error" in msg.lower():
-            return None, f"Syntax error. Check commas, keywords, and parentheses. Raw error: {msg}", None, note
+            return None, f"Syntax error. Check commas, keywords, parentheses. Raw error: {msg}", None, note
         return None, f"SQL error: {msg}", None, note
 
 def answer_df(answer_sql: List[str]) -> Optional[pd.DataFrame]:
     for sql in answer_sql:
         try:
             low = sql.strip().lower()
-            if low.startswith("select"):
-                return run_df(CONN, sql)
+            if low.startswith("select"): return run_df(CONN, sql)
             if low.startswith("create view"):
                 m = re.match(r"(?is)^\s*create\s+view\s+(if\s+not\s+exists\s+)?([a-z_]\w*)\s+as\s+select.*$", low)
                 view_name = m.group(2) if m else "vw_tmp"
                 with DB_LOCK:
                     cur = CONN.cursor()
                     cur.execute(f"DROP VIEW IF EXISTS {view_name}")
-                    cur.execute(sql)
-                    CONN.commit()
+                    cur.execute(sql); CONN.commit()
                 return run_df(CONN, f"SELECT * FROM {view_name}")
             if low.startswith("create table"):
                 m = re.match(r"(?is)^\s*create\s+table\s+(if\s+not\s+exists\s+)?([a-z_]\w*)\s+as\s+select.*$", low)
                 tbl = m.group(2) if m else None
                 with DB_LOCK:
                     cur = CONN.cursor()
-                    if tbl:
-                        cur.execute(f"DROP TABLE IF EXISTS {tbl}")
-                    cur.execute(sql)
-                    CONN.commit()
-                if tbl:
-                    return run_df(CONN, f"SELECT * FROM {tbl}")
+                    if tbl: cur.execute(f"DROP TABLE IF EXISTS {tbl}")
+                    cur.execute(sql); CONN.commit()
+                if tbl: return run_df(CONN, f"SELECT * FROM {tbl}")
         except Exception:
             continue
     return None
@@ -771,7 +814,18 @@ def validate_answer(q: Dict[str,Any], student_sql: str, df_student: Optional[pd.
         return (df_student is not None), f"**Explanation:** Your statement executed successfully for this task."
     if df_student is None:
         return False, f"**Explanation:** Expected data result differs."
-    return results_equal(df_student, df_expected), f"**Explanation:** Compare your result to a canonical solution."
+    if should_enforce_columns(q):
+        ok, note = results_equal_or_superset(df_student, df_expected)
+        if ok and note == "extra_columns":
+            return True, "**Note:** You returned extra columns. The rows match; try selecting only the requested columns next time."
+        if ok:
+            return True, "**Explanation:** Your result matches a canonical solution."
+        return False, f"**Explanation:** Compare your result to a canonical solution."
+    else:
+        ok = results_equal_rowcount_only(df_student, df_expected)
+        if ok:
+            return True, "**Explanation:** Columns weren’t specified for this task; row count matches the canonical answer."
+        return False, "**Explanation:** For this task we compared row counts (projection not enforced) and they didn’t match."
 
 def log_attempt(user_id: str, qid: str, category: str, correct: bool, sql_text: str,
                 time_taken: float, difficulty: int, source: str, notes: str):
@@ -792,6 +846,7 @@ def start_session(name: str, session: dict):
                 gr.update(value="Please enter your name to begin.", visible=True),
                 gr.update(visible=False),
                 gr.update(visible=False),
+                draw_dynamic_erd(CURRENT_SCHEMA),
                 gr.update(visible=False),
                 pd.DataFrame(),
                 pd.DataFrame())
@@ -804,21 +859,25 @@ def start_session(name: str, session: dict):
 
     prompt = q["prompt_md"]
     stats = topic_stats(fetch_attempts(CONN, user_id))
+    erd = draw_dynamic_erd(CURRENT_SCHEMA)
     return (session,
             gr.update(value=f"**Question {q['id']}**\n\n{prompt}", visible=True),
-            gr.update(visible=True),  # show SQL input
-            gr.update(value="", visible=True),  # preview block
-            gr.update(visible=False),  # next btn hidden until submit
+            gr.update(visible=True),
+            gr.update(value="", visible=True),
+            erd,
+            gr.update(visible=False),
             stats,
             pd.DataFrame())
 
 def render_preview(sql_text: str, session: dict):
     if not session or "q" not in session:
-        return gr.update(value="", visible=False)
+        return gr.update(value="", visible=False), draw_dynamic_erd(CURRENT_SCHEMA)
     s = (sql_text or "").strip()
     if not s:
-        return gr.update(value="", visible=False)
-    return gr.update(value=f"**Preview:**\n\n```sql\n{s}\n```", visible=True)
+        return gr.update(value="", visible=False), draw_dynamic_erd(CURRENT_SCHEMA)
+    hi_tables, hi_edges = sql_highlights(s, CURRENT_SCHEMA)
+    erd = draw_dynamic_erd(CURRENT_SCHEMA, highlight_tables=hi_tables, highlight_edges=hi_edges)
+    return gr.update(value=f"**Preview:**\n\n```sql\n{s}\n```", visible=True), erd
 
 def submit_answer(sql_text: str, session: dict):
     if not session or "user_id" not in session or "q" not in session:
@@ -826,7 +885,6 @@ def submit_answer(sql_text: str, session: dict):
     user_id = session["user_id"]
     q = session["q"]
     elapsed = max(0.0, time.time() - session.get("start_ts", time.time()))
-
     df, err, warn, note = exec_student_sql(sql_text)
     details = []
     if note: details.append(f"ℹ️ {note}")
@@ -836,35 +894,27 @@ def submit_answer(sql_text: str, session: dict):
         log_attempt(user_id, q.get("id","?"), q.get("category","?"), False, sql_text, elapsed, int(q.get("difficulty",1)), "bank", " | ".join([err] + details))
         stats = topic_stats(fetch_attempts(CONN, user_id))
         return gr.update(value=fb, visible=True), pd.DataFrame(), gr.update(visible=True), stats
-
     alias_msg = None
-    if q.get("requires_aliases"):
-        if not aliases_present(sql_text, q.get("required_aliases", [])):
-            alias_msg = f"⚠️ This task asked for aliases {q.get('required_aliases', [])}. I didn’t detect them."
-
+    if q.get("requires_aliases") and not aliases_present(sql_text, q.get("required_aliases", [])):
+        alias_msg = f"⚠️ This task asked for aliases {q.get('required_aliases', [])}. I didn’t detect them."
     is_correct, explanation = validate_answer(q, sql_text, df)
     if warn: details.append(f"⚠️ {warn}")
     if alias_msg: details.append(alias_msg)
-
     prefix = "✅ **Correct!**" if is_correct else "❌ **Not quite.**"
     feedback = prefix
-    if details:
-        feedback += "\n\n" + "\n".join(details)
+    if details: feedback += "\n\n" + "\n".join(details)
     feedback += "\n\n" + explanation + "\n\n**One acceptable solution:**\n```sql\n" + q["answer_sql"][0].rstrip(";") + ";\n```"
-
     log_attempt(user_id, q["id"], q.get("category","?"), bool(is_correct), sql_text, elapsed, int(q.get("difficulty",1)), "bank", " | ".join(details))
     stats = topic_stats(fetch_attempts(CONN, user_id))
     return gr.update(value=feedback, visible=True), (df if df is not None else pd.DataFrame()), gr.update(visible=True), stats
 
 def next_question(session: dict):
     if not session or "user_id" not in session:
-        return session, gr.update(value="Start a session first.", visible=True), gr.update(visible=False), gr.update(visible=False)
+        return session, gr.update(value="Start a session first.", visible=True), gr.update(visible=False), draw_dynamic_erd(CURRENT_SCHEMA), gr.update(visible=False)
     user_id = session["user_id"]
     q = pick_next_question(user_id)
-    session["qid"] = q["id"]
-    session["q"] = q
-    session["start_ts"] = time.time()
-    return session, gr.update(value=f"**Question {q['id']}**\n\n{q['prompt_md']}", visible=True), gr.update(value="", visible=True), gr.update(visible=False)
+    session["qid"] = q["id"]; session["q"] = q; session["start_ts"] = time.time()
+    return session, gr.update(value=f"**Question {q['id']}**\n\n{q['prompt_md']}", visible=True), gr.update(value="", visible=True), draw_dynamic_erd(CURRENT_SCHEMA), gr.update(visible=False)
 
 def show_hint(session: dict):
     if not session or "q" not in session:
@@ -885,8 +935,7 @@ def show_hint(session: dict):
 
 def export_progress(user_name: str):
     slug = "-".join((user_name or "").lower().split())
-    if not slug:
-        return None
+    if not slug: return None
     user_id = slug[:64]
     with DB_LOCK:
         df = pd.read_sql_query("SELECT * FROM attempts WHERE user_id=? ORDER BY id DESC", CONN, params=(user_id,))
@@ -897,25 +946,20 @@ def export_progress(user_name: str):
 
 def _domain_status_md():
     if CURRENT_INFO.get("source","") in ("openai","openai+fallback-questions"):
-        note = ""
-        if CURRENT_INFO.get("source") == "openai+fallback-questions":
-            note = " (LLM domain ok; used fallback questions)"
-        accepted = CURRENT_INFO.get("accepted",0)
-        dropped = CURRENT_INFO.get("dropped",0)
-        return (
-            f"✅ **Domain via OpenAI** `{CURRENT_INFO.get('model','?')}` → **{CURRENT_SCHEMA.get('domain','?')}**{note}. "
-            f"Accepted questions: {accepted}, dropped: {dropped}.  \n"
-            f"Tables: {', '.join(t['name'] for t in CURRENT_SCHEMA.get('tables', []))}."
-        )
-    err = CURRENT_INFO.get("error","")
-    err_short = (err[:160] + "…") if len(err) > 160 else err
+        note = " (LLM domain ok; used fallback questions)" if CURRENT_INFO.get("source") == "openai+fallback-questions" else ""
+        accepted = CURRENT_INFO.get("accepted",0); dropped = CURRENT_INFO.get("dropped",0)
+        return (f"✅ **Domain via OpenAI** `{CURRENT_INFO.get('model','?')}` → **{CURRENT_SCHEMA.get('domain','?')}**{note}. "
+                f"Accepted questions: {accepted}, dropped: {dropped}.  \n"
+                f"Tables: {', '.join(t['name'] for t in CURRENT_SCHEMA.get('tables', []))}.")
+    err = CURRENT_INFO.get("error",""); err_short = (err[:160] + "…") if len(err) > 160 else err
     return f"⚠️ **OpenAI randomization unavailable** → using fallback **{CURRENT_SCHEMA.get('domain','?')}**.\n\n> Reason: {err_short}"
 
 def regenerate_domain():
     global CURRENT_SCHEMA, CURRENT_QS, CURRENT_INFO
     prev = CURRENT_SCHEMA.get("domain") if CURRENT_SCHEMA else None
     CURRENT_SCHEMA, CURRENT_QS, CURRENT_INFO = install_schema_and_prepare_questions(prev_domain=prev)
-    return gr.update(value=_domain_status_md(), visible=True)
+    erd = draw_dynamic_erd(CURRENT_SCHEMA)
+    return gr.update(value=_domain_status_md(), visible=True), erd
 
 def preview_table(tbl: str):
     try:
@@ -925,8 +969,7 @@ def preview_table(tbl: str):
 
 def list_tables_for_preview():
     df = run_df(CONN, "SELECT name, type FROM sqlite_master WHERE type in ('table','view') AND name NOT IN ('users','attempts','session_meta') ORDER BY type, name")
-    if df.empty:
-        return ["(no tables)"]
+    if df.empty: return ["(no tables)"]
     return df["name"].tolist()
 
 # -------------------- UI --------------------
@@ -937,14 +980,11 @@ with gr.Blocks(title="Adaptive SQL Trainer — Randomized Domains") as demo:
         - Uses **OpenAI** (if configured) to randomize a domain (bookstore, retail sales, wholesaler,
           sales tax, oil & gas wells, marketing), generate **3–4 tables** and **8–12** questions.
         - Practice `SELECT`, `WHERE`, `JOIN` (INNER/LEFT), **aliases**, **views**, and **CTAS / SELECT INTO**.
-        - The app explains **SQLite quirks** (no RIGHT/FULL JOIN) and flags likely **cartesian products**.
-
-        > Set your `OPENAI_API_KEY` in Space secrets to enable randomization.
+        - **ERD highlights your JOINs** as you type; all FK edges remain visible in light gray.
         """
     )
 
     with gr.Row():
-        # -------- Left column: controls + quick preview ----------
         with gr.Column(scale=1):
             name_box = gr.Textbox(label="Your Name", placeholder="e.g., Jordan Alvarez")
             start_btn = gr.Button("Start / Resume Session", variant="primary")
@@ -967,12 +1007,11 @@ with gr.Blocks(title="Adaptive SQL Trainer — Randomized Domains") as demo:
             tbl_btn = gr.Button("Preview")
             preview_df = gr.Dataframe(value=pd.DataFrame(), interactive=False)
 
-        # -------- Right column: task + feedback + mastery + results ----------
         with gr.Column(scale=2):
             prompt_md = gr.Markdown(visible=False)
             sql_input = gr.Textbox(label="Your SQL", placeholder="Type SQL here (end ; optional).", lines=6, visible=False)
-
             preview_md = gr.Markdown(visible=False)
+            er_image = gr.Image(label="Entity Diagram", value=draw_dynamic_erd(CURRENT_SCHEMA), height=PLOT_HEIGHT)
 
             with gr.Row():
                 submit_btn = gr.Button("Run & Submit", variant="primary")
@@ -983,12 +1022,8 @@ with gr.Blocks(title="Adaptive SQL Trainer — Randomized Domains") as demo:
 
             gr.Markdown("---")
             gr.Markdown("### Your Progress by Category")
-            mastery_df = gr.Dataframe(
-                headers=["category","attempts","correct","accuracy"],
-                col_count=(4, "dynamic"),
-                row_count=(0, "dynamic"),
-                interactive=False
-            )
+            mastery_df = gr.Dataframe(headers=["category","attempts","correct","accuracy"],
+                                      col_count=(4,"dynamic"), row_count=(0,"dynamic"), interactive=False)
 
             gr.Markdown("---")
             gr.Markdown("### Result Preview")
@@ -998,12 +1033,12 @@ with gr.Blocks(title="Adaptive SQL Trainer �� Randomized Domains") as demo:
     start_btn.click(
         start_session,
         inputs=[name_box, session_state],
-        outputs=[session_state, prompt_md, sql_input, preview_md, next_btn, mastery_df, result_df],
+        outputs=[session_state, prompt_md, sql_input, preview_md, er_image, next_btn, mastery_df, result_df],
     )
     sql_input.change(
         render_preview,
         inputs=[sql_input, session_state],
-        outputs=[preview_md],
+        outputs=[preview_md, er_image],
     )
     submit_btn.click(
         submit_answer,
@@ -1013,7 +1048,7 @@ with gr.Blocks(title="Adaptive SQL Trainer — Randomized Domains") as demo:
     next_btn.click(
         next_question,
         inputs=[session_state],
-        outputs=[session_state, prompt_md, sql_input, next_btn],
+        outputs=[session_state, prompt_md, sql_input, er_image, next_btn],
     )
     hint_btn.click(
         show_hint,
@@ -1028,15 +1063,14 @@ with gr.Blocks(title="Adaptive SQL Trainer — Randomized Domains") as demo:
     regen_btn.click(
         regenerate_domain,
         inputs=[],
-        outputs=[regen_fb],
+        outputs=[regen_fb, er_image],
     )
     tbl_btn.click(
         lambda name: preview_table(name),
         inputs=[tbl_dd],
         outputs=[preview_df]
     )
-    # Keep dropdown fresh after regeneration
-    regen_btn.click(
+    regen_btn.click(  # refresh list after regeneration
         lambda: gr.update(choices=list_tables_for_preview()),
         inputs=[],
         outputs=[tbl_dd]