From 09d54d09bc0b73f00b024b52fb625e75d5011796 Mon Sep 17 00:00:00 2001
From: t8y2 <1156263951@qq.com>
Date: Mon, 22 Jun 2026 16:42:38 +0800
Subject: [PATCH] feat(mongodb): support use command to switch database

---
 apps/desktop/src/lib/mongoShellCommand.ts | 22 ++++++++++++++++++
 apps/desktop/src/stores/queryStore.ts     | 28 +++++++++++++++++++++++
 2 files changed, 50 insertions(+)

diff --git a/apps/desktop/src/lib/mongoShellCommand.ts b/apps/desktop/src/lib/mongoShellCommand.ts
index 6b491e98b8..7d6bc34974 100644
--- a/apps/desktop/src/lib/mongoShellCommand.ts
+++ b/apps/desktop/src/lib/mongoShellCommand.ts
@@ -22,6 +22,10 @@ export interface MongoGetIndexesCommand {
   collection: string;
 }
 
+export interface MongoUseCommand {
+  database: string;
+}
+
 export type MongoWriteCommand = { kind: "insert"; collection: string; docsJson: string } | { kind: "update"; collection: string; filter: string; update: string; many: boolean } | { kind: "delete"; collection: string; filter: string; many: boolean };
 
 export interface MongoAggregateSafetyOptions {
@@ -130,6 +134,15 @@ export function parseMongoGetIndexesCommand(input: string): MongoGetIndexesComma
   };
 }
 
+export function parseMongoUseCommand(input: string): MongoUseCommand | null {
+  const source = input.trim().replace(/;$/, "").trim();
+  const match = /^use\s+([a-zA-Z0-9_-]+)$/i.exec(source);
+  if (!match) return null;
+  return {
+    database: match[1],
+  };
+}
+
 export function parseMongoWriteCommand(input: string): MongoWriteCommand | null {
   const source = input.trim().replace(/;$/, "").trim();
   const insertOne = parseCollectionMethodTarget(source, "insertOne");
@@ -251,6 +264,15 @@ export function mongoWriteToQueryResult(affectedRows: number, executionTimeMs: n
   };
 }
 
+export function mongoUseToQueryResult(database: string, executionTimeMs: number): QueryResult {
+  return {
+    columns: ["message"],
+    rows: [[`switched to db ${database}`]],
+    affected_rows: 0,
+    execution_time_ms: Math.max(0, Math.round(executionTimeMs)),
+  };
+}
+
 export function mongoIndexesToQueryResult(
   indexes: {
     name: string;
diff --git a/apps/desktop/src/stores/queryStore.ts b/apps/desktop/src/stores/queryStore.ts
index 8da38dcd25..225b20b1f4 100644
--- a/apps/desktop/src/stores/queryStore.ts
+++ b/apps/desktop/src/stores/queryStore.ts
@@ -15,11 +15,13 @@ import {
   mongoDocumentsToQueryResult,
   mongoIndexesToQueryResult,
   mongoWriteToQueryResult,
+  mongoUseToQueryResult,
   parseMongoAggregateCommand,
   parseMongoCountDocumentsCommand,
   parseMongoFindCommand,
   parseMongoGetIndexesCommand,
   parseMongoWriteCommand,
+  parseMongoUseCommand,
   type MongoAggregateSafetyOptions,
 } from "@/lib/mongoShellCommand";
 import { redisCommandResultToQueryResult } from "@/lib/redisQueryResult";
@@ -1538,6 +1540,32 @@ export const useQueryStore = defineStore("query", () => {
         return;
       }
 
+      const mongoUse = conn?.db_type === "mongodb" ? parseMongoUseCommand(sql) : null;
+      if (mongoUse) {
+        console.info("[DBX][executeTabSql:mongo-use:start]", { traceId, database: mongoUse.database });
+        const current = tabs.value.find((t) => t.id === id);
+        if (current?.executionId === executionId) {
+          current.database = mongoUse.database;
+          current.results = undefined;
+          current.activeResultIndex = undefined;
+          current.result = markQueryResultRowsRaw(mongoUseToQueryResult(mongoUse.database, performance.now() - startedAt));
+          touchResult(current);
+          current.queryAnalysis = undefined;
+          current.querySourceColumns = undefined;
+          current.queryEditabilityReason = undefined;
+          current.tableMeta = undefined;
+          current.resultBaseSql = options?.resultBaseSql ?? sql;
+          current.resultSortedSql = options?.resultSortedSql;
+          syncDisplayedResultRun(current, options?.resultBaseSql ?? sql);
+        }
+        console.info("[DBX][executeTabSql:mongo-use:done]", {
+          traceId,
+          database: mongoUse.database,
+          elapsed: elapsed(),
+        });
+        return;
+      }
+
       console.info("[DBX][executeTabSql:execute-multi:start]", { traceId, elapsed: elapsed() });
       const clientSessionId = tab.mode === "query" ? tab.id : undefined;
       const executionOptions = {