fix(mongodb): complete setwindowfields densify and fill bodies

2026-10-02 02:34:42 +08:00 · 2026-10-01 23:30:31 +08:00
parent 1539c9ac34
commit 192313d33f
3 changed files with 209 additions and 1 deletions
@@ -17,6 +17,7 @@ import {
  TOP_LEVEL_QUERY_OPERATORS,
  UPDATE_OPERATORS,
  VALUE_SNIPPETS,
+  WINDOW_FUNCTION_OPERATORS,
  mongoOperatorItemType,
  type MongoOperatorSpec,
 } from "@/lib/mongo/mongoCompletionTables";
@@ -51,6 +52,7 @@ export type MongoCompletionMode =
  | "pushModifier"
  | "expression"
  | "accumulator"
+  | "windowOperator"
  | "stage"
  | "stageOption"
  | "methodOption"
@@ -318,6 +320,8 @@ const STAGE_OPTION_VALUE_MODES: Record<string, Record<string, MongoCompletionMod
  $setWindowFields: { partitionBy: "fieldRef" },
  $geoNear: { key: "fieldPath" },
  $replaceRoot: { newRoot: "fieldRef" },
+  $densify: { field: "fieldPath" },
+  $fill: { partitionBy: "fieldRef" },
 };

 export function getMongoCompletionContext(text: string, cursor: number): MongoCompletionContext {
@@ -482,6 +486,9 @@ export function buildMongoCompletionItemsFromContext(context: MongoCompletionCon
    case "accumulator":
      items = specItems(ACCUMULATORS, prefix, "accumulator", 100);
      break;
+    case "windowOperator":
+      items = specItems(WINDOW_FUNCTION_OPERATORS, prefix, "window operator", 100);
+      break;
    case "stage": {
      const kind = context.pipelineKind;
      const detail = kind === "update" ? "update stage" : "aggregation stage";
@@ -1252,6 +1259,56 @@ function classifyStageOptions(stage: string, scan: MongoCallScan, bodyIndex: num
    };
  }

+  const optionHolder = scan.stack[bodyIndex + 1];
+  const optionKey = optionHolder?.key ?? "";
+
+  if (optionKey === "sortBy") {
+    return { ...classifyKeyMap(scan, bodyIndex + 1, "sort"), stage };
+  }
+
+  if (optionKey === "output") {
+    if (stage === "$setWindowFields") {
+      if (depth === 1) {
+        if (!scan.inValue) return { mode: "field", stage };
+        return { mode: "none", stage };
+      }
+      if (depth === 2) {
+        if (!scan.inValue) return { mode: "windowOperator", stage };
+        return { mode: "fieldRef", stage };
+      }
+      if (scan.inValue) return { mode: "fieldRef", stage };
+      return { mode: innermost(scan)?.kind === "object" ? "expression" : "none", stage };
+    }
+
+    if (stage === "$fill") {
+      if (depth === 1) {
+        if (!scan.inValue) return { mode: "field", stage };
+        return { mode: "none", stage };
+      }
+      if (depth === 2) {
+        if (!scan.inValue) return { mode: "operatorField", operator: "fillOutput", stage };
+        if (scan.valueKey === "method") return { mode: "enumValue", enumKey: "fillMethod", stage };
+        if (scan.valueKey === "value") return { mode: "fieldRef", stage };
+        return { mode: "none", stage };
+      }
+    }
+  }
+
+  if (optionKey === "range" && stage === "$densify") {
+    if (depth === 1) {
+      if (!scan.inValue) return { mode: "operatorField", operator: "range", stage };
+      if (scan.valueKey === "unit") return { mode: "enumValue", enumKey: "unit", stage };
+      if (scan.valueKey === "bounds") return { mode: "enumValue", enumKey: "bounds", stage };
+      if (scan.valueKey === "step") return { mode: scan.inString ? "none" : "value", stage };
+      return { mode: "none", stage };
+    }
+    return { mode: "none", stage };
+  }
+
+  if (optionKey === "partitionByFields") {
+    if (innermost(scan)?.kind === "array") return { mode: "fieldPath", stage };
+  }
+
  if (scan.inValue) return { mode: "fieldRef", stage };
  return { mode: innermost(scan)?.kind === "object" ? "expression" : "none", stage };
 }
@@ -136,7 +136,7 @@ export const PIPELINE_STAGES: MongoOperatorSpec[] = specs([
  ["$graphLookup", "Performs a recursive search on a collection", '$graphLookup: { from: "${collection}", startWith: "$${field}", connectFromField: "${field}", connectToField: "_id", as: "${as}" }'],
  ["$geoNear", "Orders documents by proximity to a point", '$geoNear: { near: { type: "Point", coordinates: [0, 0] }, distanceField: "${distance}" }'],
  ["$setWindowFields", "Adds fields computed over a window of documents", '$setWindowFields: { partitionBy: "$${field}", sortBy: { ${field}: 1 }, output: {} }'],
-  ["$densify", "Fills gaps in a sequence of field values", '$densify: { field: "${field}", range: { step: 1, unit: "${unit}" } }'],
+  ["$densify", "Fills gaps in a sequence of field values", '$densify: { field: "${field}", range: { step: 1, bounds: "${full}" } }'],
  ["$fill", "Fills null or missing field values", "$fill: { output: { ${field}: { value: ${} } } }"],
  ["$documents", "Returns literal documents", "$documents: [${}]"],
  ["$redact", "Restricts document content based on a condition", "$redact: { ${} }"],
@@ -172,6 +172,37 @@ export const ACCUMULATORS: MongoOperatorSpec[] = specs([
  ["$percentile", "Approximates the requested percentiles", '$percentile: { input: "$${field}", p: [0.95], method: "approximate" }'],
 ]);

+export const WINDOW_OPERATORS: MongoOperatorSpec[] = specs([
+  ["$rank", "Returns the document's rank relative to other documents", "$rank: {}"],
+  ["$denseRank", "Returns the document's dense rank relative to other documents", "$denseRank: {}"],
+  ["$documentNumber", "Returns the document's position within the window", "$documentNumber: {}"],
+  ["$shift", "Returns a value from a document at a given offset", '$shift: { output: "$${field}", by: 1, default: null }'],
+  ["$expMovingAvg", "Exponential moving average", '$expMovingAvg: { input: "$${field}", N: 3 }'],
+  ["$derivative", "Average rate of change over an interval", '$derivative: { input: "$${field}", unit: "${unit}" }'],
+  ["$integral", "Approximation of the area under the curve", '$integral: { input: "$${field}", unit: "${unit}" }'],
+  ["$covariancePop", "Population covariance between two numeric expressions", '$covariancePop: ["$${field1}", "$${field2}"]'],
+  ["$covarianceSamp", "Sample covariance between two numeric expressions", '$covarianceSamp: ["$${field1}", "$${field2}"]'],
+  ["$locf", "Last observation carried forward", '$locf: "$${field}"'],
+  ["$linearFill", "Fills null and missing values by linear interpolation", '$linearFill: "$${field}"'],
+]);
+
+const WINDOW_SPEC: Spec = ["window", "Window boundary specification", 'window: { documents: ["unbounded", "current"] }'];
+
+function dedupeSpecsByLabel(...lists: readonly MongoOperatorSpec[][]): MongoOperatorSpec[] {
+  const seen = new Set<string>();
+  const deduped: MongoOperatorSpec[] = [];
+  for (const list of lists) {
+    for (const spec of list) {
+      if (seen.has(spec.label)) continue;
+      seen.add(spec.label);
+      deduped.push(spec);
+    }
+  }
+  return deduped;
+}
+
+export const WINDOW_FUNCTION_OPERATORS: MongoOperatorSpec[] = [...dedupeSpecsByLabel(WINDOW_OPERATORS, ACCUMULATORS), ...specs([WINDOW_SPEC])];
+
 export const EXPRESSION_OPERATORS: MongoOperatorSpec[] = specs([
  // Conditional
  ["$cond", "Returns one of two values from a condition", "$cond: { if: ${}, then: ${}, else: ${} }"],
@@ -416,6 +447,15 @@ export const OPERATOR_SUB_KEYS: Record<string, MongoOperatorSpec[]> = {
    ["backwards", "Compare secondary differences from the end", "backwards: false"],
    ["normalization", "Normalize text to Unicode NFD first", "normalization: false"],
  ]),
+  range: specs([
+    ["step", "Step size to densify by", "step: 1"],
+    ["unit", "Time unit for date densification", 'unit: "${unit}"'],
+    ["bounds", "Boundary condition for the range", 'bounds: "${full}"'],
+  ]),
+  fillOutput: specs([
+    ["value", "Value or expression to fill with", "value: ${}"],
+    ["method", "Fill method", 'method: "${linear}"'],
+  ]),
 };

 const BSON_TYPE_ALIASES: Spec[] = [
@@ -510,6 +550,25 @@ export const ENUM_VALUES: Record<string, MongoOperatorSpec[]> = {
    ["4", "Also punctuation, with alternate: shifted", "4"],
    ["5", "Identical, including code points", "5"],
  ]),
+  unit: specs([
+    ["millisecond", "Densify by milliseconds", '"millisecond"'],
+    ["second", "Densify by seconds", '"second"'],
+    ["minute", "Densify by minutes", '"minute"'],
+    ["hour", "Densify by hours", '"hour"'],
+    ["day", "Densify by days", '"day"'],
+    ["week", "Densify by weeks", '"week"'],
+    ["month", "Densify by months", '"month"'],
+    ["quarter", "Densify by quarters", '"quarter"'],
+    ["year", "Densify by years", '"year"'],
+  ]),
+  bounds: specs([
+    ["full", "Densify across entire range", '"full"'],
+    ["partition", "Densify within each partition", '"partition"'],
+  ]),
+  fillMethod: specs([
+    ["linear", "Linear interpolation between surrounding values", '"linear"'],
+    ["locf", "Last observation carried forward", '"locf"'],
+  ]),
 };

 export const BULK_WRITE_OPERATIONS: MongoOperatorSpec[] = specs([
@@ -710,6 +769,17 @@ export const STAGE_OPTION_KEYS: Record<string, MongoOperatorSpec[]> = {
    ["includeLocs", "Field holding the matched location", 'includeLocs: "${location}"'],
  ]),
  $replaceRoot: specs([["newRoot", "Expression producing the new root document", 'newRoot: "$${field}"']]),
+  $densify: specs([
+    ["field", "Field whose values will be densified", 'field: "${field}"'],
+    ["partitionByFields", "Fields that partition the documents", 'partitionByFields: ["${field}"]'],
+    ["range", "Range specification to densify", 'range: { step: 1, bounds: "${full}" }'],
+  ]),
+  $fill: specs([
+    ["partitionBy", "Expression to partition by", 'partitionBy: "$${field}"'],
+    ["partitionByFields", "Fields that partition the documents", 'partitionByFields: ["${field}"]'],
+    ["sortBy", "Sort order for evaluating values", "sortBy: { ${field}: 1 }"],
+    ["output", "Fields to fill", "output: { ${field}: { value: ${} } }"],
+  ]),
 };

 /** Literal values that are worth completing in value position. */
@@ -17,6 +17,7 @@ import {
  STAGE_OPTION_KEYS,
  UPDATE_OPERATORS,
  VALUE_SNIPPETS,
+  WINDOW_OPERATORS,
 } from "../../apps/desktop/src/lib/mongo/mongoCompletionTables.ts";
 import { parseMongoCommand } from "../../apps/desktop/src/lib/mongo/mongoShellCommand.ts";

@@ -1098,6 +1099,8 @@ test("every suggested sub-document key and enumerated value parses in the positi
    collation: (body) => `db.users.find({}).collation({ ${body} })`,
    timeseries: (body) => `db.createCollection("x", { timeseries: { ${body} } })`,
    clusteredIndex: (body) => `db.createCollection("x", { clusteredIndex: { ${body} } })`,
+    range: (body) => `db.users.aggregate([{ $densify: { field: "t", range: { ${body} } } }])`,
+    fillOutput: (body) => `db.users.aggregate([{ $fill: { output: { score: { ${body} } } } }])`,
  };
  for (const [operator, keys] of Object.entries(OPERATOR_SUB_KEYS)) {
    const build = keyCommands[operator];
@@ -1123,6 +1126,9 @@ test("every suggested sub-document key and enumerated value parses in the positi
    validationLevel: (value) => `db.createCollection("x", { validationLevel: ${value} })`,
    validationAction: (value) => `db.createCollection("x", { validationAction: ${value} })`,
    granularity: (value) => `db.createCollection("x", { timeseries: { timeField: "t", granularity: ${value} } })`,
+    unit: (value) => `db.users.aggregate([{ $densify: { field: "t", range: { step: 1, unit: ${value} } } }])`,
+    bounds: (value) => `db.users.aggregate([{ $densify: { field: "t", range: { step: 1, bounds: ${value} } } }])`,
+    fillMethod: (value) => `db.users.aggregate([{ $fill: { output: { score: { method: ${value} } } } }])`,
  };
  for (const [enumKey, values] of Object.entries(ENUM_VALUES)) {
    const build = valueCommands[enumKey];
@@ -1221,6 +1227,7 @@ test("snippet templates use placeholder syntax CodeMirror actually honours", ()
    ...PUSH_MODIFIERS,
    ...PIPELINE_STAGES,
    ...ACCUMULATORS,
+    ...WINDOW_OPERATORS,
    ...EXPRESSION_OPERATORS,
    ...VALUE_SNIPPETS,
    ...EXTENDED_JSON_VALUES,
@@ -1279,6 +1286,80 @@ test("infers dotted MongoDB fields from sampled documents", () => {
  assert.ok(inferred.find((field) => field.name === "tags.label" && field.type === "string"));
 });

+test("completes $setWindowFields stage options, sortBy key map, and output window operators", () => {
+  assert.deepEqual(labels("db.users.aggregate([{ $setWindowFields: { "), ["output", "partitionBy", "sortBy"]);
+  assert.deepEqual(labels("db.users.aggregate([{ $setWindowFields: { partitionBy: '", { fields }), ["$_id", "$createdAt", "$name", "$profile.email"]);
+  assert.deepEqual(labels("db.users.aggregate([{ $setWindowFields: { sortBy: { ", { fields }), ["_id", "createdAt", "name", "profile.email"]);
+  assert.deepEqual(labels("db.users.aggregate([{ $setWindowFields: { sortBy: { name: ", { fields }), ["-1", "1"]);
+  assert.deepEqual(labels("db.users.aggregate([{ $setWindowFields: { output: { ", { fields }), ["_id", "createdAt", "name", "profile.email"]);
+
+  const windowOps = labels("db.users.aggregate([{ $setWindowFields: { output: { r: { $");
+  const expectedOperators = [
+    "$rank",
+    "$denseRank",
+    "$documentNumber",
+    "$shift",
+    "$expMovingAvg",
+    "$derivative",
+    "$integral",
+    "$covariancePop",
+    "$covarianceSamp",
+    "$locf",
+    "$linearFill",
+    "$sum",
+    "$avg",
+    "$min",
+    "$max",
+    "$count",
+    "$first",
+    "$last",
+    "$push",
+    "$addToSet",
+    "$stdDevPop",
+    "$stdDevSamp",
+    "$top",
+    "$bottom",
+    "$topN",
+    "$bottomN",
+    "$firstN",
+    "$lastN",
+    "$maxN",
+    "$minN",
+    "$median",
+    "$percentile",
+  ];
+  for (const op of expectedOperators) {
+    assert.ok(windowOps.includes(op), op);
+  }
+  assert.equal(windowOps.includes("window"), false);
+
+  const bareOutputBody = labels("db.users.aggregate([{ $setWindowFields: { output: { r: { ");
+  assert.ok(bareOutputBody.includes("window"));
+  assert.ok(bareOutputBody.includes("$rank"));
+
+  const windowPrefix = labels("db.users.aggregate([{ $setWindowFields: { output: { r: { win");
+  assert.ok(windowPrefix.includes("window"));
+});
+
+test("completes $densify and $fill stage options and value modes", () => {
+  assert.deepEqual(labels("db.users.aggregate([{ $densify: { "), ["field", "partitionByFields", "range"]);
+  assert.deepEqual(labels("db.users.aggregate([{ $densify: { field: '", { fields }), ["_id", "createdAt", "name", "profile.email"]);
+  assert.deepEqual(labels("db.users.aggregate([{ $densify: { partitionByFields: ['", { fields }), ["_id", "createdAt", "name", "profile.email"]);
+  assert.deepEqual(labels("db.users.aggregate([{ $densify: { range: { "), ["bounds", "step", "unit"]);
+  assert.ok(labels("db.users.aggregate([{ $densify: { range: { unit: '").includes("hour"));
+  assert.ok(labels("db.users.aggregate([{ $densify: { range: { bounds: '").includes("full"));
+
+  assert.deepEqual(labels("db.users.aggregate([{ $fill: { "), ["output", "partitionBy", "partitionByFields", "sortBy"]);
+  assert.deepEqual(labels("db.users.aggregate([{ $fill: { partitionBy: '", { fields }), ["$_id", "$createdAt", "$name", "$profile.email"]);
+  assert.deepEqual(labels("db.users.aggregate([{ $fill: { partitionByFields: ['", { fields }), ["_id", "createdAt", "name", "profile.email"]);
+  assert.deepEqual(labels("db.users.aggregate([{ $fill: { sortBy: { ", { fields }), ["_id", "createdAt", "name", "profile.email"]);
+  assert.deepEqual(labels("db.users.aggregate([{ $fill: { sortBy: { name: ", { fields }), ["-1", "1"]);
+  assert.deepEqual(labels("db.users.aggregate([{ $fill: { output: { ", { fields }), ["_id", "createdAt", "name", "profile.email"]);
+  assert.deepEqual(labels("db.users.aggregate([{ $fill: { output: { amount: { "), ["method", "value"]);
+  assert.deepEqual(labels("db.users.aggregate([{ $fill: { output: { amount: { method: '"), ["linear", "locf"]);
+  assert.deepEqual(labels("db.users.aggregate([{ $fill: { output: { amount: { value: '", { fields }), ["$_id", "$createdAt", "$name", "$profile.email"]);
+});
+
 test("completes joined-collection fields inside $lookup, $graphLookup and $unionWith", () => {
  // $lookup foreignField targets joined collection
  const lookupForeign = getMongoCompletionContext('db.users.aggregate([{ $lookup: { from: "orders", localField: "name", foreignField: "', 'db.users.aggregate([{ $lookup: { from: "orders", localField: "name", foreignField: "'.length);