From 294c0a4ad6d0ba2d3369a43e41552ea5b0acc039 Mon Sep 17 00:00:00 2001 From: zavier <765324639@qq.com> Date: Sat, 1 Aug 2026 21:12:41 +0800 Subject: [PATCH] =?UTF-8?q?refactor(schema):=20getTableSchema=20=E8=BF=94?= =?UTF-8?q?=E5=9B=9E=E7=B1=BB=E5=9E=8B=E5=8C=96=E8=A1=8C=EF=BC=8C=E7=A7=BB?= =?UTF-8?q?=E9=99=A4=20mermaid=20ER=20=E5=9B=BE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 候选 5：getTableSchema 此前把 MySQL information_schema 裸行原样 导出，COLUMN_NAME / COLUMN_TYPE 的知识散布在 4+ 命令层站点，cast 遍地。在 driver 边界一次收窄为 SchemaColumn[]/SchemaIndex[]。 - types.ts 新增 SchemaColumn（nullable 布尔化）/ SchemaIndex（按 INDEX_NAME 聚合）/ TableSchema；裸行不再出边界 - formatSchemaMarkdown 改收类型化参数，cast 与 idxMap 分组逻辑消失 - 移除 mermaid ER 图（pi 不渲染图，用户只能看到源码文本）： - 删 /db relations er-diagram 命令及 completions - discover 的 AI 消息改为分表小节结构化文本（### 表 + 列列表）， 更直接、更省 token、无 mermaid 语法噪音 - commands/relations.ts / schema.ts / db-tools.ts 消费点 cast 全部封口 --- __tests__/db-command.test.ts | 10 +-- __tests__/schema-table.test.ts | 54 +++++++-------- commands/db.ts | 10 ++- commands/relations.ts | 119 ++++----------------------------- commands/schema.ts | 4 +- commands/utils.ts | 2 +- connection/db-manager.ts | 27 +++++++- formatting/schema-table.ts | 25 +++---- state/workspace.ts | 4 +- types.ts | 26 +++++++ 10 files changed, 110 insertions(+), 171 deletions(-) diff --git a/__tests__/db-command.test.ts b/__tests__/db-command.test.ts index 665192e..1221bf6 100644 --- a/__tests__/db-command.test.ts +++ b/__tests__/db-command.test.ts @@ -14,7 +14,7 @@ function stubWs(opts: { isReady?: boolean; tables?: string[] } = {}) { describe("getCompletions (/db 参数补全)", () => { it("expands second-level subcommands for 'relations'", async () => { const result = await getCompletions("relations", stubWs()); - expect(result?.map((c) => c.label).sort()).toEqual(["add", "discover", "er-diagram", "remove"]); + expect(result?.map((c) => c.label).sort()).toEqual(["add", "discover", "remove"]); expect(result?.[0].value).toBe("relations add "); }); @@ -25,7 +25,7 @@ describe("getCompletions (/db 参数补全)", () => { it("handles trailing space after 'relations' (Tab-completed state)", async () => { const result = await getCompletions("relations ", stubWs()); - expect(result?.length).toBe(4); + expect(result?.length).toBe(3); }); it("filters sub-subcommands by partial input", async () => { @@ -57,12 +57,6 @@ describe("getCompletions (/db 参数补全)", () => { expect(result?.map((c) => c.label).sort()).toEqual(["T_ORDERS", "t_customers"]); }); - it("completes table names for 'relations er-diagram'", async () => { - const ws = stubWs({ tables: ["t_orders"] }); - const result = await getCompletions("relations er-diagram ", ws); - expect(result?.map((c) => c.value)).toEqual(["relations er-diagram t_orders"]); - }); - it("falls back to subcommand prefix matching for partial first word", async () => { const result = await getCompletions("s", stubWs()); expect(result?.map((c) => c.label)).toContain("schema"); diff --git a/__tests__/schema-table.test.ts b/__tests__/schema-table.test.ts index e58b059..84bccbe 100644 --- a/__tests__/schema-table.test.ts +++ b/__tests__/schema-table.test.ts @@ -1,31 +1,31 @@ import { describe, it, expect } from "vitest"; import { formatSchemaMarkdown } from "../formatting/schema-table"; -import type { SqlRow } from "../types"; +import type { SchemaColumn, SchemaIndex } from "../types"; -const columns: SqlRow[] = [ +const columns: SchemaColumn[] = [ { - COLUMN_NAME: "id", - COLUMN_TYPE: "bigint(20)", - IS_NULLABLE: "NO", - COLUMN_KEY: "PRI", - COLUMN_DEFAULT: null, - EXTRA: "auto_increment", - COLUMN_COMMENT: "", + name: "id", + type: "bigint(20)", + nullable: false, + key: "PRI", + default: null, + extra: "auto_increment", + comment: "", }, { - COLUMN_NAME: "user_id", - COLUMN_TYPE: "bigint(20)", - IS_NULLABLE: "YES", - COLUMN_KEY: "MUL", - COLUMN_DEFAULT: null, - EXTRA: "", - COLUMN_COMMENT: "下单用户", + name: "user_id", + type: "bigint(20)", + nullable: true, + key: "MUL", + default: null, + extra: "", + comment: "下单用户", }, ]; -const indexes: SqlRow[] = [ - { INDEX_NAME: "PRIMARY", COLUMN_NAME: "id", NON_UNIQUE: 0, SEQ_IN_INDEX: 1 }, - { INDEX_NAME: "idx_user", COLUMN_NAME: "user_id", NON_UNIQUE: 1, SEQ_IN_INDEX: 1 }, +const indexes: SchemaIndex[] = [ + { name: "PRIMARY", columns: ["id"], unique: true }, + { name: "idx_user", columns: ["user_id"], unique: false }, ]; describe("formatSchemaMarkdown", () => { @@ -42,15 +42,15 @@ describe("formatSchemaMarkdown", () => { }); it("escapes pipe characters in comments so the table doesn't break", () => { - const cols: SqlRow[] = [ + const cols: SchemaColumn[] = [ { - COLUMN_NAME: "type", - COLUMN_TYPE: "varchar(10)", - IS_NULLABLE: "YES", - COLUMN_KEY: "", - COLUMN_DEFAULT: null, - EXTRA: "", - COLUMN_COMMENT: "a|b", + name: "type", + type: "varchar(10)", + nullable: true, + key: "", + default: null, + extra: "", + comment: "a|b", }, ]; const out = formatSchemaMarkdown("t", "db", cols, []); diff --git a/commands/db.ts b/commands/db.ts index ce9596c..c8fa960 100644 --- a/commands/db.ts +++ b/commands/db.ts @@ -176,7 +176,7 @@ export async function getCompletions( const subSubs: Record = { favorite: ["add"], - relations: ["add", "remove", "discover", "er-diagram"], + relations: ["add", "remove", "discover"], }; // 当第一个词与拥有子子命令的子命令完全匹配时，立即显示第二层。 @@ -190,14 +190,12 @@ export async function getCompletions( // 表名参数（schema、query）——在第一级部分匹配之前触发， // 避免精确子命令匹配自引用。 - const takesTable = - sub === "schema" || sub === "query" || (sub === "relations" && parts[1] === "er-diagram"); + const takesTable = sub === "schema" || sub === "query"; if (takesTable && ws.isReady) { try { const tables = await ws.getTables(); - const tablePartial = - sub === "relations" ? (hasTrailingSpace ? "" : (parts[2] ?? "")) : partial; - const valuePrefix = sub === "relations" ? "relations er-diagram " : `${sub} `; + const tablePartial = partial; + const valuePrefix = `${sub} `; return tables .filter((t) => t.toLowerCase().startsWith(tablePartial.toLowerCase())) .map((t) => ({ value: `${valuePrefix}${t}`, label: t })); diff --git a/commands/relations.ts b/commands/relations.ts index 75f1f99..3176b07 100644 --- a/commands/relations.ts +++ b/commands/relations.ts @@ -1,12 +1,12 @@ /** * /db relations —— 表关系管理。 * - * 子命令：add、remove、discover（FK 同步）、er-diagram。 + * 子命令：add、remove、discover（FK 同步）。 */ import type { ExtensionAPI, ExtensionCommandContext } from "@earendil-works/pi-coding-agent"; import type { DatabaseWorkspaceService } from "../state/workspace"; -import type { SqlRow, StoredRelation } from "../types"; +import type { StoredRelation } from "../types"; import { pickTableFuzzy, withLoader } from "./utils"; // ── 列表格式化 ───────────────────────────────────────────── @@ -43,8 +43,6 @@ export async function handleRelations( return await handleRelationsRemove(ctx, ws, pi); case "discover": return await handleRelationsDiscover(ctx, ws, pi); - case "er-diagram": - return await handleRelationsERDiagram(ctx, ws, pi, rest[1]); default: return await handleRelationsList(ctx, ws, pi); } @@ -115,7 +113,7 @@ async function handleRelationsAdd( let srcColumns: string[] = []; try { const schemaInfo = await ws.getTableSchema(srcTable); - srcColumns = schemaInfo.columns.map((c: SqlRow) => c.COLUMN_NAME as string); + srcColumns = schemaInfo.columns.map((c) => c.name); } catch { ctx.ui.notify(`无法获取 ${srcTable} 的列信息`, "error"); return; @@ -130,7 +128,7 @@ async function handleRelationsAdd( let refColumns: string[] = []; try { const schemaInfo = await ws.getTableSchema(refTable); - refColumns = schemaInfo.columns.map((c: SqlRow) => c.COLUMN_NAME as string); + refColumns = schemaInfo.columns.map((c) => c.name); } catch { ctx.ui.notify(`无法获取 ${refTable} 的列信息`, "error"); return; @@ -235,31 +233,25 @@ async function handleRelationsDiscover( } if (tables.length > 0) { - const erLines: string[] = ["erDiagram"]; const MAX_TABLES = 30; const sampleTables = tables.slice(0, MAX_TABLES); + const schemaBlocks: string[] = []; for (const t of sampleTables) { try { const info = await ws.getTableSchema(t); - erLines.push(` "${t}" {`); - for (const col of info.columns) { - const colName = col.COLUMN_NAME as string; - const colType = col.COLUMN_TYPE as string; - const comment = col.COLUMN_COMMENT ? ` "${col.COLUMN_COMMENT}"` : ""; - erLines.push(` ${colType} ${colName}${comment}`); - } - erLines.push(` }`); + const cols = info.columns.map((c) => `- ${c.name} (${c.type})`).join("\n"); + schemaBlocks.push(`### ${t}\n${cols}`); } catch { // 跳过 } } if (tables.length > MAX_TABLES) { - erLines.push(` "…还有${tables.length - MAX_TABLES}张表" {}`); + schemaBlocks.push(`…还有 ${tables.length - MAX_TABLES} 张表未展示`); } - const erDiagram = erLines.join("\n"); + const schemaText = schemaBlocks.join("\n\n"); ctx.ui.notify("正在通过 AI 分析表关系…", "info"); @@ -267,7 +259,7 @@ async function handleRelationsDiscover( { customType: "db-relation-discover", content: [ - `请分析以下数据库 ${schema} 的 mermaid ER 图，找出表之间可能的关联关系。`, + `请分析以下数据库 ${schema} 的表结构，找出表之间可能的关联关系。`, ``, `规则：`, `1. 根据列名匹配（如 users.id ↔ orders.user_id, dept_no ↔ dept_no）`, @@ -280,10 +272,8 @@ async function handleRelationsDiscover( `仅当 db_tools 也不可用时，再以 JSON 数组格式输出，每个元素：`, `{"table":"源表","column":"源列","refTable":"目标表","refColumn":"目标列","relationType":"MANY_TO_ONE","condition":""}`, ``, - `ER 图：`, - "```mermaid", - erDiagram, - "```", + `表结构：`, + schemaText, ].join("\n"), display: true, }, @@ -294,88 +284,3 @@ async function handleRelationsDiscover( ctx.ui.notify(parts.join("\n"), "info"); } - -// ── ER 图 ────────────────────────────────────────────────── - -async function handleRelationsERDiagram( - ctx: ExtensionCommandContext, - ws: DatabaseWorkspaceService, - pi: ExtensionAPI, - table?: string, -): Promise { - if (!ws.isReady) { - ctx.ui.notify("未选择数据库，请先执行 /db switch", "warning"); - return; - } - - if (!table) { - const picked = await pickTableFuzzy(ctx, ws, "选择表"); - if (!picked) return; - table = picked; - } - - let tableColumns: SqlRow[] = []; - try { - const info = await ws.getTableSchema(table); - tableColumns = info.columns; - } catch { - ctx.ui.notify(`无法获取 ${table} 的表结构`, "error"); - return; - } - - const relations = ws.listRelations(table); - const relatedTableNames = new Set(); - for (const r of relations) { - relatedTableNames.add(r.refTable); - relatedTableNames.add(r.table); - } - - const allColumns = new Map(); - allColumns.set(table, tableColumns); - for (const relatedTable of relatedTableNames) { - if (relatedTable === table) continue; - try { - const info = await ws.getTableSchema(relatedTable); - allColumns.set(relatedTable, info.columns); - } catch { - // 跳过 - } - } - - const lines: string[] = ["erDiagram"]; - - for (const r of relations) { - const label = r.condition - ? `${r.column} → ${r.refColumn} [${r.condition}]` - : `${r.column} → ${r.refColumn}`; - lines.push(` "${r.table}" ||--o{ "${r.refTable}" : "${label}"`); - } - - const drawn = new Set(); - for (const [tbl, cols] of allColumns) { - if (drawn.has(tbl)) continue; - drawn.add(tbl); - lines.push(` "${tbl}" {`); - for (const col of cols) { - const colName = col.COLUMN_NAME as string; - const colType = col.COLUMN_TYPE as string; - const comment = col.COLUMN_COMMENT ? ` "${col.COLUMN_COMMENT}"` : ""; - lines.push(` ${colType} ${colName}${comment}`); - } - lines.push(` }`); - } - - const erDiagram = lines.join("\n"); - - // display: true → 在聊天中持久显示；默认 markdown 渲染器 - // 把 mermaid 源码显示为代码块，LLM 也能读取。 - // deliverAs "followUp" 在 agent 空闲时立即提交。 - pi.sendMessage( - { - customType: "db-er-diagram", - content: [`## ER 图 — ${table}`, "", "```mermaid", erDiagram, "```"].join("\n"), - display: true, - }, - { deliverAs: "followUp", triggerTurn: false }, - ); -} diff --git a/commands/schema.ts b/commands/schema.ts index 6d9fa01..0afb9b9 100644 --- a/commands/schema.ts +++ b/commands/schema.ts @@ -4,7 +4,7 @@ import type { ExtensionAPI, ExtensionCommandContext } from "@earendil-works/pi-coding-agent"; import type { DatabaseWorkspaceService } from "../state/workspace"; -import type { SqlRow } from "../types"; +import type { TableSchema } from "../types"; import { formatSchemaMarkdown } from "../formatting/schema-table"; import { pickTableFuzzy } from "./utils"; @@ -25,7 +25,7 @@ export async function handleSchema( table = picked; } - let result: { columns: SqlRow[]; indexes: SqlRow[] }; + let result: TableSchema; try { result = await ws.getTableSchema(table); } catch (err: any) { diff --git a/commands/utils.ts b/commands/utils.ts index 2743b3c..569709a 100644 --- a/commands/utils.ts +++ b/commands/utils.ts @@ -21,7 +21,7 @@ import { createFilterReducer } from "./filter-input"; * @param extraItems —— 前置在表项之前（如 "✏️ 输入 SQL…"） * 用于合并的查询入口）。 * - * 由 schema、query、relations-add 和 relations-er-diagram 处理器共享。 + * 由 schema、query、relations-add 处理器共享。 */ export async function pickTableFuzzy( ctx: ExtensionCommandContext, diff --git a/connection/db-manager.ts b/connection/db-manager.ts index 85a4721..72c9694 100644 --- a/connection/db-manager.ts +++ b/connection/db-manager.ts @@ -9,6 +9,7 @@ import mysql, { type Pool, type RowDataPacket, type ResultSetHeader } from "mysql2/promise"; import type { ResolvedConnectionConfig } from "./db-config"; import { prepareReadOnlyQuery, DEFAULT_QUERY_LIMIT } from "./sql-policy"; +import type { SchemaIndex, TableSchema } from "../types"; export interface QueryOptions { limit?: number; // row cap for SELECTs without trailing LIMIT (default: connection's queryLimit) @@ -107,7 +108,7 @@ export class DatabaseConnectionManager { connectionId: string, database: string, table: string, - ): Promise<{ columns: RowDataPacket[]; indexes: RowDataPacket[] }> { + ): Promise { const pool = this.getPool(connectionId); const [columns] = await pool.query( @@ -126,7 +127,29 @@ export class DatabaseConnectionManager { [database, table], ); - return { columns, indexes }; + return { + columns: columns.map((c) => ({ + name: c.COLUMN_NAME as string, + type: c.COLUMN_TYPE as string, + nullable: c.IS_NULLABLE === "YES", + key: (c.COLUMN_KEY as string) ?? "", + default: (c.COLUMN_DEFAULT as string | null) ?? null, + extra: (c.EXTRA as string) ?? "", + comment: (c.COLUMN_COMMENT as string) ?? "", + })), + indexes: this.aggregateIndexes(indexes), + }; + } + + /** 将 information_schema.STATISTICS 行按索引名聚合成 SchemaIndex[]。 */ + private aggregateIndexes(rows: RowDataPacket[]): SchemaIndex[] { + const map = new Map(); + for (const idx of rows) { + const name = idx.INDEX_NAME as string; + if (!map.has(name)) map.set(name, { cols: [], unique: idx.NON_UNIQUE === 0 }); + map.get(name)!.cols.push(idx.COLUMN_NAME as string); + } + return [...map.entries()].map(([name, { cols, unique }]) => ({ name, columns: cols, unique })); } /** diff --git a/formatting/schema-table.ts b/formatting/schema-table.ts index 445a8e9..1d5c8d5 100644 --- a/formatting/schema-table.ts +++ b/formatting/schema-table.ts @@ -2,15 +2,15 @@ * schema markdown 格式化 —— 由 /db schema 命令和 db_tables LLM 工具共享的纯函数。 */ -import type { SqlRow } from "../types"; +import type { SchemaColumn, SchemaIndex } from "../types"; /** 转义竖线，避免列注释破坏 markdown 表格。 */ -function esc(val: unknown): string { +function esc(val: string | null): string { const s = String(val ?? ""); return s.replace(/\|/g, "\\|"); } -function keyLabel(key: unknown): string { +function keyLabel(key: string): string { switch (key) { case "PRI": return "PK"; @@ -26,30 +26,23 @@ function keyLabel(key: unknown): string { export function formatSchemaMarkdown( table: string, database: string, - columns: SqlRow[], - indexes: SqlRow[], + columns: SchemaColumn[], + indexes: SchemaIndex[], ): string { const lines: string[] = [`### ${table} — ${database}`, ""]; lines.push("| 列 | 类型 | Null | Key | 默认 | Extra | 注释 |"); lines.push("| --- | --- | --- | --- | --- | --- | --- |"); for (const c of columns) { - const nullable = c.IS_NULLABLE === "YES" ? "YES" : ""; + const nullable = c.nullable ? "YES" : ""; lines.push( - `| ${esc(c.COLUMN_NAME)} | ${esc(c.COLUMN_TYPE)} | ${nullable} | ${keyLabel(c.COLUMN_KEY)} | ${esc(c.COLUMN_DEFAULT)} | ${esc(c.EXTRA)} | ${esc(c.COLUMN_COMMENT)} |`, + `| ${esc(c.name)} | ${esc(c.type)} | ${nullable} | ${keyLabel(c.key)} | ${esc(c.default)} | ${esc(c.extra)} | ${esc(c.comment)} |`, ); } - const idxMap = new Map(); + lines.push("", `**索引（${indexes.length}）**`, ""); for (const idx of indexes) { - const name = idx.INDEX_NAME as string; - if (!idxMap.has(name)) idxMap.set(name, { cols: [], unique: idx.NON_UNIQUE === 0 }); - idxMap.get(name)!.cols.push(idx.COLUMN_NAME as string); - } - - lines.push("", `**索引（${idxMap.size}）**`, ""); - for (const [name, { cols, unique }] of idxMap) { - lines.push(`- \`${name}\`${unique ? " [UNIQUE]" : ""}: ${cols.join(", ")}`); + lines.push(`- \`${idx.name}\`${idx.unique ? " [UNIQUE]" : ""}: ${idx.columns.join(", ")}`); } return lines.join("\n"); diff --git a/state/workspace.ts b/state/workspace.ts index bd0fc09..485168f 100644 --- a/state/workspace.ts +++ b/state/workspace.ts @@ -24,7 +24,7 @@ import { type HistoryFilter, } from "../history/store"; import { RelationGraph } from "../relation-graph"; -import type { RelatedResult, SqlRow, StoredRelation } from "../types"; +import type { RelatedResult, SqlRow, StoredRelation, TableSchema } from "../types"; import { StateStore } from "./state-store"; // ====== 内部类型 ====== @@ -306,7 +306,7 @@ export class DatabaseWorkspaceService { async getTableSchema( table: string, opts?: { connectionId?: string; database?: string }, - ): Promise<{ columns: SqlRow[]; indexes: SqlRow[] }> { + ): Promise { const target = this.resolveTarget(opts); return this.manager.getTableSchema(target.connectionId, target.database, table); } diff --git a/types.ts b/types.ts index 62d4b36..d7c52b6 100644 --- a/types.ts +++ b/types.ts @@ -44,3 +44,29 @@ export interface RelatedResult { joinPath: string; elapsed: string; } + +// ====== schema 类型 ====== + +/** information_schema 收窄后的列定义——driver 边界一次映射。 */ +export interface SchemaColumn { + name: string; + type: string; + nullable: boolean; + key: string; + default: string | null; + extra: string; + comment: string; +} + +/** 收窄并聚合后的索引定义（按 INDEX_NAME 分组）。 */ +export interface SchemaIndex { + name: string; + columns: string[]; + unique: boolean; +} + +/** getTableSchema 的返回类型。 */ +export interface TableSchema { + columns: SchemaColumn[]; + indexes: SchemaIndex[]; +}