From dd82d948a94ada372c4d811e88094fde2da3d508 Mon Sep 17 00:00:00 2001 From: Loc Nguyen <113083148+thailoc-dev@users.noreply.github.com> Date: Sat, 26 Sep 2026 20:48:22 +0700 Subject: [PATCH] feat(mongodb): complete collection method option completion keys --- apps/desktop/src/lib/mongo/mongoCompletion.ts | 47 ++++++++--- .../src/lib/mongo/mongoCompletionTables.ts | 64 +++++++++++++++ packages/app-tests/mongoCompletion.test.ts | 80 ++++++++++++++++++- 3 files changed, 178 insertions(+), 13 deletions(-) diff --git a/apps/desktop/src/lib/mongo/mongoCompletion.ts b/apps/desktop/src/lib/mongo/mongoCompletion.ts index 2b954571e9..110d11156d 100644 --- a/apps/desktop/src/lib/mongo/mongoCompletion.ts +++ b/apps/desktop/src/lib/mongo/mongoCompletion.ts @@ -7,6 +7,7 @@ import { FIELD_QUERY_OPERATORS, PIPELINE_STAGES, PUSH_MODIFIERS, + METHOD_OPTION_KEYS, STAGE_OPTION_KEYS, TOP_LEVEL_QUERY_OPERATORS, UPDATE_OPERATORS, @@ -46,7 +47,8 @@ export type MongoCompletionMode = | "expression" | "accumulator" | "stage" - | "stageOption"; + | "stageOption" + | "methodOption"; export interface MongoCompletionField { name: string; @@ -74,6 +76,8 @@ export interface MongoCompletionContext { collection?: string; /** Enclosing aggregation stage (`$lookup`, `$group`, …), when inside one. */ stage?: string; + /** Collection method whose options object the cursor sits in. */ + method?: string; } export interface MongoCompletionInput { @@ -191,12 +195,12 @@ const ROOT_SNIPPET_BOOST: Record<(typeof ROOT_SNIPPETS)[number]["label"], number type MongoArgRole = "filter" | "update" | "replacement" | "document" | "documents" | "operations" | "pipeline" | "projection" | "keys" | "sortKeys" | "fieldName" | "options"; const METHOD_ARG_ROLES: Record = { - find: ["filter", "projection", "options"], + find: ["filter", "projection"], findOne: ["filter", "projection", "options"], - countDocuments: ["filter", "options"], - count: ["filter", "options"], - deleteOne: ["filter", "options"], - deleteMany: ["filter", "options"], + countDocuments: ["filter"], + count: ["filter"], + deleteOne: ["filter"], + deleteMany: ["filter"], findOneAndDelete: ["filter", "options"], updateOne: ["filter", "update", "options"], updateMany: ["filter", "update", "options"], @@ -204,8 +208,8 @@ const METHOD_ARG_ROLES: Record = { bulkWrite: ["operations", "options"], findOneAndUpdate: ["filter", "update", "options"], findOneAndReplace: ["filter", "replacement", "options"], - insertOne: ["document", "options"], - insertMany: ["documents", "options"], + insertOne: ["document"], + insertMany: ["documents"], aggregate: ["pipeline", "options"], createIndex: ["keys", "options"], distinct: ["fieldName", "filter"], @@ -249,7 +253,7 @@ export function getMongoCompletionContext(text: string, cursor: number): MongoCo const collection = extractActiveCollection(text, safeCursor); const { prefix, from } = readMongoPropertyPrefix(text, safeCursor); const replaceClosingQuote = closingQuoteAtCursor(prefix, text, safeCursor); - const at = (mode: MongoCompletionMode, stage?: string): MongoCompletionContext => ({ mode, prefix, from, replaceClosingQuote, collection, stage }); + const at = (mode: MongoCompletionMode, stage?: string, method?: string): MongoCompletionContext => ({ mode, prefix, from, replaceClosingQuote, collection, stage, method }); if (isInsideMongoComment(text, safeCursor)) return { mode: "none", prefix: "", from: safeCursor }; @@ -294,7 +298,7 @@ export function getMongoCompletionContext(text: string, cursor: number): MongoCo if (!scan) return at("root"); const classified = classifyCursorInCall(call.method, scan); - return { ...at(classified.mode, classified.stage), collection: classified.collection ?? collection }; + return { ...at(classified.mode, classified.stage, classified.method), collection: classified.collection ?? collection }; } export function buildMongoCompletionItems(text: string, cursor: number, input: MongoCompletionInput = {}): MongoCompletionItem[] { @@ -371,6 +375,9 @@ export function buildMongoCompletionItemsFromContext(context: MongoCompletionCon case "stageOption": items = specItems(STAGE_OPTION_KEYS[context.stage ?? ""] ?? [], prefix, `${context.stage} option`, 100); break; + case "methodOption": + items = specItems(METHOD_OPTION_KEYS[context.method ?? ""] ?? [], prefix, `${context.method}() option`, 100); + break; default: items = []; } @@ -547,6 +554,7 @@ interface MongoCursorClass { mode: MongoCompletionMode; stage?: string; collection?: string; + method?: string; } /** @@ -649,8 +657,9 @@ function classifyCursorInCall(method: string, scan: MongoCallScan): MongoCursorC return classifyPipeline(scan); // bulkWrite operations are `{ : { filter, update, … } }` entries; completing inside them is a follow-up. case "operations": - case "options": return { mode: "none" }; + case "options": + return classifyMethodOptions(method, scan); default: return { mode: "none" }; } @@ -720,6 +729,22 @@ function classifyKeyMap(scan: MongoCallScan, rootIndex: number): MongoCompletion return "field"; } +/** Option keys whose value is a field-to-value map, so the cursor completes field names there. */ +const FIELD_MAP_OPTION_KEYS = new Set(["sort", "projection"]); + +function classifyMethodOptions(method: string, scan: MongoCallScan): MongoCursorClass { + const inner = innermost(scan); + const depth = innerDepth(scan, 0); + if (!inner || depth < 0) return { mode: "none" }; + + // `{ sort: { … } }` and `{ projection: { … } }` are field maps one level in. + if (depth === 1 && inner.kind === "object" && FIELD_MAP_OPTION_KEYS.has(inner.key ?? "")) { + return { mode: scan.inValue ? "none" : "field", method }; + } + if (depth !== 0 || scan.inValue) return { mode: "none" }; + return { mode: inner.kind === "object" ? "methodOption" : "none", method }; +} + function classifyPipeline(scan: MongoCallScan): MongoCursorClass { const pipelineIndex = findPipelineArrayIndex(scan.stack); if (pipelineIndex < 0) return { mode: "none" }; diff --git a/apps/desktop/src/lib/mongo/mongoCompletionTables.ts b/apps/desktop/src/lib/mongo/mongoCompletionTables.ts index 92c4e528ca..c6f03610ac 100644 --- a/apps/desktop/src/lib/mongo/mongoCompletionTables.ts +++ b/apps/desktop/src/lib/mongo/mongoCompletionTables.ts @@ -294,6 +294,70 @@ export const EXPRESSION_OPERATORS: MongoOperatorSpec[] = specs([ ["$rand", "Returns a random float between 0 and 1", "$rand: {}"], ]); +/** + * Option keys accepted by a collection method's trailing options argument. + * + * Only methods whose options DBX actually applies appear here, and only with the keys they + * accept: several of these are deserialized with `deny_unknown_fields`, so suggesting a key the + * driver does not know would produce a completion that fails at Run. Methods with no options + * support (`find`, `insertOne`, `deleteOne`, `countDocuments`, …) are absent by design. + * + * `createIndex` and `aggregate` pass their options through to the server command, so those lists + * are the common documented ones rather than an exhaustive set. + */ +export const METHOD_OPTION_KEYS: Record = { + findOne: specs([["sort", "Sort order used to pick the single document", "sort: { ${field}: -1 }"]]), + updateOne: specs([ + ["upsert", "Insert the document when the filter matches nothing", "upsert: true"], + ["arrayFilters", "Conditions for the $[] positional operator", 'arrayFilters: [{ "${elem}.${field}": ${} }]'], + ]), + updateMany: specs([ + ["upsert", "Insert the document when the filter matches nothing", "upsert: true"], + ["arrayFilters", "Conditions for the $[] positional operator", 'arrayFilters: [{ "${elem}.${field}": ${} }]'], + ]), + replaceOne: specs([["upsert", "Insert the replacement when the filter matches nothing", "upsert: true"]]), + findOneAndUpdate: specs([ + ["returnDocument", "Return the document 'before' or 'after' the update", 'returnDocument: "after"'], + ["returnNewDocument", "Return the updated document instead of the original", "returnNewDocument: true"], + ["new", "Legacy alias for returnNewDocument", "new: true"], + ["upsert", "Insert the document when the filter matches nothing", "upsert: true"], + ["projection", "Fields to return", "projection: { ${field}: 1 }"], + ["sort", "Sort order used to pick the single document", "sort: { ${field}: -1 }"], + ["arrayFilters", "Conditions for the $[] positional operator", 'arrayFilters: [{ "${elem}.${field}": ${} }]'], + ]), + findOneAndReplace: specs([ + ["returnDocument", "Return the document 'before' or 'after' the replacement", 'returnDocument: "after"'], + ["returnNewDocument", "Return the replacement instead of the original", "returnNewDocument: true"], + ["new", "Legacy alias for returnNewDocument", "new: true"], + ["upsert", "Insert the replacement when the filter matches nothing", "upsert: true"], + ["projection", "Fields to return", "projection: { ${field}: 1 }"], + ["sort", "Sort order used to pick the single document", "sort: { ${field}: -1 }"], + ]), + findOneAndDelete: specs([ + ["projection", "Fields to return", "projection: { ${field}: 1 }"], + ["sort", "Sort order used to pick the single document", "sort: { ${field}: -1 }"], + ]), + bulkWrite: specs([["ordered", "Stop at the first failed operation", "ordered: false"]]), + createIndex: specs([ + ["name", "Index name", 'name: "${name}"'], + ["unique", "Reject duplicate values", "unique: true"], + ["sparse", "Index only documents that have the field", "sparse: true"], + ["expireAfterSeconds", "TTL index lifetime in seconds", "expireAfterSeconds: 3600"], + ["partialFilterExpression", "Index only documents matching this filter", "partialFilterExpression: { ${} }"], + ["collation", "Locale-aware comparison rules", 'collation: { locale: "${en}" }'], + ["hidden", "Keep the index but hide it from the planner", "hidden: true"], + ]), + aggregate: specs([ + ["allowDiskUse", "Let stages write temporary files", "allowDiskUse: true"], + ["maxTimeMS", "Server-side time limit in milliseconds", "maxTimeMS: 5000"], + ["collation", "Locale-aware comparison rules", 'collation: { locale: "${en}" }'], + ["hint", "Index to use", 'hint: "${index}"'], + ["comment", "Comment recorded in the server logs and profiler", 'comment: "${comment}"'], + ["let", "Variables available to the pipeline as $$name", "let: { ${name}: ${} }"], + ["explain", "Return the query plan instead of the results", "explain: true"], + ]), +}; + /** Option keys accepted by the stages whose shape is a fixed set of names. */ export const STAGE_OPTION_KEYS: Record = { $lookup: specs([ diff --git a/packages/app-tests/mongoCompletion.test.ts b/packages/app-tests/mongoCompletion.test.ts index 10ff0416e6..4ff9cd3b9e 100644 --- a/packages/app-tests/mongoCompletion.test.ts +++ b/packages/app-tests/mongoCompletion.test.ts @@ -1,7 +1,8 @@ import { strict as assert } from "node:assert"; import { test } from "vitest"; import { buildMongoCompletionItems, getMongoCompletionContext, getMongoCompletionResultValidFor, inferMongoCompletionFields, shouldAutoOpenMongoCompletion } from "../../apps/desktop/src/lib/mongo/mongoCompletion.ts"; -import { ACCUMULATORS, EXPRESSION_OPERATORS, EXTENDED_JSON_VALUES, PIPELINE_STAGES, PUSH_MODIFIERS, QUERY_OPERATORS, STAGE_OPTION_KEYS, UPDATE_OPERATORS, VALUE_SNIPPETS } from "../../apps/desktop/src/lib/mongo/mongoCompletionTables.ts"; +import { ACCUMULATORS, EXPRESSION_OPERATORS, EXTENDED_JSON_VALUES, METHOD_OPTION_KEYS, PIPELINE_STAGES, PUSH_MODIFIERS, QUERY_OPERATORS, STAGE_OPTION_KEYS, UPDATE_OPERATORS, VALUE_SNIPPETS } from "../../apps/desktop/src/lib/mongo/mongoCompletionTables.ts"; +import { parseMongoCommand } from "../../apps/desktop/src/lib/mongo/mongoShellCommand.ts"; const collections = ["users", "user_events", "order-items", "audit.logs"]; const fields = [ @@ -492,7 +493,8 @@ test("stays quiet inside string values, comments and unmodelled arguments", () = assert.deepEqual(labels('db.users.find({ name: "Ad', { fields }), []); assert.deepEqual(labels("// db.users.fi", { fields }), []); assert.deepEqual(labels('db.users.find({ _id: ObjectId("6a04', { fields }), []); - assert.deepEqual(labels("db.users.updateOne({}, { $set: {} }, { up", { fields }), []); + // bulkWrite operation entries are still unmodelled; its options argument is not. + assert.deepEqual(labels("db.users.bulkWrite([{ insertOne: { do", { fields }), []); }); test("suggests only helpers the shell parser accepts", () => { @@ -513,6 +515,80 @@ test("suggests only helpers the shell parser accepts", () => { assert.equal(methodLabels.includes("limit"), false); }); +test("suggests the option keys of methods that take an options argument", () => { + assert.deepEqual(labels("db.users.updateOne({}, {$set:{a:1}}, { "), ["arrayFilters", "upsert"]); + assert.deepEqual(labels("db.users.updateOne({}, {$set:{a:1}}, { ups"), ["upsert"]); + assert.deepEqual(labels("db.users.replaceOne({}, {}, { "), ["upsert"]); + assert.deepEqual(labels("db.users.findOneAndDelete({}, { "), ["projection", "sort"]); + assert.deepEqual(labels("db.users.bulkWrite([], { "), ["ordered"]); + assert.deepEqual(labels("db.users.findOne({}, {}, { "), ["sort"]); + assert.ok(labels("db.users.findOneAndUpdate({}, {$set:{a:1}}, { ").includes("returnDocument")); + assert.ok(labels("db.users.createIndex({a:1}, { ").includes("expireAfterSeconds")); + assert.ok(labels("db.users.aggregate([], { ").includes("allowDiskUse")); + + // A `sort` or `projection` option holds field names. + assert.deepEqual(labels("db.users.findOneAndUpdate({}, {$set:{a:1}}, { sort: { na", { fields }), ["name"]); + assert.deepEqual(labels("db.users.findOneAndDelete({}, { projection: { na", { fields }), ["name"]); +}); + +test("stays quiet in the trailing argument of methods that take no options", () => { + // Suggesting an option key here would hand the user a command the parser rejects. + for (const text of ["db.users.find({}, {}, { ", "db.users.insertOne({}, { ", "db.users.insertMany([], { ", "db.users.deleteOne({}, { ", "db.users.deleteMany({}, { ", "db.users.countDocuments({}, { "]) { + assert.deepEqual(labels(text, { fields, collections }), [], text); + } +}); + +test("every suggested option key parses on the method that offers it", () => { + // The option sets mirror the driver's own structs, most of which reject unknown fields, so a + // key that only exists in this table would complete into a command that fails at Run. + const values: Record = { + upsert: "true", + arrayFilters: '[{ "e.f": 1 }]', + returnDocument: '"after"', + returnNewDocument: "true", + new: "true", + projection: "{ name: 1 }", + sort: "{ name: -1 }", + ordered: "false", + name: '"idx"', + unique: "true", + sparse: "true", + expireAfterSeconds: "3600", + partialFilterExpression: "{ name: { $exists: true } }", + collation: '{ locale: "en" }', + hidden: "true", + allowDiskUse: "true", + maxTimeMS: "5000", + hint: '"idx"', + comment: '"why"', + let: "{ n: 1 }", + explain: "true", + }; + const callArgs: Record = { + findOne: "{}, {}", + updateOne: "{}, {$set:{a:1}}", + updateMany: "{}, {$set:{a:1}}", + replaceOne: "{}, {b:1}", + findOneAndUpdate: "{}, {$set:{a:1}}", + findOneAndReplace: "{}, {b:1}", + findOneAndDelete: "{}", + bulkWrite: "[{ insertOne: { document: { a: 1 } } }]", + createIndex: "{a:1}", + aggregate: "[]", + }; + + for (const [method, options] of Object.entries(METHOD_OPTION_KEYS)) { + const args = callArgs[method]; + assert.ok(args !== undefined, `${method} needs sample arguments in this test`); + for (const option of options) { + const value = values[option.label]; + assert.ok(value !== undefined, `${option.label} needs a sample value in this test`); + const command = `db.users.${method}(${args}, { ${option.label}: ${value} })`; + assert.ok(parseMongoCommand(command), `${command} must parse`); + } + } +}); + test("completes both arguments of distinct", () => { const method = buildMongoCompletionItems("db.users.dist", "db.users.dist".length).find((item) => item.label === "distinct"); assert.equal(method?.apply, 'users.distinct("${field}")');