fix(mongodb): offer $comment, skip used keys and keep $set values quiet

This commit is contained in:
Loc Nguyen
2026-10-04 17:23:43 +08:00
committed by GitHub
parent 09a7d21510
commit 7af2ee1fce
3 changed files with 186 additions and 5 deletions
+122 -2
View File
@@ -17,6 +17,7 @@ import {
OPERATOR_SUB_KEYS,
STAGE_OPTION_KEYS,
TOP_LEVEL_QUERY_OPERATORS,
UPDATE_OPERATOR_LABELS,
UPDATE_OPERATORS,
VALUE_SNIPPETS,
WINDOW_FUNCTION_OPERATORS,
@@ -121,6 +122,8 @@ export interface MongoCompletionContext {
enumKey?: string;
/** User variables in scope for aggregation expressions. */
variables?: string[];
/** Keys already defined in the innermost object container (before and after cursor). */
usedKeys?: string[];
}
export interface MongoCompletionInput {
@@ -445,6 +448,7 @@ export function getMongoCompletionContext(text: string, cursor: number): MongoCo
...(classified.enumKey ? { enumKey: classified.enumKey } : {}),
...(classified.pipelineKind ? { pipelineKind: classified.pipelineKind } : {}),
...(variables ? { variables } : {}),
...(scan.usedKeys && scan.usedKeys.length > 0 ? { usedKeys: scan.usedKeys } : {}),
collection: classified.collection ?? collection,
};
}
@@ -453,6 +457,24 @@ export function buildMongoCompletionItems(text: string, cursor: number, input: M
return buildMongoCompletionItemsFromContext(getMongoCompletionContext(text, cursor), input);
}
const KEY_POSITION_MODES: ReadonlySet<MongoCompletionMode> = new Set([
"filterField",
"field",
"queryOperator",
"updateOperator",
"pullCondition",
"pushModifier",
"projectionOperator",
"stage",
"stageOption",
"methodOption",
"bulkWriteOperation",
"bulkWriteField",
"operatorField",
"accumulator",
"windowOperator",
]);
export function buildMongoCompletionItemsFromContext(context: MongoCompletionContext, input: MongoCompletionInput = {}): MongoCompletionItem[] {
const { mode, prefix } = context;
const collections = input.collections ?? [];
@@ -570,6 +592,11 @@ export function buildMongoCompletionItemsFromContext(context: MongoCompletionCon
default:
items = [];
}
if (context.usedKeys && context.usedKeys.length > 0 && KEY_POSITION_MODES.has(mode)) {
const used = new Set(context.usedKeys);
items = items.filter((item) => !used.has(item.label));
}
return finalizeQuotedMongoCompletionItems(context, items);
}
@@ -647,7 +674,14 @@ export function getMongoDocumentQueryCompletionContext(text: string, cursor: num
const { prefix, from } = readMongoPropertyPrefix(text, safeCursor);
const variables = classified.mode === "fieldRef" || classified.mode === "expression" ? collectScopeVariables(scan) : undefined;
return { ...classified, prefix, from, replaceClosingQuote: closingQuoteAtCursor(prefix, text, safeCursor), ...(variables ? { variables } : {}) };
return {
...classified,
prefix,
from,
replaceClosingQuote: closingQuoteAtCursor(prefix, text, safeCursor),
...(variables ? { variables } : {}),
...(scan.usedKeys && scan.usedKeys.length > 0 ? { usedKeys: scan.usedKeys } : {}),
};
}
/**
@@ -784,6 +818,8 @@ interface MongoCallScan {
valueKey: string | null;
inValue: boolean;
inString: boolean;
/** Keys already defined in the innermost object container (before and after cursor). */
usedKeys?: string[];
}
interface MongoCursorClass {
@@ -903,7 +939,87 @@ function scanMongoCallArguments(text: string, start: number, cursor: number): Mo
if (!/\s/.test(char)) token += char;
}
return { argIndex, stack, valueKey, inValue, inString: quote !== null };
const innermostAtCursor = stack[stack.length - 1];
const usedKeys = innermostAtCursor && innermostAtCursor.kind === "object" && !inValue ? collectInnermostObjectKeys(text, cursor, stack, innermostAtCursor, token, quote) : undefined;
return { argIndex, stack, valueKey, inValue, inString: quote !== null, ...(usedKeys && usedKeys.length > 0 ? { usedKeys } : {}) };
}
function collectInnermostObjectKeys(text: string, cursor: number, stack: MongoContainer[], innermostObj: MongoContainer, tokenAtCursor: string, quoteAtCursor: string | null): string[] {
const keys = new Set(innermostObj.keys ?? []);
const continuationStack: MongoContainer[] = [...stack];
let token = tokenAtCursor;
let quote = quoteAtCursor;
let skippingCursorKey = true;
for (let i = cursor; i < text.length; i++) {
const char = text[i] ?? "";
if (quote) {
if (char === "\\") {
i++;
if (i < text.length) token += text[i];
continue;
}
if (char === quote) {
quote = null;
} else {
token += char;
}
continue;
}
if ((char === "/" && (text[i + 1] === "/" || text[i + 1] === "*")) || (char === "-" && text[i + 1] === "-")) {
const skipped = skipMongoStringOrComment(text, i, text.length);
if (skipped > i) {
i = skipped - 1;
continue;
}
}
if (char === '"' || char === "'") {
quote = char;
token = "";
continue;
}
if (char === "{" || char === "[" || char === "(") {
skippingCursorKey = false;
continuationStack.push({
kind: char === "{" ? "object" : char === "[" ? "array" : "call",
key: null,
});
token = "";
continue;
}
if (char === "}" || char === "]" || char === ")") {
skippingCursorKey = false;
if (continuationStack.length === 0) break;
const popped = continuationStack.pop();
if (popped === innermostObj) {
break;
}
token = "";
continue;
}
if (char === ":") {
if (skippingCursorKey) {
skippingCursorKey = false;
token = "";
continue;
}
if (continuationStack[continuationStack.length - 1] === innermostObj) {
const cleanKey = token.trim().replace(/^["']|["']$/g, "");
if (cleanKey) keys.add(cleanKey);
}
token = "";
continue;
}
if (char === ",") {
skippingCursorKey = false;
token = "";
continue;
}
if (!/\s/.test(char)) token += char;
}
return [...keys];
}
function isExpressionArray(scan: MongoCallScan): boolean {
@@ -1142,6 +1258,10 @@ function classifyUpdate(scan: MongoCallScan, rootIndex: number): MongoCursorClas
const parent = scan.stack[scan.stack.length - 2];
if (scan.inValue) {
if (innerDepth(scan, rootIndex) === 0 && scan.valueKey) {
const cleanKey = scan.valueKey.replace(/^["']|["']$/g, "");
if (UPDATE_OPERATOR_LABELS.has(cleanKey)) return { mode: "none" };
}
// `$currentDate: { at: { $type: "timestamp" } }`.
if (scan.valueKey === "$type" && parent?.key === "$currentDate") return { mode: "enumValue", enumKey: "currentDateType" };
return { mode: scan.inString ? "none" : "value" };
@@ -78,6 +78,7 @@ export const TOP_LEVEL_QUERY_OPERATORS: MongoOperatorSpec[] = specs([
["$and", "Joins clauses with a logical AND", "$and: [${}]"],
["$or", "Joins clauses with a logical OR", "$or: [${}]"],
["$nor", "Joins clauses with a logical NOR", "$nor: [${}]"],
["$comment", "Attaches a comment to the query, visible in logs and profiler", '$comment: "${comment}"'],
["$expr", "Uses aggregation expressions in a query", "$expr: { ${} }"],
["$text", "Performs a text search", '$text: { $search: "${text}" }'],
["$where", "Matches with a JavaScript predicate", '$where: "${expression}"'],
+63 -3
View File
@@ -392,7 +392,7 @@ test("treats -- as a line comment, matching the editor's SQL language mode", ()
assert.deepEqual(labels("db.users.find({ -- fi", { fields }), []);
assert.deepEqual(labels("db.users.find({ age: { -- $", { fields }), []);
// Code after a closed `--` line still completes.
assert.deepEqual(labels("db.users.find({ name: 1, -- note\n na", { fields }), ["name"]);
assert.deepEqual(labels("db.users.find({ createdAt: 1, -- note\n na", { fields }), ["name"]);
// A double dash inside a string value is not a comment.
assert.deepEqual(labels('db.users.find({ note: "a--b", na', { fields }), ["name"]);
});
@@ -1627,8 +1627,9 @@ test("completes runCommand command names and their collection arguments", () =>
assert.deepEqual(labels('db.runCommand({ collStats: "', { collections }), collections);
assert.deepEqual(labels('db.runCommand({ find: "us', { collections }), ["users", "user_events"]);
assert.deepEqual(labels('db.runCommand({ ping: "', { collections }), []);
// A second key still completes commands, and `db.` offers the helper.
assert.ok(labels('db.runCommand({ find: "users", ').includes("find"));
// A second key still completes commands (and drops the used command), and `db.` offers the helper.
assert.ok(labels('db.runCommand({ find: "users", ').includes("ping"));
assert.equal(labels('db.runCommand({ find: "users", ').includes("find"), false);
assert.ok(labels("db.").includes("runCommand"));
});
@@ -2071,6 +2072,65 @@ 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("offers $comment as a top-level query operator", () => {
const items = labels("db.users.find({ $c", { fields });
assert.ok(items.includes("$comment"), "db.users.find({ $c should offer $comment");
const snippet = buildMongoCompletionItems("db.users.find({ $comment", "db.users.find({ $comment".length, { fields }).find((item) => item.label === "$comment");
assert.equal(snippet?.apply, '$comment: "${comment}"');
assert.ok(parseMongoCommand('db.users.find({ status: "A", $comment: "audit" })'));
});
test("skips keys already present in the same object", () => {
// Query operator keys
const queryOps = labels("db.users.find({ age: { $gt: 1, $g", { fields });
assert.equal(queryOps.includes("$gt"), false, "expected no $gt in second operator position");
assert.ok(queryOps.includes("$gte"), "expected $gte to be offered");
// Filter fields
const filterKeys = labels('db.users.find({ name: "a", n', { fields });
assert.equal(filterKeys.includes("name"), false, "expected no `name` when name already present");
// Update operators
const updateOps = labels("db.users.updateOne({}, { $set: { a: 1 }, $s", { fields });
assert.equal(updateOps.includes("$set"), false, "expected no $set when $set already present");
assert.ok(updateOps.includes("$setOnInsert"), "expected $setOnInsert to remain available");
// Guards:
// find({ still lists all fields
assert.ok(labels("db.users.find({ ", { fields }).includes("name"));
// find({ age: { $ still lists $gt
assert.ok(labels("db.users.find({ age: { $", { fields }).includes("$gt"));
// Key after the cursor is also excluded
const textWithKeyAfterCursor = "db.users.find({ n , name: 1 })";
const cursor = textWithKeyAfterCursor.indexOf("n") + 1;
const itemsAfterCursor = buildMongoCompletionItems(textWithKeyAfterCursor, cursor, { fields }).map((item) => item.label);
assert.equal(itemsAfterCursor.includes("name"), false, "expected key after cursor to be excluded");
});
test("withholds scalar values in value position right after a top-level update operator", () => {
// At the root of an update document, $set: requires an object, so value position must be quiet (mode none)
const context = getMongoCompletionContext("db.users.updateOne({}, { $set: ", "db.users.updateOne({}, { $set: ".length);
assert.equal(context.mode, "none");
assert.deepEqual(labels("db.users.updateOne({}, { $set: ", { fields }), []);
// Applies to every UPDATE_OPERATORS key at the root of updateOne / updateMany / findOneAndUpdate and bulkWrite
assert.equal(getMongoCompletionContext("db.users.updateMany({}, { $inc: ", "db.users.updateMany({}, { $inc: ".length).mode, "none");
assert.equal(getMongoCompletionContext("db.users.findOneAndUpdate({}, { $unset: ", "db.users.findOneAndUpdate({}, { $unset: ".length).mode, "none");
assert.equal(getMongoCompletionContext("db.users.bulkWrite([{ updateOne: { filter: {}, update: { $set: ", "db.users.bulkWrite([{ updateOne: { filter: {}, update: { $set: ".length).mode, "none");
// Guard: values INSIDE the operator object ($set: { name: ) must keep the value list unchanged
const insideOperatorValues = labels("db.users.updateOne({}, { $set: { name: ", { fields });
assert.ok(insideOperatorValues.includes("ObjectId"));
assert.ok(insideOperatorValues.includes("ISODate"));
assert.ok(insideOperatorValues.includes("true"));
// Guard: pipeline-style updates are unaffected
assert.ok(labels("db.users.updateOne({}, [{ $", { fields }).includes("$set"));
assert.deepEqual(labels("db.users.updateOne({}, [{ $set: { ", { fields }), ["_id", "createdAt", "name", "profile.email"]);
});
test("completes bracket collection references db['name'] and db[\"name\"]", () => {
// Case 1: db["orders-2024"]. -> method mode, collection methods
const afterBracketDot = 'db["orders-2024"].';