feat(mongodb): complete collection method option completion keys

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<string, readonly MongoArgRole[]> = {
-  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<string, readonly MongoArgRole[]> = {
   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 `{ <op>: { 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" };

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<string, MongoOperatorSpec[]> = {
+  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 $[<identifier>] positional operator", 'arrayFilters: [{ "${elem}.${field}": ${} }]'],
+  ]),
+  updateMany: specs([
+    ["upsert", "Insert the document when the filter matches nothing", "upsert: true"],
+    ["arrayFilters", "Conditions for the $[<identifier>] 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 $[<identifier>] 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<string, MongoOperatorSpec[]> = {
   $lookup: specs([

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<string, string> = {
+    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<string, string> = {
+    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}")');