sql-server-query-ai/lib/digest.ts

import { XMLParser } from "fast-xml-parser";

export class PlanError extends Error {}
export const MAX_XML_BYTES = 12_000_000;

const parser = new XMLParser({ ignoreAttributes: false, attributeNamePrefix: "", removeNSPrefix: true });

export type Operator = { id: string; op: string; cost: number; estRows: number; actualRows: number | null;
  execs: number | null; io: number | null; cpu: number | null; spillLevels: number | null; parallel: boolean };
export type RowGuess = { id: string; op: string; est: number; actual: number; ratio: number };
export type MissingIndex = { table: string; equality: string[]; inequality: string[]; included: string[]; impact: number };
export type Wait = { type: string; ms: number };
export type Digest = { sql: string; bytes: number; subtreeCost: number | null; parallel: boolean;
  nonParallelReason: string | null; earlyAbort: string | null; tables: string[]; topOperators: Operator[];
  rowGuessErrors: RowGuess[]; missingIndexes: MissingIndex[]; warnings: string[]; waits: Wait[] };

function collect(node: any, name: string, out: any[] = []): any[] {
  if (Array.isArray(node)) { for (const n of node) collect(n, name, out); return out; }
  if (node && typeof node === "object") for (const [k, v] of Object.entries(node)) {
    if (k === "#text") continue;
    if (k === name) (Array.isArray(v) ? v : [v]).forEach((n) => out.push(n));
    collect(v, name, out);
  }
  return out;
}
const num = (v: any): number | null =>
  (v === undefined || v === null || v === "" || Number.isNaN(Number(v)) ? null : Number(v));
const text = (v: any): string | null =>
  (typeof v === "object" && v !== null ? (String(v["#text"] ?? "") || null) : (v == null ? null : String(v)));
const unbrack = (v: any) => String(v ?? "").replace(/[[\]]/g, "");
const colNames = (node: any) => collect(node, "Column").map((c: any) => unbrack(c.Name)).filter(Boolean);

export function digestPlan(raw: string): Digest {
  const xml = raw.replace(/^\uFEFF/, "");        // SSMS writes a BOM; the XML parser does not want it
  if (!xml || xml.trim().length < 40) throw new PlanError("Empty input. Paste the ShowPlan XML, or drop a .sqlplan file.");
  if (xml.length > MAX_XML_BYTES)
    throw new PlanError(`Plan is ${Math.round(xml.length / 1e6)} MB, over the ${MAX_XML_BYTES / 1e6} MB cap.`);
  let doc: any;
  // Name the parser's complaint. "re-export the plan" is what a user gets told either way, so a truncated
  // file, a stray entity and a genuinely wrong file have to stay distinguishable.
  try { doc = parser.parse(xml); }
  catch (e: any) {
    const why = String(e?.message ?? e).split("\n")[0].slice(0, 140);
    throw new PlanError(`The XML stopped parsing: ${why}. Check the file was saved completely (SSMS: Save As, not copy-paste).`);
  }
  if (!doc?.ShowPlanXML) throw new PlanError("Not a SQL Server ShowPlan XML document. In SSMS: press Ctrl+M, run the query, then save the Execution Plan tab as .sqlplan.");

  const stmts = collect(doc, "StmtSimple");
  const stmt = stmts.find((s: any) => s?.QueryBlock || s?.RelOp) ?? stmts[0] ?? {};
  const operators: Operator[] = collect(doc, "RelOp").map((o: any) => {
    const rt = collect(o, "RunTimeCountersPerThread")[0];
    return {
      id: String(o.NodeId ?? o["Node_id"] ?? "?"),
      op: String(o.LogicalOp ?? o.PhysicalOp ?? "Unknown"),
      // Real ShowPlan spells these EstimateIO / EstimateCPU; builds differ, so fall through every spelling.
      cost: num(o.EstimatedOperatorCost) ?? num(o.EstimatedTotalSubtreeCost) ?? num(o.EstimateIO) ?? 0,
      estRows: num(o.EstimateRows) ?? 0,
      actualRows: rt ? num(rt.ActualRows) : null,
      execs: rt ? num(rt.ActualExecs) : null,
      io: num(o.EstimateIO) ?? num(o.EstimatedIO), cpu: num(o.EstimateCPU) ?? num(o.EstimatedCPU),
      spillLevels: num(collect(o, "SpillToTempDb")[0]?.SpillLevel),
      parallel: o.Parallel === "true",
    };
  });

  // Database/Schema/Table arrive as separate attributes, each bracketed: rebuild the three-part name.
  const fullName = (n: any) => [n.Database, n.Schema, n.Table].map(unbrack).filter(Boolean).join(".");
  const tables = [...new Set([...collect(doc, "IndexScan"), ...collect(doc, "TableScan"), ...collect(doc, "Table")]
    .map(fullName).filter(Boolean))];

  const missingIndexes: MissingIndex[] = collect(doc, "MissingIndexGroup").map((g: any) => {
    const mi = collect(g, "MissingIndex")[0] ?? {};
    const byUsage = (u: string) => collect(mi, "ColumnGroup").filter((c: any) => c.Usage === u).flatMap(colNames);
    return {
      table: fullName(mi),
      equality: byUsage("EQUALITY"), inequality: byUsage("INEQUALITY"),
      included: collect(mi, "IncludedColumns").flatMap(colNames),
      impact: num(g.Impact) ?? 0,
    };
  }).sort((a, b) => b.impact - a.impact);

  // Generic on purpose: SQL Server adds warning elements over time, unknown ones still surface.
  const warnings = [...new Set(collect(doc, "Warnings").flatMap((w: any) =>
    Object.entries(w).filter(([k]) => k !== "#text").map(([k, v]: [string, any]) => {
      const attrs = v && typeof v === "object"
        ? Object.entries(v).filter(([, x]) => typeof x !== "object").map(([ak, x]) => `${ak}=${x}`).join(" ")
        : "";
      return attrs ? `${k} ${attrs}` : k;
    })))];

  const waits: Wait[] = collect(doc, "WaitTime")
    .map((w: any) => ({ type: String(w.WaitType ?? "unknown"), ms: num(w.WaitTime) ?? 0 }))
    .filter((w) => w.ms > 0).sort((a, b) => b.ms - a.ms).slice(0, 5);

  const rowGuessErrors: RowGuess[] = operators
    .filter((o) => o.actualRows !== null && o.estRows > 0)
    .map((o) => ({ id: o.id, op: o.op, est: o.estRows, actual: o.actualRows as number,
      ratio: Math.max(o.estRows, o.actualRows as number) / Math.min(o.estRows, o.actualRows as number) }))
    .filter((r) => r.ratio >= 10).sort((a, b) => b.ratio - a.ratio).slice(0, 5);

  return {
    sql: text(stmt.StatementText) ?? "", bytes: xml.length,
    subtreeCost: num(stmt.StatementSubtreeCost), parallel: operators.some((o) => o.parallel),
    nonParallelReason: text(collect(doc, "NonParallelPlanReason")[0]),
    earlyAbort: text(collect(doc, "StatementOptmEarlyAbortReason")[0]),
    tables, topOperators: operators.sort((a, b) => b.cost - a.cost).slice(0, 12),
    rowGuessErrors, missingIndexes, warnings, waits,
  };
}

// Tokens are the constraint: models and Jev read this, never the raw XML.
export const digestForModel = (d: Digest) => ({
  sql: d.sql.slice(0, 4000), subtreeCost: d.subtreeCost, parallel: d.parallel,
  nonParallelReason: d.nonParallelReason, earlyAbort: d.earlyAbort, tables: d.tables.slice(0, 20),
  topOperators: d.topOperators, rowGuessErrors: d.rowGuessErrors,
  missingIndexes: d.missingIndexes.slice(0, 5), warnings: d.warnings.slice(0, 20), waits: d.waits,
});