fix(mongodb): withhold rejected stages in sub pipelines

apps/desktop/src/lib/mongo/mongoCompletion.ts

@@ -60,6 +60,9 @@ export type MongoCompletionMode =
   | "operatorField"
   | "enumValue";
 
+/** Kind of aggregation or update pipeline holding the cursor. */
+export type MongoPipelineKind = "aggregate" | "update" | "facet" | "join" | "view";
+
 export interface MongoCompletionField {
   name: string;
   type?: string;
@@ -88,6 +91,8 @@ export interface MongoCompletionContext {
   database?: string;
   /** Enclosing aggregation stage (`$lookup`, `$group`, …), when inside one. */
   stage?: string;
+  /** Kind of pipeline array holding the cursor, when inside a pipeline. */
+  pipelineKind?: MongoPipelineKind;
   /** Collection method whose options object the cursor sits in. */
   method?: string;
   /** bulkWrite operation (`updateOne`, `deleteMany`, …) whose body the cursor sits in. */
@@ -270,6 +275,32 @@ const OPTION_STAGES = new Set(Object.keys(STAGE_OPTION_KEYS));
 /** The stages an update pipeline accepts: those that rewrite the document without reshaping the result set. */
 const UPDATE_PIPELINE_STAGES = PIPELINE_STAGES.filter((stage) => ["$set", "$addFields", "$unset", "$project", "$replaceRoot", "$replaceWith"].includes(stage.label));
 
+/** Stages that write or merge results, illegal inside sub-pipelines and view definitions. */
+const SUB_PIPELINE_FORBIDDEN_STAGES = new Set(["$out", "$merge"]);
+
+/** Stages MongoDB rejects inside a $facet branch: writes, nesting, metadata/source stages, and $geoNear. */
+const FACET_FORBIDDEN_STAGES = new Set(["$out", "$merge", "$facet", "$collStats", "$indexStats", "$planCacheStats", "$geoNear", "$documents", "$changeStream"]);
+
+/** Stages accepted in join sub-pipelines ($lookup, $unionWith) and view definitions. */
+const SUB_PIPELINE_STAGES = PIPELINE_STAGES.filter((stage) => !SUB_PIPELINE_FORBIDDEN_STAGES.has(stage.label));
+
+/** Stages accepted inside a $facet branch. */
+const FACET_PIPELINE_STAGES = PIPELINE_STAGES.filter((stage) => !FACET_FORBIDDEN_STAGES.has(stage.label));
+
+function pipelineStagesFor(kind?: MongoPipelineKind): MongoOperatorSpec[] {
+  switch (kind) {
+    case "update":
+      return UPDATE_PIPELINE_STAGES;
+    case "facet":
+      return FACET_PIPELINE_STAGES;
+    case "join":
+    case "view":
+      return SUB_PIPELINE_STAGES;
+    default:
+      return PIPELINE_STAGES;
+  }
+}
+
 /** Stages taking a bare `"$field"` string, completed as a field reference. */
 const FIELD_REF_STAGES = new Set(["$unwind", "$sortByCount", "$replaceWith"]);
 
@@ -362,6 +393,7 @@ export function getMongoCompletionContext(text: string, cursor: number): MongoCo
     ...at(classified.mode, classified.stage, classified.method, classified.bulkWriteOperation, classified.keyMap),
     ...(classified.operator ? { operator: classified.operator } : {}),
     ...(classified.enumKey ? { enumKey: classified.enumKey } : {}),
+    ...(classified.pipelineKind ? { pipelineKind: classified.pipelineKind } : {}),
     collection: classified.collection ?? collection,
   };
 }
@@ -437,9 +469,12 @@ export function buildMongoCompletionItemsFromContext(context: MongoCompletionCon
     case "accumulator":
       items = specItems(ACCUMULATORS, prefix, "accumulator", 100);
       break;
-    case "stage":
-      items = context.stage === "update" ? specItems(UPDATE_PIPELINE_STAGES, prefix, "update stage", 100) : specItems(PIPELINE_STAGES, prefix, "aggregation stage", 100);
+    case "stage": {
+      const kind = context.pipelineKind;
+      const detail = kind === "update" ? "update stage" : "aggregation stage";
+      items = specItems(pipelineStagesFor(kind), prefix, detail, 100);
       break;
+    }
     case "stageOption":
       items = specItems(STAGE_OPTION_KEYS[context.stage ?? ""] ?? [], prefix, `${context.stage} option`, 100);
       break;
@@ -654,6 +689,7 @@ interface MongoCallScan {
 interface MongoCursorClass {
   mode: MongoCompletionMode;
   stage?: string;
+  pipelineKind?: MongoPipelineKind;
   collection?: string;
   method?: string;
   bulkWriteOperation?: string;
@@ -967,7 +1003,9 @@ function classifyMethodOptions(method: string, scan: MongoCallScan): MongoCursor
     return { mode: scan.inString ? "none" : "keyMapValue", method, keyMap: option };
   }
   if (FILTER_OPTION_KEYS.has(option)) return { ...classifyFilter(scan, 1), method };
-  if (option === "pipeline" && scan.stack[1]?.kind === "array") return { ...classifyPipeline(scan, 1), method };
+  if (option === "pipeline" && scan.stack[1]?.kind === "array") {
+    return { ...classifyPipeline(scan, findPipelineArrayIndex(scan.stack), method === "createCollection" ? "view" : "aggregate"), method };
+  }
   const valueEnums = SUB_DOCUMENT_OPTION_KEYS[option];
   if (valueEnums) return { ...classifySubDocument(scan, 1, option, valueEnums), method };
   return { mode: "none" };
@@ -1032,21 +1070,34 @@ function classifyBulkWriteOperations(scan: MongoCallScan): MongoCursorClass {
   }
 }
 
+function detectPipelineKind(scan: MongoCallScan, pipelineIndex: number, defaultKind: MongoPipelineKind = "aggregate"): MongoPipelineKind {
+  if (defaultKind === "update") return "update";
+  if (pipelineIndex > 0 && scan.stack[pipelineIndex - 1]?.key === "$facet") return "facet";
+  if (scan.stack[pipelineIndex]?.key === "pipeline") {
+    const parentKey = scan.stack[pipelineIndex - 1]?.key;
+    if (parentKey === "$lookup" || parentKey === "$unionWith") return "join";
+    return defaultKind;
+  }
+  return defaultKind;
+}
+
 /**
- * `kind` is `update` for the pipeline form of an update, which accepts only the
- * stages that rewrite a document; the marker rides in `stage` so the item builder
- * can narrow the list.
+ * `kind` carries the enclosing pipeline flavor (top-level aggregate, update,
+ * $facet branch, join sub-pipeline or view pipeline) so the item builder can
+ * narrow the stage list.
  */
-function classifyPipeline(scan: MongoCallScan, pipelineIndex = findPipelineArrayIndex(scan.stack), kind: "aggregate" | "update" = "aggregate"): MongoCursorClass {
+function classifyPipeline(scan: MongoCallScan, pipelineIndex = findPipelineArrayIndex(scan.stack), kind?: MongoPipelineKind): MongoCursorClass {
   if (pipelineIndex < 0) return { mode: "none" };
 
   const stageHolder = scan.stack[pipelineIndex + 1];
   if (!stageHolder) return { mode: "none" }; // directly inside the array, no stage object yet
   if (stageHolder.kind !== "object") return { mode: "none" };
 
+  const pipelineKind = detectPipelineKind(scan, pipelineIndex, kind);
+
   // `[{ … }]` — the cursor is in the stage object itself.
   if (scan.stack.length - 1 === pipelineIndex + 1) {
-    if (!scan.inValue) return { mode: "stage", stage: kind === "update" ? "update" : undefined };
+    if (!scan.inValue) return { mode: "stage", pipelineKind };
     const stage = scan.valueKey ?? "";
     return { mode: stageStringValueMode(stage), stage };
   }

packages/app-tests/mongoCompletion.test.ts

@@ -523,6 +523,108 @@ test("suggests stages only at pipeline level, including nested pipelines", () =>
   assert.equal(labels("db.users.aggregate([{ $match: { tags: { $in: [{ $m").includes("$match"), false);
 });
 
+test("withholds illegal stages in sub-pipelines and view definitions", () => {
+  // Top-level aggregate includes $out and $merge.
+  const topLevel = labels("db.users.aggregate([{ $");
+  assert.ok(topLevel.includes("$out"));
+  assert.ok(topLevel.includes("$merge"));
+  assert.ok(topLevel.includes("$match"));
+  assert.ok(topLevel.includes("$project"));
+  assert.ok(topLevel.includes("$group"));
+
+  // $facet branch rejects $out, $merge, $facet, $collStats, $indexStats, $planCacheStats, $geoNear, $documents.
+  const facetStages = labels("db.users.aggregate([{ $facet: { a: [{ $");
+  for (const forbidden of ["$out", "$merge", "$facet", "$collStats", "$indexStats", "$planCacheStats", "$geoNear", "$documents"]) {
+    assert.equal(facetStages.includes(forbidden), false, `$facet branch must not offer ${forbidden}`);
+  }
+  assert.ok(facetStages.includes("$match"));
+  assert.ok(facetStages.includes("$project"));
+  assert.ok(facetStages.includes("$group"));
+
+  // $lookup sub-pipeline rejects $out and $merge, but permits $documents, $facet, $geoNear.
+  const lookupStages = labels('db.users.aggregate([{ $lookup: { from: "orders", pipeline: [{ $');
+  assert.equal(lookupStages.includes("$out"), false, "$lookup must not offer $out");
+  assert.equal(lookupStages.includes("$merge"), false, "$lookup must not offer $merge");
+  assert.ok(lookupStages.includes("$match"));
+  assert.ok(lookupStages.includes("$project"));
+  assert.ok(lookupStages.includes("$group"));
+  assert.ok(lookupStages.includes("$documents"));
+  assert.ok(lookupStages.includes("$facet"));
+  assert.ok(lookupStages.includes("$geoNear"));
+
+  // $unionWith sub-pipeline rejects $out and $merge, but permits $documents.
+  const unionStages = labels('db.users.aggregate([{ $unionWith: { coll: "orders", pipeline: [{ $');
+  assert.equal(unionStages.includes("$out"), false, "$unionWith must not offer $out");
+  assert.equal(unionStages.includes("$merge"), false, "$unionWith must not offer $merge");
+  assert.ok(unionStages.includes("$match"));
+  assert.ok(unionStages.includes("$project"));
+  assert.ok(unionStages.includes("$group"));
+  assert.ok(unionStages.includes("$documents"));
+
+  // View definition pipeline rejects $out and $merge, but permits $documents.
+  const viewStages = labels('db.createCollection("v", { viewOn: "users", pipeline: [{ $');
+  assert.equal(viewStages.includes("$out"), false, "view pipeline must not offer $out");
+  assert.equal(viewStages.includes("$merge"), false, "view pipeline must not offer $merge");
+  assert.ok(viewStages.includes("$match"));
+  assert.ok(viewStages.includes("$project"));
+  assert.ok(viewStages.includes("$group"));
+  assert.ok(viewStages.includes("$documents"));
+
+  // Sub-pipeline nested inside a view pipeline applies sub-pipeline restrictions.
+  const viewFacetStages = labels('db.createCollection("v", { viewOn: "users", pipeline: [{ $facet: { a: [{ $');
+  assert.equal(viewFacetStages.includes("$out"), false);
+  assert.equal(viewFacetStages.includes("$merge"), false);
+  assert.equal(viewFacetStages.includes("$facet"), false);
+  assert.equal(viewFacetStages.includes("$geoNear"), false);
+  assert.ok(viewFacetStages.includes("$match"));
+
+  // runCommand aggregate pipeline is a top-level aggregate and offers $out and $merge.
+  const runCommandStages = labels('db.runCommand({ aggregate: "users", pipeline: [{ $');
+  assert.ok(runCommandStages.includes("$out"));
+  assert.ok(runCommandStages.includes("$merge"));
+  assert.ok(runCommandStages.includes("$match"));
+
+  // Update pipeline remains unchanged.
+  const updateStages = labels("db.users.updateOne({}, [{ $", { fields });
+  assert.deepEqual([...updateStages].sort(), ["$addFields", "$project", "$replaceRoot", "$replaceWith", "$set", "$unset"]);
+});
+
+test("classifies pipeline kinds across top-level, sub-pipelines, views and updates", () => {
+  const topText = "db.users.aggregate([{ $";
+  assert.equal(getMongoCompletionContext(topText, topText.length).pipelineKind, "aggregate");
+  assert.equal(getMongoCompletionContext(topText, topText.length).stage, undefined);
+
+  const facetText = "db.users.aggregate([{ $facet: { a: [{ $";
+  const facetCtx = getMongoCompletionContext(facetText, facetText.length);
+  assert.equal(facetCtx.pipelineKind, "facet");
+  assert.equal(facetCtx.stage, undefined);
+
+  const lookupText = 'db.users.aggregate([{ $lookup: { from: "orders", pipeline: [{ $';
+  const lookupCtx = getMongoCompletionContext(lookupText, lookupText.length);
+  assert.equal(lookupCtx.pipelineKind, "join");
+  assert.equal(lookupCtx.stage, undefined);
+
+  const unionText = 'db.users.aggregate([{ $unionWith: { coll: "orders", pipeline: [{ $';
+  const unionCtx = getMongoCompletionContext(unionText, unionText.length);
+  assert.equal(unionCtx.pipelineKind, "join");
+  assert.equal(unionCtx.stage, undefined);
+
+  const viewText = 'db.createCollection("v", { viewOn: "users", pipeline: [{ $';
+  const viewCtx = getMongoCompletionContext(viewText, viewText.length);
+  assert.equal(viewCtx.pipelineKind, "view");
+  assert.equal(viewCtx.stage, undefined);
+
+  const runCommandText = 'db.runCommand({ aggregate: "users", pipeline: [{ $';
+  const runCommandCtx = getMongoCompletionContext(runCommandText, runCommandText.length);
+  assert.equal(runCommandCtx.pipelineKind, "aggregate");
+  assert.equal(runCommandCtx.stage, undefined);
+
+  const updateText = "db.users.updateOne({}, [{ $";
+  const updateCtx = getMongoCompletionContext(updateText, updateText.length);
+  assert.equal(updateCtx.pipelineKind, "update");
+  assert.equal(updateCtx.stage, undefined);
+});
+
 test("suggests values, not fields, after a filter key", () => {
   const items = labels("db.users.find({ _id: ", { fields });