fix(mongodb): complete against the database chosen by use

2026-10-02 02:34:42 +08:00 · 2026-10-01 22:38:55 +08:00
parent 3979fcf9ec
commit 1d28475c81
2 changed files with 137 additions and 9 deletions
@@ -87,8 +87,10 @@ export interface MongoCompletionContext {
  replaceClosingQuote?: '"' | "'";
  /** Collection the cursor's command targets, used to load field metadata. */
  collection?: string;
-  /** Database the cursor's command targets when reached through `db.getSiblingDB(…)`, used instead of the editor's active database. */
+  /** Database the cursor's command targets when reached through `db.getSiblingDB(…)` or `use <db>`, used instead of the editor's active database. */
  database?: string;
+  /** Whether the command's root is an explicit `db.getSiblingDB(…)`, which disallows chaining another getSiblingDB. */
+  siblingRoot?: boolean;
  /** Enclosing aggregation stage (`$lookup`, `$group`, …), when inside one. */
  stage?: string;
  /** Kind of pipeline array holding the cursor, when inside a pipeline. */
@@ -332,7 +334,16 @@ export function getMongoCompletionContext(text: string, cursor: number): MongoCo
  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 };
+  if (endsAtDbRootDot(beforeCursor)) {
+    return {
+      mode: "collection",
+      prefix: "",
+      from: safeCursor,
+      collection,
+      database,
+      siblingRoot: endsAtSiblingRootDot(beforeCursor),
+    };
+  }

  const getSiblingDbPrefix = matchGetSiblingDbPrefix(beforeCursor);
  if (getSiblingDbPrefix) {
@@ -358,12 +369,14 @@ export function getMongoCompletionContext(text: string, cursor: number): MongoCo

  const collectionPrefix = matchDbCollectionPrefix(beforeCursor);
  if (collectionPrefix) {
+    const isSibling = matchSiblingCollectionPrefix(beforeCursor);
    return {
      mode: collectionPrefix.prefix.includes(".") ? "collectionOrMethod" : "collection",
      prefix: collectionPrefix.prefix,
      from: collectionPrefix.from,
      collection,
      database,
+      ...(isSibling ? { siblingRoot: true } : {}),
    };
  }

@@ -419,7 +432,7 @@ export function buildMongoCompletionItemsFromContext(context: MongoCompletionCon
      items = databaseItems(prefix, input.databases ?? []);
      break;
    case "collection":
-      items = collectionItems(prefix, collections, context.database !== undefined);
+      items = collectionItems(prefix, collections, context.siblingRoot ?? false);
      break;
    case "collectionOrMethod":
      items = collectionOrMethodItems(prefix, collections);
@@ -1481,6 +1494,7 @@ function readMethodPrefix(beforeCursor: string): { prefix: string; from: number
 * accepts either root rather than only a literal `db.`.
 */
 const DB_ROOT = String.raw`db(?:\s*\.\s*getSiblingDB\s*\(\s*(?:"[^"]*"|'[^']*')\s*\))?`;
+const SIBLING_ROOT_PATTERN = String.raw`db\s*\.\s*getSiblingDB\s*\(\s*(?:"[^"]*"|'[^']*')\s*\)`;
 const COLLECTION_REF = String.raw`(?:[A-Za-z_][\w$-]*|getCollection\(["'][^"']+["']\))`;

 /** `db.` or `db.getSiblingDB("other").` immediately before the cursor. */
@@ -1488,6 +1502,10 @@ function endsAtDbRootDot(beforeCursor: string): boolean {
  return new RegExp(String.raw`(?:^|[\s;(])${DB_ROOT}\s*\.$`).test(beforeCursor);
 }

+function endsAtSiblingRootDot(beforeCursor: string): boolean {
+  return new RegExp(String.raw`(?:^|[\s;(])${SIBLING_ROOT_PATTERN}\s*\.$`).test(beforeCursor);
+}
+
 function matchDbCollectionPrefix(beforeCursor: string): { prefix: string; from: number } | null {
  const match = new RegExp(String.raw`(?:^|[\s;(])${DB_ROOT}\.([A-Za-z_][\w$-]*(?:\.[\w$-]*)*)$`).exec(beforeCursor);
  if (!match) return null;
@@ -1495,6 +1513,17 @@ function matchDbCollectionPrefix(beforeCursor: string): { prefix: string; from:
  return { prefix, from: beforeCursor.length - prefix.length };
 }

+function matchSiblingCollectionPrefix(beforeCursor: string): boolean {
+  return new RegExp(String.raw`(?:^|[\s;(])${SIBLING_ROOT_PATTERN}\s*\.([A-Za-z_][\w$-]*(?:\.[\w$-]*)*)$`).test(beforeCursor);
+}
+
+const MONGO_COMMAND_LINE_START_PATTERN = /(?:use\b|show\s+(?:dbs|databases|collections)\b|db(?:\s*\.|\b))/iy;
+
+function isMongoCommandLineStart(text: string, index: number): boolean {
+  MONGO_COMMAND_LINE_START_PATTERN.lastIndex = index;
+  return MONGO_COMMAND_LINE_START_PATTERN.test(text);
+}
+
 /** 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);
@@ -1708,15 +1737,76 @@ function extractActiveCollection(text: string, cursor: number): string | undefin
  return lastDirect?.[1];
 }

+const USE_COMMAND_PATTERN = /use\s+([a-zA-Z0-9_-]+)(?=[\s;]|$)/iy;
+
 /**
- * The last `db.getSiblingDB("name")` before the cursor decides which database the
- * command targets; plain `db.` references leave it unset so the editor's active
- * database keeps applying.
+ * Resolves the database targeted by the command at the cursor.
+ *
+ * If the current command explicitly addresses another database via `db.getSiblingDB("name")`,
+ * that database takes precedence. Otherwise, the database set by the last preceding top-level
+ * `use <name>` command applies. If neither is present, returns undefined so the editor's active
+ * database continues to apply.
 */
 function extractActiveDatabase(text: string, cursor: number): string | undefined {
-  const before = text.slice(0, cursor);
-  const matches = [...before.matchAll(new RegExp(String.raw`(?:^|[\s;(])db\s*\.\s*getSiblingDB\s*\(\s*(["'])([^"']*)\1\s*\)`, "g"))];
-  return matches[matches.length - 1]?.[2] || undefined;
+  const safeCursor = Math.max(0, Math.min(cursor, text.length));
+  const before = text.slice(0, safeCursor);
+  const masked = maskMongoLiterals(before);
+
+  let parenDepth = 0;
+  let bracketDepth = 0;
+  let braceDepth = 0;
+  let currentCommandStart = 0;
+  let lastUseDb: string | undefined = undefined;
+
+  let i = 0;
+  while (i < masked.length) {
+    const isTopLevel = parenDepth === 0 && bracketDepth === 0 && braceDepth === 0;
+
+    if (isTopLevel) {
+      const prevChar = i > 0 ? masked[i - 1] : "\n";
+      if (/[\s;]/.test(prevChar)) {
+        USE_COMMAND_PATTERN.lastIndex = i;
+        const useMatch = USE_COMMAND_PATTERN.exec(masked);
+        if (useMatch) {
+          lastUseDb = useMatch[1];
+          currentCommandStart = i;
+          i += useMatch[0].length;
+          continue;
+        }
+      }
+
+      if (masked[i] === ";") {
+        let next = i + 1;
+        while (next < masked.length && /\s/.test(masked[next])) next++;
+        currentCommandStart = next;
+      } else if (masked[i] === "\n") {
+        let next = i + 1;
+        while (next < masked.length && (masked[next] === " " || masked[next] === "\t")) next++;
+        if (next < masked.length && isMongoCommandLineStart(masked, next)) {
+          currentCommandStart = next;
+        }
+      }
+    }
+
+    const char = masked[i];
+    if (char === "(") parenDepth++;
+    else if (char === ")") parenDepth = Math.max(0, parenDepth - 1);
+    else if (char === "[") bracketDepth++;
+    else if (char === "]") bracketDepth = Math.max(0, bracketDepth - 1);
+    else if (char === "{") braceDepth++;
+    else if (char === "}") braceDepth = Math.max(0, braceDepth - 1);
+
+    i++;
+  }
+
+  const currentCommandText = before.slice(currentCommandStart);
+  const siblingMatches = [...currentCommandText.matchAll(/(?:^|[\s;(])db\s*\.\s*getSiblingDB\s*\(\s*(?:(["'])([^"']*)\1|[^\s)]+)?\s*\)/gi)];
+  if (siblingMatches.length > 0) {
+    const lastMatch = siblingMatches[siblingMatches.length - 1];
+    return lastMatch?.[2] || undefined;
+  }
+
+  return lastUseDb;
 }

 function collectFieldTypes(value: unknown, prefix: string, out: Map<string, Set<string>>, depth: number) {
@@ -934,6 +934,44 @@ test("completes database names after use and inside getSiblingDB", () => {
  assert.deepEqual(labels("// use ", { databases }), []);
 });

+test("honours a preceding use command and resolves active database", () => {
+  // `use <db>` switches the active database for following commands in the script.
+  const context = getMongoCompletionContext("use analytics\ndb.", "use analytics\ndb.".length);
+  assert.equal(context.database, "analytics");
+  assert.equal(context.mode, "collection");
+
+  // Plain db root after `use` still offers getSiblingDB, unlike an explicit db.getSiblingDB(…) root.
+  const items = labels("use x\ndb.", { collections });
+  assert.ok(items.includes("getSiblingDB"));
+  assert.ok(items.includes("users"));
+
+  // Field completion inside a command carries the database set by preceding `use`.
+  const fieldContext = getMongoCompletionContext("use analytics\ndb.users.find({ ", "use analytics\ndb.users.find({ ".length);
+  assert.equal(fieldContext.database, "analytics");
+  assert.equal(fieldContext.collection, "users");
+  assert.ok(labels("use analytics\ndb.users.find({ na", { fields }).includes("name"));
+
+  // An explicit db.getSiblingDB("b") on the current command wins over an earlier `use a`.
+  const siblingOverride = getMongoCompletionContext('use a\ndb.getSiblingDB("b").users.find({ ', 'use a\ndb.getSiblingDB("b").users.find({ '.length);
+  assert.equal(siblingOverride.database, "b");
+  assert.equal(siblingOverride.collection, "users");
+
+  // `use` inside comments or argument objects does not set the active database.
+  assert.equal(getMongoCompletionContext("db.users.find({ use: 1 })", "db.users.find({ use: 1 })".length).database, undefined);
+  assert.equal(getMongoCompletionContext("// use x\ndb.", "// use x\ndb.".length).database, undefined);
+  assert.equal(getMongoCompletionContext("/* use x */\ndb.", "/* use x */\ndb.".length).database, undefined);
+
+  // Across multiple `use` commands, the last one before the cursor wins.
+  assert.equal(getMongoCompletionContext("use first\nuse second\ndb.", "use first\nuse second\ndb.".length).database, "second");
+
+  // getSiblingDB on a previous command does not leak into a later plain db command.
+  assert.equal(getMongoCompletionContext('use a\ndb.getSiblingDB("b").users.find({});\ndb.', 'use a\ndb.getSiblingDB("b").users.find({});\ndb.'.length).database, "a");
+
+  // Long scripts with thousands of commands resolve without O(n^2) slicing overhead.
+  const longScript = "db.users.find({});\n".repeat(2000) + "use big\ndb.";
+  assert.equal(getMongoCompletionContext(longScript, longScript.length).database, "big");
+});
+
 test("completes the keys of an operator's own sub-document", () => {
  // `{ field: { $… } }` takes the field operators, but `$text` and the geo operators hold
  // documents with their own keys, where `$gt` and friends are nonsense.