fix(mongodb): complete collections reached through db["name"]

2026-10-19 16:16:45 +00:00 · 2026-10-04 17:01:14 +08:00
parent b12e71455c
commit 09a7d21510
2 changed files with 119 additions and 5 deletions
@@ -1913,6 +1913,7 @@ function readMethodPrefix(beforeCursor: string): { prefix: string; from: number
 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\(["'][^"']+["']\))`;
+const COLLECTION_RECEIVER = String.raw`(?:${DB_ROOT}\.${COLLECTION_REF}|${DB_ROOT}\s*\[\s*(?:"[^"]*"|'[^']*')\s*\])`;

 /** `db.` or `db.getSiblingDB("other").` immediately before the cursor. */
 function endsAtDbRootDot(beforeCursor: string): boolean {
@@ -1973,14 +1974,14 @@ function matchShowSubcommandPrefix(beforeCursor: string): { prefix: string; from
 }

 function matchGetCollectionPrefix(beforeCursor: string): { prefix: string; from: number } | null {
-  const match = lastCodeMatch(new RegExp(String.raw`(?:^|[\s;(])${DB_ROOT}\.getCollection\(\s*(["'][^"'\\]*)$`), beforeCursor);
+  const match = lastCodeMatch(new RegExp(String.raw`(?:^|[\s;(])${DB_ROOT}(?:\.getCollection\(|\s*\[)\s*(["'][^"'\\]*)$`), beforeCursor);
  if (!match) return null;
  const prefix = match[1] ?? "";
  return { prefix, from: beforeCursor.length - prefix.length };
 }

 function isAfterCollectionDot(beforeCursor: string): boolean {
-  return !!lastCodeMatch(new RegExp(String.raw`(?:^|[\s;(])${DB_ROOT}\.${COLLECTION_REF}\.[\w$-]*$`), beforeCursor);
+  return !!lastCodeMatch(new RegExp(String.raw`(?:^|[\s;(])${COLLECTION_RECEIVER}\.[\w$-]*$`), beforeCursor);
 }

 /**
@@ -1989,7 +1990,7 @@ function isAfterCollectionDot(beforeCursor: string): boolean {
 * chains (which only accept `toArray()` and `pretty()`).
 */
 function matchCursorMethodDot(beforeCursor: string): { find: boolean; countable: boolean; terminal?: boolean } | null {
-  const collectionCall = new RegExp(String.raw`(?:^|[\s;(])${DB_ROOT}\.${COLLECTION_REF}\.(find|aggregate)\s*\(`, "g");
+  const collectionCall = new RegExp(String.raw`(?:^|[\s;(])${COLLECTION_RECEIVER}\.(find|aggregate)\s*\(`, "g");
  const lastMatch = lastCodeMatch(collectionCall, beforeCursor);
  if (!lastMatch) return null;

@@ -2240,13 +2241,28 @@ function extractActiveCollection(before: string): string | undefined {
  const literals = mongoLiteralRanges(before);
  const isCode = (match: RegExpMatchArray) => match.index === undefined || !isInsideMongoLiteral(literals, match.index);
  const getCollectionMatches = [...before.matchAll(new RegExp(String.raw`${DB_ROOT}\.getCollection\(["']([^"']+)["']\)`, "g"))].filter(isCode);
+  const bracketMatches = [...before.matchAll(new RegExp(String.raw`${DB_ROOT}\s*\[\s*(["'])([^"']+)\1\s*\]`, "g"))].filter(isCode);
  const directMatches = [...before.matchAll(new RegExp(String.raw`${DB_ROOT}\.([A-Za-z_][\w$-]*)\s*\.`, "g"))].filter(isCode).filter((match) => match[1] !== "getCollection");
  const lastGetCollection = getCollectionMatches[getCollectionMatches.length - 1];
+  const lastBracket = bracketMatches[bracketMatches.length - 1];
  const lastDirect = directMatches[directMatches.length - 1];
  const getCollectionIndex = lastGetCollection?.index ?? -1;
+  const bracketIndex = lastBracket?.index ?? -1;
  const directIndex = lastDirect?.index ?? -1;
-  if (getCollectionIndex > directIndex) return lastGetCollection?.[1];
-  return lastDirect?.[1];
+
+  let lastIndex = directIndex;
+  let activeCollection = lastDirect?.[1];
+
+  if (getCollectionIndex > lastIndex) {
+    lastIndex = getCollectionIndex;
+    activeCollection = lastGetCollection?.[1];
+  }
+  if (bracketIndex > lastIndex) {
+    lastIndex = bracketIndex;
+    activeCollection = lastBracket?.[2];
+  }
+
+  return activeCollection;
 }

 const USE_COMMAND_PATTERN = /use\s+([a-zA-Z0-9_-]+)(?=[\s;]|$)/iy;
@@ -2071,6 +2071,104 @@ test("keeps the find chain completable when a comment holds an unbalanced parent
  assert.deepEqual(labels('db.users.find({ note: "(x" }).li', { fields }), ["limit"]);
 });

+test("completes bracket collection references db['name'] and db[\"name\"]", () => {
+  // Case 1: db["orders-2024"]. -> method mode, collection methods
+  const afterBracketDot = 'db["orders-2024"].';
+  const ctx1 = getMongoCompletionContext(afterBracketDot, afterBracketDot.length);
+  assert.equal(ctx1.mode, "method");
+  assert.equal(ctx1.collection, "orders-2024");
+  const labels1 = labels(afterBracketDot);
+  assert.ok(labels1.includes("find"));
+  assert.ok(labels1.includes("findOne"));
+  assert.ok(labels1.includes("aggregate"));
+
+  // Case 2: db['users'].fi (single quotes, prefix) -> method mode, filtered collection methods
+  const prefixedSingleQuote = "db['users'].fi";
+  const ctx2 = getMongoCompletionContext(prefixedSingleQuote, prefixedSingleQuote.length);
+  assert.equal(ctx2.mode, "method");
+  assert.equal(ctx2.collection, "users");
+  assert.equal(ctx2.prefix, "fi");
+  const labels2 = labels(prefixedSingleQuote);
+  assert.ok(labels2.includes("find"));
+  assert.ok(labels2.includes("findOne"));
+  assert.ok(labels2.includes("findOneAndUpdate"));
+  assert.equal(labels2.includes("db.collection.find"), false);
+
+  // Case 3: db["orders-2024"].find({ -> filterField mode, collection resolved for field completion
+  const findFilter = 'db["orders-2024"].find({';
+  const ctx3 = getMongoCompletionContext(findFilter, findFilter.length);
+  assert.equal(ctx3.mode, "filterField");
+  assert.equal(ctx3.collection, "orders-2024");
+  const labels3 = labels(findFilter, { fields });
+  assert.ok(labels3.includes("name"));
+  assert.ok(labels3.includes("createdAt"));
+
+  // Case 4: db["orders-2024"].find({}). -> cursor methods
+  const cursorChain = 'db["orders-2024"].find({}).';
+  const ctx4 = getMongoCompletionContext(cursorChain, cursorChain.length);
+  assert.equal(ctx4.mode, "cursorMethod");
+  assert.equal(ctx4.collection, "orders-2024");
+  const labels4 = labels(cursorChain);
+  assert.ok(labels4.includes("limit"));
+  assert.ok(labels4.includes("sort"));
+  assert.ok(labels4.includes("skip"));
+  assert.ok(labels4.includes("count"));
+  assert.ok(labels4.includes("toArray"));
+
+  // Case 5: db["orders-2024"].aggregate([{ $ -> aggregation stages with active collection
+  const aggStage = 'db["orders-2024"].aggregate([{ $';
+  const ctx5 = getMongoCompletionContext(aggStage, aggStage.length);
+  assert.equal(ctx5.mode, "stage");
+  assert.equal(ctx5.collection, "orders-2024");
+  const labels5 = labels(aggStage);
+  assert.ok(labels5.includes("$match"));
+  assert.ok(labels5.includes("$group"));
+  assert.ok(labels5.includes("$project"));
+
+  // Case 6: db[" -> collectionRef mode, quoted collection names
+  const bracketOpen = 'db["';
+  const ctx6 = getMongoCompletionContext(bracketOpen, bracketOpen.length);
+  assert.equal(ctx6.mode, "collectionRef");
+  assert.equal(ctx6.prefix, '"');
+  const items6 = buildMongoCompletionItems(bracketOpen, bracketOpen.length, { collections });
+  assert.ok(items6.some((item) => item.label === "users" && item.apply === '"users"'));
+  assert.ok(items6.some((item) => item.label === "order-items" && item.apply === '"order-items"'));
+
+  // Case 6b: db[""] with cursor between quotes consumes closing quote
+  const textQuoted = 'db[""]';
+  const cursorQuoted = textQuoted.indexOf('""') + 1;
+  const ctxQuoted = getMongoCompletionContext(textQuoted, cursorQuoted);
+  assert.equal(ctxQuoted.mode, "collectionRef");
+  assert.equal(ctxQuoted.replaceClosingQuote, '"');
+  const itemQuoted = buildMongoCompletionItems(textQuoted, cursorQuoted, { collections }).find((c) => c.label === "users");
+  assert.equal(itemQuoted?.replaceClosingQuote, '"');
+
+  // Case 7 & 8: Name containing '.' and '-' with both quote styles
+  const dottedName = "db['audit.logs'].find({";
+  const ctxDot = getMongoCompletionContext(dottedName, dottedName.length);
+  assert.equal(ctxDot.mode, "filterField");
+  assert.equal(ctxDot.collection, "audit.logs");
+  const dottedMethod = "db['audit.logs'].";
+  const ctxDotMethod = getMongoCompletionContext(dottedMethod, dottedMethod.length);
+  assert.equal(ctxDotMethod.mode, "method");
+  assert.equal(ctxDotMethod.collection, "audit.logs");
+  assert.ok(labels(dottedMethod).includes("find"));
+
+  // Case 9: db.getSiblingDB("shop")["orders"].find({ -> database "shop", collection "orders"
+  const siblingText = 'db.getSiblingDB("shop")["orders"].find({';
+  const ctxSibling = getMongoCompletionContext(siblingText, siblingText.length);
+  assert.equal(ctxSibling.mode, "filterField");
+  assert.equal(ctxSibling.collection, "orders");
+  assert.equal(ctxSibling.database, "shop");
+
+  const siblingMethod = 'db.getSiblingDB("shop")["orders"].';
+  const ctxSiblingMethod = getMongoCompletionContext(siblingMethod, siblingMethod.length);
+  assert.equal(ctxSiblingMethod.mode, "method");
+  assert.equal(ctxSiblingMethod.collection, "orders");
+  assert.equal(ctxSiblingMethod.database, "shop");
+  assert.ok(labels(siblingMethod).includes("find"));
+});
+
 test("suppresses completion after a dot following findOne", () => {
  const text = "db.users.findOne({}).";
  const context = getMongoCompletionContext(text, text.length);