From 3d56f3b0dcd94912b96a47cbfb022815d13d1cb3 Mon Sep 17 00:00:00 2001
From: Kevin Messiaen <114553769+kevinmessiaen@users.noreply.github.com>
Date: Mon, 6 Jul 2026 16:15:40 +0700
Subject: [PATCH] feat(mongodb): add aggregate query surface to the connector

---
 .../cli/src/connectors/mongodb/connector.ts   | 45 +++++++++++++++++++
 .../test/connectors/mongodb/connector.test.ts | 37 +++++++++++++++
 2 files changed, 82 insertions(+)

diff --git a/packages/cli/src/connectors/mongodb/connector.ts b/packages/cli/src/connectors/mongodb/connector.ts
index f46d6038..12350835 100644
--- a/packages/cli/src/connectors/mongodb/connector.ts
+++ b/packages/cli/src/connectors/mongodb/connector.ts
@@ -39,6 +39,20 @@ export interface KtxMongoListedCollection {
   type?: string;
 }
 
+export interface KtxMongoQueryInput {
+  connectionId: string;
+  collection: string;
+  database?: string;
+  pipeline: Record<string, unknown>[];
+  limit: number;
+}
+
+export interface KtxMongoQueryResult {
+  headers: string[];
+  rows: unknown[][];
+  rowCount: number;
+}
+
 interface KtxMongoFindOptions {
   sort: Record<string, 1 | -1>;
   limit: number;
@@ -50,6 +64,12 @@ export interface KtxMongoClient {
   listCollections(databaseName: string): Promise<KtxMongoListedCollection[]>;
   estimatedDocumentCount(databaseName: string, collectionName: string): Promise<number>;
   find(databaseName: string, collectionName: string, options: KtxMongoFindOptions): Promise<KtxMongoDocument[]>;
+  aggregate(
+    databaseName: string,
+    collectionName: string,
+    pipeline: Record<string, unknown>[],
+    options: { limit: number },
+  ): Promise<KtxMongoDocument[]>;
   ping(databaseName: string): Promise<void>;
   close(): Promise<void>;
 }
@@ -106,6 +126,20 @@ class DefaultMongoClient implements KtxMongoClient {
       .toArray() as Promise<KtxMongoDocument[]>;
   }
 
+  async aggregate(
+    databaseName: string,
+    collectionName: string,
+    pipeline: Record<string, unknown>[],
+    options: { limit: number },
+  ): Promise<KtxMongoDocument[]> {
+    const client = await this.connectedClient();
+    return client
+      .db(databaseName)
+      .collection(collectionName)
+      .aggregate([...pipeline, { $limit: options.limit }], { maxTimeMS: SAMPLE_MAX_TIME_MS })
+      .toArray() as Promise<KtxMongoDocument[]>;
+  }
+
   async ping(databaseName: string): Promise<void> {
     const client = await this.connectedClient();
     await client.db(databaseName).command({ ping: 1 });
@@ -362,6 +396,17 @@ export class KtxMongoDbScanConnector implements KtxScanConnector {
     return { values, nullCount, distinctCount: null };
   }
 
+  async executeQuery(input: KtxMongoQueryInput, _ctx: KtxScanContext): Promise<KtxMongoQueryResult> {
+    this.assertConnection(input.connectionId);
+    const database = input.database ?? this.databases[0]!;
+    const documents = await this.clientForQuery().aggregate(database, input.collection, input.pipeline, {
+      limit: input.limit,
+    });
+    const headers = unionDocumentKeys(documents);
+    const rows = documents.map((document) => headers.map((header) => normalizeSampleValue(document[header])));
+    return { headers, rows, rowCount: rows.length };
+  }
+
   async listSchemas(): Promise<string[]> {
     return [...this.databases];
   }
diff --git a/packages/cli/test/connectors/mongodb/connector.test.ts b/packages/cli/test/connectors/mongodb/connector.test.ts
index 43d60013..3fd25039 100644
--- a/packages/cli/test/connectors/mongodb/connector.test.ts
+++ b/packages/cli/test/connectors/mongodb/connector.test.ts
@@ -40,6 +40,7 @@ function fakeClientFactory(): { factory: KtxMongoClientFactory; client: KtxMongo
       collectionName === 'users' ? 2 : 1,
     ),
     find: vi.fn(async (_databaseName: string, collectionName: string) => DOCUMENTS[collectionName] ?? []),
+    aggregate: vi.fn(async (_databaseName: string, collectionName: string) => DOCUMENTS[collectionName] ?? []),
     ping: vi.fn(async () => undefined),
     close: vi.fn(async () => undefined),
   };
@@ -184,6 +185,7 @@ describe('KtxMongoDbScanConnector.introspect', () => {
         return 2;
       }),
       find: vi.fn(async (_db: string, name: string) => DOCUMENTS[name] ?? DOCUMENTS.users!),
+      aggregate: vi.fn(async (_db: string, name: string) => DOCUMENTS[name] ?? DOCUMENTS.users!),
       ping: vi.fn(async () => undefined),
       close: vi.fn(async () => undefined),
     };
@@ -239,3 +241,38 @@ describe('ktx sql against a MongoDB connection', () => {
     ).rejects.toThrow(/does not support read-only SQL execution/);
   });
 });
+
+describe('KtxMongoDbScanConnector.executeQuery', () => {
+  it('runs an aggregation pipeline and returns normalized rows keyed by the union of document fields', async () => {
+    const { factory, client } = fakeClientFactory();
+    const result = await connector(baseConnection, factory).executeQuery(
+      {
+        connectionId: 'mongo-prod',
+        collection: 'users',
+        pipeline: [{ $match: { age: { $gte: 18 } } }],
+        limit: 50,
+      },
+      { runId: 't' },
+    );
+
+    expect(client.aggregate).toHaveBeenCalledWith('app', 'users', [{ $match: { age: { $gte: 18 } } }], {
+      limit: 50,
+    });
+    expect(result.headers).toEqual(['_id', 'email', 'age', 'address']);
+    // _id is a BSON wrapper → String(...); address is a sub-document → JSON string.
+    expect(result.rows[0]).toEqual(['a1', 'a@x.com', 31, JSON.stringify({ city: 'NY' })]);
+    // Second document is missing age and address → null-filled to keep rows rectangular.
+    expect(result.rows[1]).toEqual(['a2', 'b@x.com', null, null]);
+    expect(result.rowCount).toBe(2);
+  });
+
+  it('rejects a query whose connection id does not match this connector', async () => {
+    const { factory } = fakeClientFactory();
+    await expect(
+      connector(baseConnection, factory).executeQuery(
+        { connectionId: 'other', collection: 'users', pipeline: [], limit: 10 },
+        { runId: 't' },
+      ),
+    ).rejects.toThrow(/cannot serve connection other/);
+  });
+});