feat(completion): sample mongo fields from random documents

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<string, Set<string>>, 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<string, unknown>)) {
     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);
   }

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<string, string>();
+  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();
+  });
+});

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);

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