diff --git a/apps/desktop/src/lib/mongo/mongoCompletion.ts b/apps/desktop/src/lib/mongo/mongoCompletion.ts index 22eaa23073..1b9b2c3bd5 100644 --- a/apps/desktop/src/lib/mongo/mongoCompletion.ts +++ b/apps/desktop/src/lib/mongo/mongoCompletion.ts @@ -1929,18 +1929,23 @@ function extractActiveDatabase(text: string, cursor: number): string | undefined return lastUseDb; } +const MAX_MONGO_COMPLETION_FIELDS = 512; + function collectFieldTypes(value: unknown, prefix: string, out: Map>, depth: number) { if (depth > 4 || value == null || typeof value !== "object") return; // A wrapper such as {$oid: "..."} is one BSON value. Walking into it would offer // `_id.$oid`, a path that exists only in transport and matches nothing on the server. if (mongoExtendedJsonValueType(value)) return; if (Array.isArray(value)) { - for (const item of value.slice(0, 3)) collectFieldTypes(item, prefix, out, depth + 1); + for (const item of value.slice(0, 10)) collectFieldTypes(item, prefix, out, depth + 1); return; } for (const [key, child] of Object.entries(value as Record)) { const path = prefix ? `${prefix}.${key}` : key; - if (!out.has(path)) out.set(path, new Set()); + if (!out.has(path)) { + if (out.size >= MAX_MONGO_COMPLETION_FIELDS) continue; + out.set(path, new Set()); + } out.get(path)?.add(describeMongoValueType(child)); collectFieldTypes(child, path, out, depth + 1); } diff --git a/apps/desktop/src/stores/__tests__/connectionStore.mongoCompletionFields.spec.ts b/apps/desktop/src/stores/__tests__/connectionStore.mongoCompletionFields.spec.ts new file mode 100644 index 0000000000..b6894c2139 --- /dev/null +++ b/apps/desktop/src/stores/__tests__/connectionStore.mongoCompletionFields.spec.ts @@ -0,0 +1,105 @@ +import { createPinia, setActivePinia } from "pinia"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import type { ConnectionConfig } from "@/types/database"; + +function installLocalStorage() { + const data = new Map(); + vi.stubGlobal("localStorage", { + getItem: vi.fn((key: string) => data.get(key) ?? null), + setItem: vi.fn((key: string, value: string) => data.set(key, value)), + removeItem: vi.fn((key: string) => data.delete(key)), + }); +} + +function mongoConnection(): ConnectionConfig { + return { + id: "mongo-1", + name: "MongoDB", + db_type: "mongodb", + host: "127.0.0.1", + port: 27017, + username: "", + password: "", + database: "app", + } as ConnectionConfig; +} + +describe("connectionStore listMongoCompletionFields", () => { + beforeEach(() => { + vi.resetModules(); + vi.unstubAllGlobals(); + installLocalStorage(); + setActivePinia(createPinia()); + }); + + it("samples fields via mongoAggregateDocuments with $sample size 100 and maxTimeMS", async () => { + const api = { + mongoAggregateDocuments: vi.fn().mockResolvedValue({ + documents: [ + { _id: "1", username: "alice", meta: { score: 95 } }, + { _id: "2", email: "alice@example.com" }, + ], + }), + mongoFindDocuments: vi.fn(), + }; + vi.doMock("@/lib/backend/tauriRuntime", () => ({ isTauriRuntime: () => false })); + vi.doMock("@/lib/backend/api", () => api); + + const { useConnectionStore } = await import("@/stores/connectionStore"); + const store = useConnectionStore(); + const config = mongoConnection(); + store.addEphemeralConnection(config); + + const fields = await store.listMongoCompletionFields(config.id, "app", "users"); + + expect(api.mongoAggregateDocuments).toHaveBeenCalledOnce(); + expect(api.mongoAggregateDocuments).toHaveBeenCalledWith(config.id, "app", "users", '[{"$sample":{"size":100}}]', 100, JSON.stringify({ maxTimeMS: 5000 })); + expect(api.mongoFindDocuments).not.toHaveBeenCalled(); + + expect(fields).toEqual([ + { name: "_id", type: "string" }, + { name: "email", type: "string" }, + { name: "meta", type: "object" }, + { name: "meta.score", type: "number" }, + { name: "username", type: "string" }, + ]); + + // Second call should hit the cache without calling api.mongoAggregateDocuments again + const cached = await store.listMongoCompletionFields(config.id, "app", "users"); + expect(cached).toEqual(fields); + expect(api.mongoAggregateDocuments).toHaveBeenCalledOnce(); + }); + + it("falls back to mongoFindDocuments with 100 documents when aggregation fails", async () => { + const api = { + mongoAggregateDocuments: vi.fn().mockRejectedValue(new Error("Aggregation not supported")), + mongoFindDocuments: vi.fn().mockResolvedValue({ + documents: [{ _id: "1", fallbackField: "val", count: 42 }], + }), + }; + vi.doMock("@/lib/backend/tauriRuntime", () => ({ isTauriRuntime: () => false })); + vi.doMock("@/lib/backend/api", () => api); + + const { useConnectionStore } = await import("@/stores/connectionStore"); + const store = useConnectionStore(); + const config = mongoConnection(); + store.addEphemeralConnection(config); + + const fields = await store.listMongoCompletionFields(config.id, "app", "users"); + + expect(api.mongoAggregateDocuments).toHaveBeenCalledOnce(); + expect(api.mongoFindDocuments).toHaveBeenCalledOnce(); + expect(api.mongoFindDocuments).toHaveBeenCalledWith(config.id, "app", "users", 0, 100, "{}"); + + expect(fields).toEqual([ + { name: "_id", type: "string" }, + { name: "count", type: "number" }, + { name: "fallbackField", type: "string" }, + ]); + + // Second call should hit cache + const cached = await store.listMongoCompletionFields(config.id, "app", "users"); + expect(cached).toEqual(fields); + expect(api.mongoFindDocuments).toHaveBeenCalledOnce(); + }); +}); diff --git a/apps/desktop/src/stores/connectionStore.ts b/apps/desktop/src/stores/connectionStore.ts index b884481d4b..8ffbbfd134 100644 --- a/apps/desktop/src/stores/connectionStore.ts +++ b/apps/desktop/src/stores/connectionStore.ts @@ -8759,7 +8759,12 @@ export const useConnectionStore = defineStore("connection", () => { if (cached) return cached; return withCompletionInFlight(`${cacheKey}:mongo-fields`, async () => { await ensureConnected(connectionId); - const result = await api.mongoFindDocuments(connectionId, database, collection, 0, 20, "{}"); + let result; + try { + result = await api.mongoAggregateDocuments(connectionId, database, collection, '[{"$sample":{"size":100}}]', 100, JSON.stringify({ maxTimeMS: 5000 })); + } catch { + result = await api.mongoFindDocuments(connectionId, database, collection, 0, 100, "{}"); + } const fields = inferMongoCompletionFields(result.documents ?? []); mongoCompletionFieldsCache.value[cacheKey] = fields; evictOldestCacheEntries(mongoCompletionFieldsCache.value, COMPLETION_CACHE_MAX); diff --git a/packages/app-tests/mongoCompletion.test.ts b/packages/app-tests/mongoCompletion.test.ts index ac94fdfdc2..13e6dd91b7 100644 --- a/packages/app-tests/mongoCompletion.test.ts +++ b/packages/app-tests/mongoCompletion.test.ts @@ -1286,6 +1286,56 @@ test("infers dotted MongoDB fields from sampled documents", () => { assert.ok(inferred.find((field) => field.name === "tags.label" && field.type === "string")); }); +test("infers MongoDB fields from up to 10 array elements and ignores later elements", () => { + const documents = [ + { + items: [{ e0: 0 }, { e1: 1 }, { e2: 2 }, { e3: 3 }, { e4: 4 }, { e5: 5 }, { e6: 6 }, { e7: 7 }, { e8: 8 }, { e9: 9 }, { e10: 10 }, { e11: 11 }], + }, + ]; + + const inferred = inferMongoCompletionFields(documents); + + // Array path itself + assert.ok(inferred.find((field) => field.name === "items" && field.type === "array")); + // Elements 0 through 9 (first 10 elements) must be sampled + assert.ok(inferred.find((field) => field.name === "items.e0" && field.type === "number")); + assert.ok(inferred.find((field) => field.name === "items.e2" && field.type === "number")); + assert.ok(inferred.find((field) => field.name === "items.e3" && field.type === "number")); + assert.ok(inferred.find((field) => field.name === "items.e9" && field.type === "number")); + // Elements 10 and 11 (beyond the 10-element limit) must NOT be sampled + assert.equal( + inferred.some((field) => field.name === "items.e10"), + false, + ); + assert.equal( + inferred.some((field) => field.name === "items.e11"), + false, + ); +}); + +test("caps inferred MongoDB distinct field paths at 512", () => { + const pathologicalDoc: Record = {}; + for (let i = 0; i < 600; i++) { + pathologicalDoc[`field_${i.toString().padStart(4, "0")}`] = i; + } + + const inferred = inferMongoCompletionFields([pathologicalDoc]); + assert.equal(inferred.length, 512); + + // A second document with existing and new fields: existing fields update types, new fields are ignored + const secondDoc: Record = { + field_0000: "string_value", + brand_new_overflow_field: 42, + }; + const inferredWithSecond = inferMongoCompletionFields([pathologicalDoc, secondDoc]); + assert.equal(inferredWithSecond.length, 512); + assert.equal(inferredWithSecond.find((field) => field.name === "field_0000")?.type, "number | string"); + assert.equal( + inferredWithSecond.some((field) => field.name === "brand_new_overflow_field"), + false, + ); +}); + 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"]);