fix(mongodb): complete dotted collection names

This commit is contained in:
zipg 2026-07-20 17:32:38 +08:00 committed by GitHub
parent 1d1959c579
commit 7aff5ea493
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 75 additions and 17 deletions

View File

@ -10,7 +10,7 @@ import { ACCUMULATORS, COMMON_OPERATORS, EXPRESSION_OPERATORS, PIPELINE_STAGES,
* the editor turns into "no popup" deliberately better than falling back to
* `root` and showing unrelated `db.…` snippets mid-document.
*/
export type MongoCompletionMode = "none" | "root" | "collection" | "collectionRef" | "method" | "cursorMethod" | "field" | "fieldPath" | "fieldRef" | "value" | "queryOperator" | "updateOperator" | "pushModifier" | "expression" | "accumulator" | "stage" | "stageOption";
export type MongoCompletionMode = "none" | "root" | "collection" | "collectionOrMethod" | "collectionRef" | "method" | "cursorMethod" | "field" | "fieldPath" | "fieldRef" | "value" | "queryOperator" | "updateOperator" | "pushModifier" | "expression" | "accumulator" | "stage" | "stageOption";
export interface MongoCompletionField {
name: string;
@ -150,8 +150,18 @@ export function getMongoCompletionContext(text: string, cursor: number): MongoCo
if (beforeCursor.endsWith("db.")) return { mode: "collection", prefix: "", from: safeCursor, collection };
const getCollectionPrefix = matchGetCollectionPrefix(beforeCursor);
if (getCollectionPrefix) return { mode: "collectionRef", prefix: getCollectionPrefix.prefix, from: getCollectionPrefix.from, collection };
const collectionPrefix = matchDbCollectionPrefix(beforeCursor);
if (collectionPrefix) return { mode: "collection", prefix: collectionPrefix.prefix, from: collectionPrefix.from, collection };
if (collectionPrefix) {
return {
mode: collectionPrefix.prefix.includes(".") ? "collectionOrMethod" : "collection",
prefix: collectionPrefix.prefix,
from: collectionPrefix.from,
collection,
};
}
if (isAfterCollectionDot(beforeCursor)) {
const methodPrefix = readMethodPrefix(beforeCursor);
@ -190,6 +200,8 @@ export function buildMongoCompletionItemsFromContext(context: MongoCompletionCon
return rootItems(prefix);
case "collection":
return collectionItems(prefix, collections);
case "collectionOrMethod":
return collectionOrMethodItems(prefix, collections);
case "collectionRef":
return collectionRefItems(prefix, collections);
case "method":
@ -230,7 +242,7 @@ export function mongoCompletionNeedsFields(mode: MongoCompletionMode): boolean {
/** Modes whose items are built from the database's collection names. */
export function mongoCompletionNeedsCollections(mode: MongoCompletionMode): boolean {
return mode === "collection" || mode === "collectionRef";
return mode === "collection" || mode === "collectionOrMethod" || mode === "collectionRef";
}
export function shouldAutoOpenMongoCompletion(text: string, cursor: number): boolean {
@ -563,16 +575,7 @@ function rootItems(prefix: string): MongoCompletionItem[] {
}
function collectionItems(prefix: string, collections: string[]): MongoCompletionItem[] {
const names = collections
.filter((collection) => matchesFuzzyPrefix(collection, prefix))
.slice(0, 100)
.map((collection) => ({
label: collection,
type: "table" as const,
detail: "collection",
apply: needsGetCollectionSyntax(collection) ? `getCollection("${escapeDoubleQuoted(collection)}")` : collection,
boost: startsWithPrefix(collection, prefix) ? 120 : 90,
}));
const names = collectionNameItems(prefix, collections);
const methods = DATABASE_METHODS.filter((method) => matchesFuzzyPrefix(method.label, prefix)).map((method) => ({
label: method.label,
type: "function" as const,
@ -583,6 +586,33 @@ function collectionItems(prefix: string, collections: string[]): MongoCompletion
return [...names, ...methods];
}
function collectionNameItems(prefix: string, collections: string[], boost = 120): MongoCompletionItem[] {
return collections
.filter((collection) => matchesFuzzyPrefix(collection, prefix))
.slice(0, 100)
.map((collection) => ({
label: collection,
type: "table" as const,
detail: "collection",
apply: needsGetCollectionSyntax(collection) ? `getCollection("${escapeDoubleQuoted(collection)}")` : collection,
boost: startsWithPrefix(collection, prefix) ? boost : boost - 30,
}));
}
function collectionOrMethodItems(prefix: string, collections: string[]): MongoCompletionItem[] {
const dot = prefix.lastIndexOf(".");
const collection = prefix.slice(0, dot);
const methodPrefix = prefix.slice(dot + 1);
const collectionRef = needsGetCollectionSyntax(collection) && collections.includes(collection) ? `getCollection("${escapeDoubleQuoted(collection)}")` : collection;
const methods = methodItems(methodPrefix).map((item) => ({
...item,
apply: `${collectionRef}.${item.apply}`,
boost: Math.min(item.boost, 110),
}));
return dedupeAndSort([...collectionNameItems(prefix, collections, 150), ...methods]);
}
function collectionRefItems(prefix: string, collections: string[]): MongoCompletionItem[] {
return collections
.filter((collection) => matchesFuzzyPrefix(collection, prefix))
@ -709,7 +739,14 @@ function readMethodPrefix(beforeCursor: string): { prefix: string; from: number
}
function matchDbCollectionPrefix(beforeCursor: string): { prefix: string; from: number } | null {
const match = /(?:^|[\s;(])db\.([A-Za-z_][\w$-]*)$/.exec(beforeCursor);
const match = /(?:^|[\s;(])db\.([A-Za-z_][\w$-]*(?:\.[\w$-]*)*)$/.exec(beforeCursor);
if (!match) return null;
const prefix = match[1] ?? "";
return { prefix, from: beforeCursor.length - prefix.length };
}
function matchGetCollectionPrefix(beforeCursor: string): { prefix: string; from: number } | null {
const match = /(?:^|[\s;(])db\.getCollection\(\s*(["'][^"'\\]*)$/.exec(beforeCursor);
if (!match) return null;
const prefix = match[1] ?? "";
return { prefix, from: beforeCursor.length - prefix.length };

View File

@ -3,7 +3,7 @@ import { test } from "vitest";
import { buildMongoCompletionItems, getMongoCompletionContext, inferMongoCompletionFields, shouldAutoOpenMongoCompletion } from "../../apps/desktop/src/lib/mongo/mongoCompletion.ts";
import { ACCUMULATORS, EXPRESSION_OPERATORS, PIPELINE_STAGES, PUSH_MODIFIERS, QUERY_OPERATORS, STAGE_OPTION_KEYS, UPDATE_OPERATORS, VALUE_SNIPPETS } from "../../apps/desktop/src/lib/mongo/mongoCompletionTables.ts";
const collections = ["users", "user_events", "order-items"];
const collections = ["users", "user_events", "order-items", "audit.logs"];
const fields = [
{ name: "_id", type: "object" },
{ name: "name", type: "string" },
@ -40,6 +40,27 @@ test("uses getCollection apply text for unsafe collection names", () => {
assert.equal(item?.apply, 'getCollection("order-items")');
});
test("continues dotted collection names after a direct db prefix", () => {
const items = buildMongoCompletionItems("db.audit.", "db.audit.".length, { collections });
const dottedCollection = items.find((candidate) => candidate.label === "audit.logs");
assert.equal(dottedCollection?.apply, 'getCollection("audit.logs")');
assert.equal(getMongoCompletionContext("db.audit.", "db.audit.".length).from, "db.".length);
});
test("keeps direct collection method completion while resolving dotted names", () => {
const item = buildMongoCompletionItems("db.users.fi", "db.users.fi".length, { collections }).find((candidate) => candidate.label === "find");
assert.equal(item?.apply, "users.find({})");
});
test("suggests dotted names inside getCollection", () => {
const text = 'db.getCollection("audit.';
const item = buildMongoCompletionItems(text, text.length, { collections }).find((candidate) => candidate.label === "audit.logs");
assert.equal(item?.apply, '"audit.logs"');
});
test("suggests collection methods after direct and getCollection references", () => {
assert.ok(labels("db.users.").includes("find"));
assert.ok(labels('db.getCollection("users").ag').includes("aggregate"));
@ -51,7 +72,7 @@ test("suggests collection stats methods after a collection reference", () => {
assert.ok(methodLabels.includes(method), `expected completion to include ${method}`);
}
const item = buildMongoCompletionItems("db.users.stat", "db.users.stat".length).find((candidate) => candidate.label === "stats");
assert.equal(item?.apply, "stats()");
assert.equal(item?.apply, "users.stats()");
});
test("suggests cursor methods after find result chains", () => {
@ -292,7 +313,7 @@ test("suggests only helpers the shell parser accepts", () => {
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, 'distinct("${field}")');
assert.equal(method?.apply, 'users.distinct("${field}")');
// First argument names a field, so it is completed as a quoted path.
const fieldArg = buildMongoCompletionItems('db.users.distinct("pro', 'db.users.distinct("pro'.length, { fields });