From 9a16aeaf8b29b5043b910a8f8b6d6220af8643e2 Mon Sep 17 00:00:00 2001
From: t8y2 <1156263951@qq.com>
Date: Tue, 11 Aug 2026 13:04:28 +0000
Subject: [PATCH] fix(mongodb): separate formatted shell statements

Closes #5909
---
 .../__tests__/mongo/mongoFormatter.spec.ts    | 71 +++++++++++++++++++
 apps/desktop/src/lib/mongo/mongoFormatter.ts  | 50 +++++++++++--
 2 files changed, 116 insertions(+), 5 deletions(-)

diff --git a/apps/desktop/src/lib/__tests__/mongo/mongoFormatter.spec.ts b/apps/desktop/src/lib/__tests__/mongo/mongoFormatter.spec.ts
index e747e81c04..7877ff054c 100644
--- a/apps/desktop/src/lib/__tests__/mongo/mongoFormatter.spec.ts
+++ b/apps/desktop/src/lib/__tests__/mongo/mongoFormatter.spec.ts
@@ -2,6 +2,77 @@ import { describe, expect, it } from "vitest";
 import { formatMongoShellText, MAX_MONGO_FORMAT_CHARS } from "@/lib/mongo/mongoFormatter";
 
 describe("mongoFormatter", () => {
+  it("separates top-level Mongo shell statements", () => {
+    const query = `db.order_full_link_trace_log.find({
+  traceId:"master-order:63907695457"
+});
+
+db.order_full_link_trace_log.find({}).sort({
+  masterOrderId:-1
+});
+
+db.order_full_link_trace_log.find({
+  masterOrderId:63907756017
+});`;
+
+    expect(formatMongoShellText(query)).toBe(`db.order_full_link_trace_log.find({
+  traceId: "master-order:63907695457"
+});
+
+db.order_full_link_trace_log.find({})
+  .sort({
+    masterOrderId: -1
+  });
+
+db.order_full_link_trace_log.find({
+  masterOrderId: 63907756017
+});`);
+  });
+
+  it("normalizes existing statement spacing and remains idempotent", () => {
+    const expected = "db.first.find({});\n\ndb.second.find({});";
+
+    expect(formatMongoShellText("db.first.find({}); db.second.find({});")).toBe(expected);
+    expect(formatMongoShellText("db.first.find({});\n\n\n db.second.find({});")).toBe(expected);
+    expect(formatMongoShellText(expected)).toBe(expected);
+  });
+
+  it("leaves a single statement unchanged", () => {
+    expect(formatMongoShellText("db.items.find({});")).toBe("db.items.find({});");
+  });
+
+  it("does not treat semicolons inside nested syntax or literals as statement boundaries", () => {
+    const formatted = formatMongoShellText(`db.items.find({
+  text:"a;b",
+  pattern:/a;b/,
+  $where:function(){const ids=[1,2];return ids.length>0;}
+});db.logs.find({});`);
+
+    expect(formatted).toContain('text: "a;b"');
+    expect(formatted).toContain("pattern: /a;b/");
+    expect(formatted).toContain("const ids=[");
+    expect(formatted).toContain("];return ids.length>0;");
+    expect(formatted.split("\n\n")).toHaveLength(2);
+  });
+
+  it("keeps trailing comments attached to their statement", () => {
+    const formatted = formatMongoShellText("db.first.find({}); // keep; this comment\ndb.second.find({}); /* keep; this too */ db.third.find({});");
+
+    expect(formatted).toBe("db.first.find({}); // keep; this comment\n\ndb.second.find({}); /* keep; this too */\n\ndb.third.find({});");
+  });
+
+  it("keeps the statement after a standalone line comment executable", () => {
+    const expected = "db.first.find({});\n\n// comment for next statement\ndb.second.find({});";
+
+    expect(formatMongoShellText("db.first.find({});\n// comment for next statement\ndb.second.find({});")).toBe(expected);
+    expect(formatMongoShellText("db.first.find({});\r\n// comment for next statement\r\ndb.second.find({});")).toBe(expected);
+    expect(formatMongoShellText(expected)).toBe(expected);
+  });
+
+  it("does not add a line break after a final line comment", () => {
+    expect(formatMongoShellText("db.first.find({}); // final comment")).toBe("db.first.find({}); // final comment");
+  });
+
   it("formats documents with many short fields without repeated whole-output rewrites", () => {
     const fields = Array.from({ length: 20_000 }, (_, index) => `"field${index}":${index}`).join(",");
     const query = `db.items.insert({${fields}});`;
diff --git a/apps/desktop/src/lib/mongo/mongoFormatter.ts b/apps/desktop/src/lib/mongo/mongoFormatter.ts
index d49291a9b5..6a5c763d1d 100644
--- a/apps/desktop/src/lib/mongo/mongoFormatter.ts
+++ b/apps/desktop/src/lib/mongo/mongoFormatter.ts
@@ -9,6 +9,7 @@ interface FormatState {
   indentLevel: number;
   atLineStart: boolean;
   pendingSpace: boolean;
+  pendingStatementBreak: boolean;
   chainIndent: boolean;
   pendingChainCall: boolean;
   stack: Array<{ char: string; expanded: boolean; chainCall?: boolean }>;
@@ -23,12 +24,37 @@ export function formatMongoShellText(text: string, settings: Partial<SqlFormatte
   }
 
   const indentUnit = settings.useTabs ? "\t" : " ".repeat(settings.tabWidth ?? DEFAULT_SQL_FORMATTER_SETTINGS.tabWidth);
-  const state: FormatState = { out: [], lastChar: "", lastNonWhitespace: "", indentLevel: 0, atLineStart: true, pendingSpace: false, chainIndent: false, pendingChainCall: false, stack: [] };
+  const state: FormatState = { out: [], lastChar: "", lastNonWhitespace: "", indentLevel: 0, atLineStart: true, pendingSpace: false, pendingStatementBreak: false, chainIndent: false, pendingChainCall: false, stack: [] };
 
   for (let index = 0; index < text.length; index++) {
     const char = text[index] ?? "";
     const next = text[index + 1] ?? "";
 
+    if (state.pendingStatementBreak) {
+      if (/\s/.test(char)) {
+        if (char === "\n" || char === "\r") separateStatements(state, indentUnit);
+        else state.pendingSpace = !state.atLineStart;
+        continue;
+      }
+
+      if (char === "/" && next === "/") {
+        const comment = readLineComment(text, index);
+        appendToken(state, comment.value, indentUnit);
+        if (comment.hasLineBreak) newline(state, indentUnit);
+        index = comment.end;
+        continue;
+      }
+
+      if (char === "/" && next === "*") {
+        const comment = readBlockComment(text, index);
+        appendToken(state, comment.value, indentUnit);
+        index = comment.end;
+        continue;
+      }
+
+      separateStatements(state, indentUnit);
+    }
+
     if (char === '"' || char === "'" || char === "`") {
       const literal = readQuotedLiteral(text, index, char);
       appendToken(state, literal.value, indentUnit);
@@ -39,6 +65,7 @@ export function formatMongoShellText(text: string, settings: Partial<SqlFormatte
     if (char === "/" && next === "/") {
       const comment = readLineComment(text, index);
       appendToken(state, comment.value, indentUnit);
+      if (comment.hasLineBreak) newline(state, indentUnit);
       index = comment.end;
       continue;
     }
@@ -110,6 +137,12 @@ export function formatMongoShellText(text: string, settings: Partial<SqlFormatte
       continue;
     }
 
+    if (char === ";" && state.stack.length === 0) {
+      appendToken(state, char, indentUnit);
+      state.pendingStatementBreak = true;
+      continue;
+    }
+
     appendToken(state, char, indentUnit);
   }
 
@@ -148,6 +181,12 @@ function newline(state: FormatState, indentUnit: string, indentDelta = 0) {
   void indentUnit;
 }
 
+function separateStatements(state: FormatState, indentUnit: string) {
+  newline(state, indentUnit);
+  appendOutput(state, "\n");
+  state.pendingStatementBreak = false;
+}
+
 function shouldInsertPendingSpace(previous: string, token: string): boolean {
   if (!previous) return false;
   if ([".", "(", "[", "{"].includes(previous)) return false;
@@ -194,10 +233,11 @@ function readQuotedLiteral(text: string, start: number, quote: string): { value:
   return { value: text.slice(start), end: text.length - 1 };
 }
 
-function readLineComment(text: string, start: number): { value: string; end: number } {
-  const end = text.indexOf("\n", start);
-  if (end < 0) return { value: text.slice(start), end: text.length - 1 };
-  return { value: text.slice(start, end), end: end - 1 };
+function readLineComment(text: string, start: number): { value: string; end: number; hasLineBreak: boolean } {
+  let end = start + 2;
+  while (end < text.length && text[end] !== "\n" && text[end] !== "\r") end++;
+  if (end === text.length) return { value: text.slice(start), end: text.length - 1, hasLineBreak: false };
+  return { value: text.slice(start, end), end: end - 1, hasLineBreak: true };
 }
 
 function readBlockComment(text: string, start: number): { value: string; end: number } {