feat: richer column value profiles in schema_inspect (ranked values, uniqueness, value shapes, MongoDB)

packages/drivers/src/mongodb.ts

@@ -389,6 +389,46 @@ export async function connect(config: ConnectionConfig): Promise<Connector> {
     return JSON.stringify(val)
   }
 
+  /**
+   * Raw documents for value profiling. `maxTimeMS` is the smaller of the caller's cap and the
+   * connection's query timeout; `signal` abandons the read and closes the cursor.
+   */
+  async function sampleDocuments(
+    schema: string,
+    table: string,
+    limit = 100,
+    options: { signal?: AbortSignal; maxTimeMS?: number } = {},
+  ): Promise<Record<string, unknown>[]> {
+    options.signal?.throwIfAborted()
+    const configured = resolveQueryTimeout(config).ms
+    const maxTimeMS =
+      options.maxTimeMS === undefined ? configured : configured === undefined ? options.maxTimeMS : Math.min(configured, options.maxTimeMS)
+    const { signal } = options
+    // The driver (6.x) cancels the operation on `signal`; the race below also covers fakes and older drivers.
+    const cursor = client
+      .db(schema)
+      .collection(table)
+      .find({}, { ...withMaxTime(maxTimeMS), ...(signal ? { signal } : {}) })
+      .limit(limit)
+    if (!signal) return cursor.toArray()
+    let abort: (() => void) | undefined
+    const aborted = new Promise<never>((_, reject) => {
+      abort = () => {
+        void Promise.resolve(cursor.close?.()).catch(() => {})
+        reject(signal.reason)
+      }
+      signal.addEventListener("abort", abort, { once: true })
+    })
+    // A read that rejects after the abort won the race must not surface as an unhandled rejection.
+    const read: Promise<Record<string, unknown>[]> = cursor.toArray()
+    read.catch(() => {})
+    try {
+      return await Promise.race([read, aborted])
+    } finally {
+      signal.removeEventListener("abort", abort!)
+    }
+  }
+
   const connector: Connector & {
     // altimate_change start
     runCommand(
@@ -910,6 +950,8 @@ export async function connect(config: ConnectionConfig): Promise<Connector> {
         .sort((a: { name: string }, b: { name: string }) => a.name.localeCompare(b.name))
     },
 
+    sampleDocuments,
+
     async describeTable(schema: string, table: string): Promise<SchemaColumn[]> {
       const db = client.db(schema)
       const coll = db.collection(table)

packages/drivers/src/types.ts

@@ -135,6 +135,13 @@ export interface Connector {
   listSchemas(): Promise<string[]>
   listTables(schema: string): Promise<Array<{ name: string; type: string }>>
   describeTable(schema: string, table: string): Promise<SchemaColumn[]>
+  /** Document stores only: raw documents, bounded by `maxTimeMS` and cancellable through `signal`. */
+  sampleDocuments?(
+    schema: string,
+    table: string,
+    limit?: number,
+    options?: { signal?: AbortSignal; maxTimeMS?: number },
+  ): Promise<Record<string, unknown>[]>
   introspectSchema?(signal?: AbortSignal): Promise<SchemaMetadata>
   close(): Promise<void>
 }

packages/drivers/test/mongodb-unit.test.ts

@@ -10,9 +10,10 @@ const { Binary } = createRequire(import.meta.resolve("mongodb"))("bson")
 
 async function fakeMongo(
   docs?: Record<string, unknown>[],
-  metadata: { collections?: Record<string, string[]>; error?: Error } = {},
+  metadata: { collections?: Record<string, string[]>; error?: Error; read?: () => Promise<Record<string, unknown>[]> } = {},
 ) {
   const finds: unknown[] = []
+  const findOptions: unknown[] = []
   const pipelines: Record<string, unknown>[][] = []
   const aggregateOptions: unknown[] = []
   const counts: { filter: unknown; options: unknown }[] = []
@@ -50,9 +51,10 @@ async function fakeMongo(
               collections[database].push(name)
             },
             collection: () => ({
-              find: (filter: unknown) => {
+              find: (filter: unknown, options: unknown) => {
                 finds.push(filter)
-                return { limit: () => ({ toArray: async () => docs ?? [], async *[Symbol.asyncIterator]() { yield* docs ?? [] } }) }
+                findOptions.push(options)
+                return { limit: () => ({ toArray: () => (metadata.read ? metadata.read() : Promise.resolve(docs ?? [])), async *[Symbol.asyncIterator]() { yield* docs ?? [] } }) }
               },
               aggregate: (pipeline: Record<string, unknown>[], options: unknown) => {
                 pipelines.push(pipeline)
@@ -81,7 +83,7 @@ async function fakeMongo(
   const connector = await connect({ type: "mongodb", database: "test" })
   await connector.connect()
   const query = (command: string, extra: object = {}) => JSON.stringify({ collection: "articles", command, ...extra })
-  return { connector, query, collections, finds, pipelines, aggregateOptions, counts, distincts, collectionLists, databaseLists, estimates }
+  return { connector, query, collections, finds, findOptions, pipelines, aggregateOptions, counts, distincts, collectionLists, databaseLists, estimates }
 }
 
 describe("MongoDB read command shapes", () => {
@@ -473,6 +475,57 @@ describe("MongoDB native command filters", () => {
   })
 })
 
+describe("MongoDB nested documents", () => {
+  const docs = [
+    { _id: 1, user: { name: "a", geo: { lat: 1.5 } }, items: [{ sku: "x" }], "a.b": 1 },
+    { _id: 2, user: null, items: [] },
+  ]
+
+  test("describeTable stays top-level: nested objects are 'object', arrays 'array', no nested paths", async () => {
+    const { connector } = await fakeMongo(docs)
+    expect(await connector.describeTable("test", "articles")).toEqual([
+      { name: "_id", data_type: "int32", nullable: false },
+      { name: "a.b", data_type: "int32", nullable: true },
+      { name: "items", data_type: "array", nullable: false },
+      { name: "user", data_type: "object", nullable: true },
+    ])
+  })
+
+  test("sampleDocuments returns raw documents under the smaller of the caller and connection deadlines", async () => {
+    const { connector, findOptions } = await fakeMongo(docs)
+    expect(await connector.sampleDocuments!("test", "articles", 50, { maxTimeMS: 15_000 })).toEqual(docs)
+    expect(findOptions).toEqual([{ maxTimeMS: 15_000 }])
+  })
+
+  test("sampleDocuments honours an aborted signal", async () => {
+    const { connector } = await fakeMongo(docs)
+    await expect(connector.sampleDocuments!("test", "articles", 50, { signal: AbortSignal.abort(new Error("stop")) })).rejects.toThrow("stop")
+  })
+
+  test("a read rejecting after the abort is handled and the abort reason wins", async () => {
+    const unhandled: unknown[] = []
+    const onUnhandled = (reason: unknown) => unhandled.push(reason)
+    process.on("unhandledRejection", onUnhandled)
+    try {
+      let rejectRead!: (e: Error) => void
+      const { connector, findOptions } = await fakeMongo(docs, {
+        read: () => new Promise((_, reject) => { rejectRead = reject }),
+      })
+      const controller = new AbortController()
+      const pending = connector.sampleDocuments!("test", "articles", 50, { signal: controller.signal })
+      const abortReason = new Error("cancelled")
+      controller.abort(abortReason)
+      await expect(pending).rejects.toBe(abortReason)
+      rejectRead(new Error("cursor closed"))
+      await new Promise((resolve) => setTimeout(resolve, 20))
+      expect(unhandled).toEqual([])
+      expect(findOptions[0]).toMatchObject({ signal: controller.signal })
+    } finally {
+      process.off("unhandledRejection", onUnhandled)
+    }
+  })
+})
+
 describe("MongoDB collection command shorthand", () => {
   test.each([
     { name: "key", aliases: { key: "status" } },

packages/opencode/src/altimate/tools/schema-inspect.ts

@@ -128,6 +128,9 @@ export const SchemaInspectTool = Tool.define("schema_inspect", {
             table: ref,
             columns: schemaResult.columns!,
             warehouse: schemaResult.warehouse ?? target.warehouse,
+            ...(schemaResult.schema_name
+              ? { document: { database: schemaResult.schema_name, collection: schemaResult.table ?? args.table } }
+              : {}),
           },
           ctx,
         )