feat(mongodb): complete database names after use and getSiblingDB

apps/desktop/src/components/editor/useQueryEditorCompletion.ts

@@ -20,7 +20,7 @@ import {
 import { originForSqlCompletionProvider, originForTypedSqlCompletionStart, shouldAllowSqlCompletionTrigger, type SqlCompletionTriggerFacts, type SqlCompletionTriggerOrigin } from "@/lib/sql/sqlCompletionTriggerPolicy";
 import { driverProfileHasCompletionCandidates } from "@/lib/database/driverProfileExtensions";
 import { buildElasticsearchCompletionItemsFromContext, elasticsearchCompletionNeedsFields, getElasticsearchCompletionContext, getElasticsearchCompletionResultValidFor, shouldAutoOpenElasticsearchCompletion, type ElasticsearchCompletionField } from "@/lib/elasticsearch/elasticsearchCompletion";
-import { buildMongoCompletionItemsFromContext, getMongoCompletionContext, getMongoCompletionResultValidFor, mongoCompletionNeedsCollections, mongoCompletionNeedsFields, shouldAutoOpenMongoCompletion } from "@/lib/mongo/mongoCompletion";
+import { buildMongoCompletionItemsFromContext, getMongoCompletionContext, getMongoCompletionResultValidFor, mongoCompletionNeedsCollections, mongoCompletionNeedsDatabases, mongoCompletionNeedsFields, shouldAutoOpenMongoCompletion } from "@/lib/mongo/mongoCompletion";
 import { buildSoqlCompletionItems, getSoqlCompletionContext, getSoqlCompletionResultValidFor, resolveSoqlFieldCandidates, resolveSoqlValueField, shouldAutoOpenSoqlCompletion, soqlCompletionNeedsObjects, type SoqlCompletionField, type SoqlCompletionObject } from "@/lib/soql/soqlCompletion";
 import {
   buildSqlServerUseDatabaseCompletionItems,
@@ -565,9 +565,19 @@ export function useQueryEditorCompletion(options: QueryEditorCompletionOptions)
     if (!explicit && !shouldAutoOpenMongoCompletion(fullDoc, position)) return null;
 
     const completionContext = getMongoCompletionContext(fullDoc, position);
+    let databases: string[] = [];
     let collections: string[] = [];
     let fields: Awaited<ReturnType<typeof connectionStore.listMongoCompletionFields>> = [];
 
+    // `use` and `getSiblingDB` name a database, which does not depend on the tab having one selected.
+    if (mongoCompletionNeedsDatabases(completionContext.mode)) {
+      try {
+        databases = await connectionStore.listCompletionDatabases(props.connectionId);
+      } catch {
+        databases = [];
+      }
+    }
+
     if (props.database && mongoCompletionNeedsCollections(completionContext.mode)) {
       try {
         collections = await connectionStore.listMongoCompletionCollections(props.connectionId, completionContext.database ?? props.database);
@@ -587,6 +597,7 @@ export function useQueryEditorCompletion(options: QueryEditorCompletionOptions)
     if (epoch !== completionEpoch) return null;
 
     const items = buildMongoCompletionItemsFromContext(completionContext, {
+      databases,
       collections,
       fields,
     });

apps/desktop/src/lib/mongo/__tests__/mongoSiblingDatabaseCompletion.spec.ts

@@ -37,6 +37,15 @@ describe("getSiblingDB completion targeting", () => {
     expect(labels("db.")).toContain("getSiblingDB");
   });
 
+  it("completes the database name inside the getSiblingDB argument", () => {
+    const databases = ["archive", "analytics", "admin"];
+    const items = (text: string) => buildMongoCompletionItems(text, text.length, { databases }).map((item) => item.label);
+    expect(getMongoCompletionContext('db.getSiblingDB("ar', 19)).toMatchObject({ mode: "database", prefix: '"ar', from: 16 });
+    expect(items('db.getSiblingDB("ar')).toEqual(["archive"]);
+    expect(items('db.getSiblingDB("')).toEqual(databases);
+    expect(buildMongoCompletionItems('db.getSiblingDB("ar', 19, { databases })[0]?.apply).toBe('"archive"');
+  });
+
   it("does not treat a variable argument as a resolved sibling database", () => {
     const text = "db.getSiblingDB(name).";
     expect(getMongoCompletionContext(text, text.length).database).toBeUndefined();

apps/desktop/src/lib/mongo/mongoCompletion.ts

@@ -33,6 +33,7 @@ import {
 export type MongoCompletionMode =
   | "none"
   | "root"
+  | "database"
   | "collection"
   | "collectionOrMethod"
   | "collectionRef"
@@ -93,6 +94,7 @@ export interface MongoCompletionContext {
 }
 
 export interface MongoCompletionInput {
+  databases?: string[];
   collections?: string[];
   fields?: MongoCompletionField[];
 }
@@ -270,8 +272,21 @@ export function getMongoCompletionContext(text: string, cursor: number): MongoCo
 
   if (isInsideMongoComment(text, safeCursor)) return { mode: "none", prefix: "", from: safeCursor };
 
+  const usePrefix = matchUseDatabasePrefix(beforeCursor);
+  if (usePrefix) return { mode: "database", prefix: usePrefix.prefix, from: usePrefix.from };
+
   if (endsAtDbRootDot(beforeCursor)) return { mode: "collection", prefix: "", from: safeCursor, collection, database };
 
+  const getSiblingDbPrefix = matchGetSiblingDbPrefix(beforeCursor);
+  if (getSiblingDbPrefix) {
+    return {
+      mode: "database",
+      prefix: getSiblingDbPrefix.prefix,
+      from: getSiblingDbPrefix.from,
+      replaceClosingQuote: closingQuoteAtCursor(getSiblingDbPrefix.prefix, text, safeCursor),
+    };
+  }
+
   const getCollectionPrefix = matchGetCollectionPrefix(beforeCursor);
   if (getCollectionPrefix) {
     return {
@@ -339,6 +354,9 @@ export function buildMongoCompletionItemsFromContext(context: MongoCompletionCon
     case "root":
       items = rootItems(prefix);
       break;
+    case "database":
+      items = databaseItems(prefix, input.databases ?? []);
+      break;
     case "collection":
       items = collectionItems(prefix, collections, context.database !== undefined);
       break;
@@ -424,10 +442,17 @@ export function mongoCompletionNeedsCollections(mode: MongoCompletionMode): bool
   return mode === "collection" || mode === "collectionOrMethod" || mode === "collectionRef";
 }
 
+/** Modes whose items are built from the connection's database names. */
+export function mongoCompletionNeedsDatabases(mode: MongoCompletionMode): boolean {
+  return mode === "database";
+}
+
 export function shouldAutoOpenMongoCompletion(text: string, cursor: number): boolean {
   const previousChar = text[cursor - 1];
   if (!previousChar) return false;
   if (text.slice(0, cursor).endsWith("db.")) return true;
+  // `use ` names a database next; open the list as soon as the space is typed.
+  if (previousChar === " " && matchUseDatabasePrefix(text.slice(0, cursor))) return true;
   if (previousChar === "$" || previousChar === "." || previousChar === '"' || previousChar === "'") return true;
   if (/[{,[:]/.test(previousChar) || /[{,[:]\s+$/.test(text.slice(0, cursor))) {
     return getMongoCompletionContext(text, cursor).mode !== "none";
@@ -983,6 +1008,24 @@ function collectionOrMethodItems(prefix: string, collections: string[]): MongoCo
   return dedupeAndSort([...collectionNameItems(prefix, collections, 150), ...methods]);
 }
 
+/**
+ * Database names for `use <name>` and `db.getSiblingDB("<name>")`. The bare form
+ * after `use` inserts the name as typed; the argument form keeps its quotes.
+ */
+function databaseItems(prefix: string, databases: string[]): MongoCompletionItem[] {
+  const quoted = prefix.startsWith('"') || prefix.startsWith("'");
+  return databases
+    .filter((database) => matchesFuzzyPrefix(database, prefix))
+    .slice(0, 100)
+    .map((database) => ({
+      label: database,
+      type: "table" as const,
+      detail: "database",
+      apply: quoted ? quoteMongoString(database, prefix) : database,
+      boost: startsWithPrefix(database, prefix) ? 120 : 90,
+    }));
+}
+
 function collectionRefItems(prefix: string, collections: string[]): MongoCompletionItem[] {
   return collections
     .filter((collection) => matchesFuzzyPrefix(collection, prefix))
@@ -1190,6 +1233,26 @@ function matchDbCollectionPrefix(beforeCursor: string): { prefix: string; from:
   return { prefix, from: beforeCursor.length - prefix.length };
 }
 
+/** Cursor inside the string argument of `db.getSiblingDB(`, with the opening quote as part of the prefix. */
+function matchGetSiblingDbPrefix(beforeCursor: string): { prefix: string; from: number } | null {
+  const match = /(?:^|[\s;(])db\s*\.\s*getSiblingDB\s*\(\s*(["'][^"'\\]*)$/.exec(beforeCursor);
+  if (!match) return null;
+  const prefix = match[1] ?? "";
+  return { prefix, from: beforeCursor.length - prefix.length };
+}
+
+/**
+ * Cursor in the bare database name after a `use` command. A quoted name is not
+ * valid there, so it stays unmatched, and a field called `use` inside an argument
+ * list is a key, not the command.
+ */
+function matchUseDatabasePrefix(beforeCursor: string): { prefix: string; from: number } | null {
+  const match = /(?:^|[\s;])use\s+([^\s;"'()]*)$/.exec(maskMongoLiterals(beforeCursor));
+  if (!match || isInsideCallArguments(beforeCursor)) return null;
+  const prefix = match[1] ?? "";
+  return { prefix: beforeCursor.slice(beforeCursor.length - prefix.length), from: beforeCursor.length - prefix.length };
+}
+
 function matchGetCollectionPrefix(beforeCursor: string): { prefix: string; from: number } | null {
   const match = new RegExp(String.raw`(?:^|[\s;(])${DB_ROOT}\.getCollection\(\s*(["'][^"'\\]*)$`).exec(beforeCursor);
   if (!match) return null;
@@ -1287,7 +1350,7 @@ function isInsideCallArguments(beforeCursor: string): boolean {
   return depth > 0;
 }
 
-/** `use <database>` takes a database name, which this engine has no list of, so it stays quiet. */
+/** After `use` with something `matchUseDatabasePrefix` rejects (a quoted or parenthesised name), so no snippets belong there. */
 function isAfterUseKeyword(beforeCursor: string): boolean {
   return /(?:^|[\s;])use\s+[\w$-]*$/.test(maskMongoLiterals(beforeCursor));
 }

docs/content/docs/mongodb.cn.mdx

@@ -38,7 +38,7 @@ MongoDB 的总数可能使用估算值。数据持续写入时，翻页过程中
 
 ## 过滤与命令
 
-可以在文档浏览器中构建字段过滤，也可以在 MongoDB 查询标签中输入 shell 风格命令。编辑器会根据 Collection、字段、查询运算符、更新运算符、聚合阶段和游标方法提供补全。
+可以在文档浏览器中构建字段过滤，也可以在 MongoDB 查询标签中输入 shell 风格命令。编辑器会根据数据库、Collection、字段、查询运算符、更新运算符、聚合阶段和游标方法提供补全。
 
 常见命令包括：
 
@@ -53,7 +53,7 @@ db.orders.aggregate([
 ])
 ```
 
-`use <database>` 会切换当前 MongoDB 查询上下文。
+`use <database>` 会切换当前 MongoDB 查询上下文。在 `use ` 之后以及 `db.getSiblingDB("…")` 内部会补全数据库名。
 
 ## 表格编辑
 

docs/content/docs/mongodb.mdx

@@ -38,7 +38,7 @@ MongoDB totals may be estimated. In actively changing collections, the count and
 
 ## Filters and Commands
 
-Build field filters in the document browser or use shell-style commands in a MongoDB query tab. Completion can suggest collections, fields, query operators, update operators, aggregation stages, and cursor methods.
+Build field filters in the document browser or use shell-style commands in a MongoDB query tab. Completion can suggest databases, collections, fields, query operators, update operators, aggregation stages, and cursor methods.
 
 Examples:
 
@@ -53,7 +53,7 @@ db.orders.aggregate([
 ])
 ```
 
-Use `use <database>` to switch the MongoDB query context.
+Use `use <database>` to switch the MongoDB query context. The database name is completed after `use ` and inside `db.getSiblingDB("…")`.
 
 ## Grid Editing
 

packages/app-tests/mongoCompletion.test.ts

@@ -641,10 +641,6 @@ test("offers nothing rather than top-level snippets inside an unmodelled argumen
     "db.users.estimatedDocumentCount(",
     "db.runCommand({ ",
     'db.createCollection("x", { ',
-    "use ",
-    "use ord",
-    // A modelled call followed by `use` takes the scan-failure path, not the unmodelled-call one.
-    "db.users.find({}); use ",
   ]) {
     assert.deepEqual(labels(text, { fields, collections }), [], text);
   }
@@ -758,6 +754,44 @@ test("completes both arguments of distinct", () => {
   assert.ok(labels('db.users.distinct("name", { age: { $g', { fields }).includes("$gte"));
 });
 
+test("completes database names after use and inside getSiblingDB", () => {
+  const databases = ["orders", "order_archive", "admin", "local"];
+
+  // `use ` takes a bare name, inserted as typed.
+  assert.deepEqual(labels("use ", { databases }), ["orders", "order_archive", "admin", "local"]);
+  assert.deepEqual(labels("use ord", { databases }), ["orders", "order_archive"]);
+  assert.deepEqual(getMongoCompletionContext("use ord", 7), { mode: "database", prefix: "ord", from: 4 });
+  assert.equal(buildMongoCompletionItems("use ord", 7, { databases })[0]?.apply, "orders");
+  // Whitespace and earlier commands do not get in the way.
+  assert.deepEqual(labels("db.users.find({}); use ord", { databases }), ["orders", "order_archive"]);
+  assert.deepEqual(labels("use   ord", { databases }), ["orders", "order_archive"]);
+  assert.deepEqual(labels("use\tord", { databases }), ["orders", "order_archive"]);
+  // The list opens as soon as the space after `use` is typed, and stays open while the name is typed.
+  assert.equal(shouldAutoOpenMongoCompletion("use ", 4), true);
+  assert.equal(shouldAutoOpenMongoCompletion("use ord", 7), true);
+  // Without a database list the position is still quiet rather than showing root snippets.
+  assert.deepEqual(labels("use ", { fields, collections }), []);
+  assert.deepEqual(labels("use ord", { fields, collections }), []);
+
+  // `getSiblingDB("…")` takes a quoted name, and the completion keeps the quotes.
+  assert.deepEqual(labels('db.getSiblingDB("', { databases }), ["orders", "order_archive", "admin", "local"]);
+  assert.deepEqual(labels('db.getSiblingDB("ord', { databases }), ["orders", "order_archive"]);
+  assert.equal(buildMongoCompletionItems('db.getSiblingDB("ord', 20, { databases })[0]?.apply, '"orders"');
+  assert.equal(buildMongoCompletionItems("db.getSiblingDB('ord", 20, { databases })[0]?.apply, "'orders'");
+  assert.deepEqual(getMongoCompletionContext('db.getSiblingDB("ord")', 20), { mode: "database", prefix: '"ord', from: 16, replaceClosingQuote: '"' });
+  assert.deepEqual(labels("db . getSiblingDB( 'ord", { databases }), ["orders", "order_archive"]);
+  // Root snippets do not leak into the argument when nothing matches.
+  assert.deepEqual(labels('db.getSiblingDB("zzz', { databases }), []);
+
+  // A field that happens to be called `use` is a key, not the command: the filter keeps its own suggestions.
+  const insideFilter = labels("db.users.find({ use ", { databases, fields });
+  assert.notEqual(getMongoCompletionContext("db.users.find({ use ", 20).mode, "database");
+  assert.ok(insideFilter.every((label) => !databases.includes(label)), insideFilter.join(", "));
+  // Neither is `use` inside a string or a comment.
+  assert.deepEqual(labels('db.users.find({ name: "use ', { databases, fields }), []);
+  assert.deepEqual(labels("// use ", { databases }), []);
+});
+
 test("suggests the use and db.version commands", () => {
   assert.equal(buildMongoCompletionItems("us", 2).find((item) => item.label === "use")?.apply, "use ${database}");
   // Database helpers stay reachable after `db.`, alongside the collection names.