From 1ded0a7bcd166950bbc4ad401a5e09ae44acc736 Mon Sep 17 00:00:00 2001
From: t8y2 <1156263951@qq.com>
Date: Sun, 20 Sep 2026 21:55:58 +0800
Subject: [PATCH] feat(diagram): export table and column comments to dbml and
 mermaid

Closes #9777
---
 apps/desktop/src/lib/export/diagramFormats.ts | 34 ++++++++++-
 packages/app-tests/diagramFormats.test.ts     | 60 +++++++++++++++++++
 2 files changed, 93 insertions(+), 1 deletion(-)

diff --git a/apps/desktop/src/lib/export/diagramFormats.ts b/apps/desktop/src/lib/export/diagramFormats.ts
index 95270437ae..1fecad509d 100644
--- a/apps/desktop/src/lib/export/diagramFormats.ts
+++ b/apps/desktop/src/lib/export/diagramFormats.ts
@@ -75,6 +75,28 @@ function mermaidCardinality(sourceCardinality: "1" | "N", targetCardinality: "1"
   return "}o--o{";
 }
 
+/** Comments are optional metadata; blank/whitespace-only values export nothing (same rule as the cards, issue #9748). */
+function commentText(value: string | null | undefined): string {
+  return typeof value === "string" ? value.trim() : "";
+}
+
+/**
+ * Escape a DBML note exactly like the official dbml-core exporter: backslashes
+ * are doubled first, then plain values stay single-quoted while values with
+ * single quotes or line breaks switch to triple quotes with `\'` escapes
+ * (https://docs.dbml.org/docs/annotation).
+ */
+function dbmlNote(value: string): string {
+  let escaped = value.replace(/\\/g, "\\\\");
+  if (!/['\n\r]/.test(escaped)) return `'${escaped}'`;
+  return `'''${escaped.replace(/'/g, "\\'").replace(/\r\n/g, "\n")}'''`;
+}
+
+/** Mermaid attribute comments are single-line and cannot contain double quotes (https://mermaid.js.org/syntax/entityRelationshipDiagram.html). */
+function mermaidCommentText(value: string): string {
+  return value.replace(/"/g, "'").replace(/\s+/g, " ").trim();
+}
+
 export function diagramExportFileName(connectionName: string, databaseName: string, mode: "table" | "engineering", format: DiagramExportFormat): string {
   const context = [connectionName, databaseName].map(fileToken).filter(Boolean);
   const modeSuffix = mode === "engineering" ? "engineering-er" : "table-structure";
@@ -96,9 +118,13 @@ export function buildDiagramDbml(tables: DiagramTable[], relationships: DiagramR
       const attrs: string[] = [];
       if (column.is_primary_key) attrs.push("pk");
       if (!column.is_nullable) attrs.push("not null");
+      const columnComment = commentText(column.comment);
+      if (columnComment) attrs.push(`note: ${dbmlNote(columnComment)}`);
       const attrSuffix = attrs.length > 0 ? ` [${attrs.join(", ")}]` : "";
       lines.push(`  ${quoteIdent(column.name)} ${dbmlType(column.data_type)}${attrSuffix}`);
     }
+    const tableComment = commentText(table.comment);
+    if (tableComment) lines.push(`  Note: ${dbmlNote(tableComment)}`);
     lines.push("}", "");
   }
 
@@ -116,11 +142,17 @@ export function buildDiagramMermaid(tables: DiagramTable[], relationships: Diagr
 
   for (const table of tables) {
     const entity = mermaidIdent(table.name);
+    // erDiagram has no entity-level note syntax, so the table comment rides on
+    // a `%%` comment line (comments must sit on their own line) above the block.
+    const tableComment = commentText(table.comment);
+    if (tableComment) lines.push(`  %% ${mermaidCommentText(table.name)}: ${mermaidCommentText(tableComment)}`);
     lines.push(`  ${entity} {`);
     for (const column of table.columns) {
       const typeToken = (column.data_type || "unknown").replace(/\s+/g, "_").replace(/[^A-Za-z0-9_]/g, "") || "unknown";
       const key = column.is_primary_key ? " PK" : "";
-      lines.push(`    ${typeToken} ${mermaidIdent(column.name)}${key}`);
+      const columnComment = commentText(column.comment);
+      const commentSuffix = columnComment ? ` "${mermaidCommentText(columnComment)}"` : "";
+      lines.push(`    ${typeToken} ${mermaidIdent(column.name)}${key}${commentSuffix}`);
     }
     lines.push("  }");
   }
diff --git a/packages/app-tests/diagramFormats.test.ts b/packages/app-tests/diagramFormats.test.ts
index bd23e7daa3..0fb6367f4f 100644
--- a/packages/app-tests/diagramFormats.test.ts
+++ b/packages/app-tests/diagramFormats.test.ts
@@ -156,6 +156,66 @@ test("buildDiagramRelationships unique FK exports as 1:1 in Mermaid/DBML", () =>
   assert.match(buildDiagramDbml(profileTables, relationships), /Ref: profiles\.user_id - users\.id/);
 });
 
+test("buildDiagramDbml emits table and column notes with dbml-core escaping", () => {
+  const commented: DiagramTable[] = [
+    {
+      name: "users",
+      comment: "应用's 用户\n表",
+      columns: [
+        { name: "id", data_type: "bigint", is_nullable: false, column_default: null, is_primary_key: true, extra: null, comment: "主键ID" },
+        { name: "path", data_type: "varchar(255)", is_nullable: true, column_default: null, is_primary_key: false, extra: null, comment: "it's a\\path" },
+      ],
+      foreignKeys: [],
+    },
+  ];
+
+  const dbml = buildDiagramDbml(commented, []);
+
+  // Plain note: single quotes, no escaping needed; note is the last bracket entry.
+  assert.ok(dbml.includes("  id bigint [pk, not null, note: '主键ID']"));
+  // Quotes/newlines/backslashes: backslashes doubled, ' escaped as \', triple-quoted.
+  assert.ok(dbml.includes(String.raw`  path varchar(255) [note: '''it\'s a\\path''']`));
+  assert.ok(dbml.includes(String.raw`  Note: '''应用\'s 用户`) && dbml.includes("\n表'''\n}"));
+});
+
+test("buildDiagramMermaid appends quoted column comments and %% table comment lines", () => {
+  const commented: DiagramTable[] = [
+    {
+      name: "users",
+      comment: "应用\n用户 表",
+      columns: [
+        { name: "id", data_type: "bigint", is_nullable: false, column_default: null, is_primary_key: true, extra: null, comment: "主键ID" },
+        { name: "path", data_type: "varchar(255)", is_nullable: true, column_default: null, is_primary_key: false, extra: null, comment: "it's a\\path" },
+        { name: "nickname", data_type: "varchar(64)", is_nullable: true, column_default: null, is_primary_key: false, extra: null, comment: '昵称 "外号"' },
+      ],
+      foreignKeys: [],
+    },
+  ];
+
+  const mermaid = buildDiagramMermaid(commented, []);
+
+  // Table comment rides on its own %% line right above the entity block; newlines collapse.
+  assert.ok(mermaid.includes("  %% users: 应用 用户 表\n  users {"));
+  // Column comments are double-quoted, follow the name/keys, and stay single-line.
+  assert.ok(mermaid.includes('    bigint id PK "主键ID"'));
+  assert.ok(mermaid.includes(String.raw`    varchar255 path "it's a\path"`));
+  // Mermaid comments cannot contain double quotes, so they are swapped for single ones.
+  assert.ok(mermaid.includes("    varchar64 nickname \"昵称 '外号'\""));
+});
+
+test("blank comments keep dbml and mermaid output byte-identical", () => {
+  const blanked = tables.map((table) => ({
+    ...table,
+    comment: "   ",
+    columns: table.columns.map((column) => ({ ...column, comment: "" })),
+  }));
+  assert.equal(buildDiagramDbml(blanked, []), buildDiagramDbml(tables, []));
+  assert.equal(buildDiagramMermaid(blanked, []), buildDiagramMermaid(tables, []));
+  const nulled = tables.map((table) => ({ ...table, comment: null, columns: table.columns.map((column) => ({ ...column, comment: null })) }));
+  assert.equal(buildDiagramDbml(nulled, []), buildDiagramDbml(tables, []));
+  assert.equal(buildDiagramMermaid(nulled, []), buildDiagramMermaid(tables, []));
+});
+
 test("diagramExportDialogFilter returns expected extensions", () => {
   assert.deepEqual(diagramExportDialogFilter("svg").extensions, ["svg"]);
   assert.deepEqual(diagramExportDialogFilter("png").extensions, ["png"]);