From d7482ebe06d29f3822ae301e40f27c3a37c83ffc Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Thu, 21 May 2026 15:11:14 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=A7=A9=20fix:=20Support=20DocumentDB=20Pr?= =?UTF-8?q?ompt=20Group=20Lookup=20(#13232)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: support DocumentDB prompt group lookup * test: address prompt group review feedback * test: prove prompt group indexed lookup --- .../src/methods/prompt.getPromptGroup.spec.ts | 229 ++++++++++++++++++ packages/data-schemas/src/methods/prompt.ts | 41 ++-- 2 files changed, 244 insertions(+), 26 deletions(-) create mode 100644 packages/data-schemas/src/methods/prompt.getPromptGroup.spec.ts diff --git a/packages/data-schemas/src/methods/prompt.getPromptGroup.spec.ts b/packages/data-schemas/src/methods/prompt.getPromptGroup.spec.ts new file mode 100644 index 0000000000..3d3b6a74bd --- /dev/null +++ b/packages/data-schemas/src/methods/prompt.getPromptGroup.spec.ts @@ -0,0 +1,229 @@ +import mongoose from 'mongoose'; +import { MongoMemoryServer } from 'mongodb-memory-server'; +import type { PipelineStage } from 'mongoose'; +import { SYSTEM_TENANT_ID, tenantStorage } from '~/config/tenantContext'; +import { createModels, logger } from '..'; +import { createMethods } from './index'; + +logger.silent = true; + +type PromptModel = mongoose.Model; +type PromptGroupModel = mongoose.Model; +type Methods = ReturnType; +type MongoExplainStage = Record; + +let mongoServer: MongoMemoryServer; +let Prompt: PromptModel; +let PromptGroup: PromptGroupModel; +let methods: Methods; + +const TENANT_A = 'tenant-a'; +const TENANT_B = 'tenant-b'; + +async function withTenant(tenantId: string | undefined, fn: () => Promise): Promise { + return tenantStorage.run({ tenantId }, async () => fn()); +} + +async function seedGroupAndPrompt(opts: { tenantId?: string; promptTenantId?: string } = {}) { + const author = new mongoose.Types.ObjectId(); + const prompt = await Prompt.create({ + groupId: new mongoose.Types.ObjectId(), + author, + prompt: 'Hello {{name}}', + type: 'text', + tenantId: opts.promptTenantId ?? opts.tenantId, + }); + const group = await PromptGroup.create({ + name: 'Repro Group', + productionId: prompt._id, + author, + authorName: 'Test Author', + tenantId: opts.tenantId, + }); + return { group, prompt }; +} + +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null && !Array.isArray(value); +} + +function collectRecords(value: unknown, predicate: (record: Record) => boolean) { + const matches: Array> = []; + + function visit(candidate: unknown) { + if (Array.isArray(candidate)) { + for (const item of candidate) { + visit(item); + } + return; + } + + if (!isRecord(candidate)) return; + if (predicate(candidate)) { + matches.push(candidate); + } + + for (const child of Object.values(candidate)) { + visit(child); + } + } + + visit(value); + return matches; +} + +function hasIndexedIdPlan(explain: unknown): boolean { + return collectRecords(explain, (record) => { + const stage = record.stage; + if (stage === 'IDHACK' || stage === 'EXPRESS_IXSCAN') return true; + if (stage !== 'IXSCAN') return false; + + const keyPattern = record.keyPattern; + return isRecord(keyPattern) && keyPattern._id === 1; + }).length > 0; +} + +beforeAll(async () => { + mongoServer = await MongoMemoryServer.create(); + await mongoose.connect(mongoServer.getUri()); + createModels(mongoose); + Prompt = mongoose.models.Prompt; + PromptGroup = mongoose.models.PromptGroup; + methods = createMethods(mongoose, { + removeAllPermissions: async () => {}, + }); +}); + +afterAll(async () => { + await mongoose.disconnect(); + await mongoServer.stop(); +}); + +afterEach(async () => { + await Prompt.deleteMany({}); + await PromptGroup.deleteMany({}); +}); + +describe('getPromptGroup', () => { + it('returns the group with productionPrompt populated when no tenant context is set', async () => { + const { group } = await seedGroupAndPrompt(); + const result = await methods.getPromptGroup({ _id: group._id }); + + expect(result).not.toBeNull(); + expect(result?.productionPrompt).toBeTruthy(); + expect((result?.productionPrompt as { prompt: string }).prompt).toBe('Hello {{name}}'); + }); + + it('returns the group with productionPrompt populated when tenant matches', async () => { + const { group } = await seedGroupAndPrompt({ tenantId: TENANT_A }); + + const result = await withTenant(TENANT_A, () => methods.getPromptGroup({ _id: group._id })); + + expect(result).not.toBeNull(); + expect(result?.productionPrompt).toBeTruthy(); + expect((result?.productionPrompt as { prompt: string }).prompt).toBe('Hello {{name}}'); + }); + + it('clears productionPrompt when the joined prompt belongs to a different tenant', async () => { + const { group } = await seedGroupAndPrompt({ + tenantId: TENANT_A, + promptTenantId: TENANT_B, + }); + + const result = await withTenant(TENANT_A, () => methods.getPromptGroup({ _id: group._id })); + + expect(result).not.toBeNull(); + expect(result?.productionPrompt).toBeNull(); + }); + + it('returns the group with productionPrompt populated under SYSTEM_TENANT_ID context', async () => { + const { group } = await seedGroupAndPrompt({ + tenantId: TENANT_A, + promptTenantId: TENANT_B, + }); + + const result = await withTenant(SYSTEM_TENANT_ID, () => + methods.getPromptGroup({ _id: group._id }), + ); + + expect(result).not.toBeNull(); + expect(result?.productionPrompt).toBeTruthy(); + expect((result?.productionPrompt as { prompt: string }).prompt).toBe('Hello {{name}}'); + }); + + it('returns null when the group does not exist', async () => { + const missingId = new mongoose.Types.ObjectId(); + const result = await methods.getPromptGroup({ _id: missingId }); + expect(result).toBeNull(); + }); + + it('accepts string _id values (route-handler shape)', async () => { + const { group } = await seedGroupAndPrompt(); + const result = await methods.getPromptGroup({ + _id: (group._id as mongoose.Types.ObjectId).toString(), + }); + expect(result).not.toBeNull(); + expect((result?.productionPrompt as { prompt: string }).prompt).toBe('Hello {{name}}'); + }); + + describe('regression: DocumentDB-incompatible aggregation form', () => { + it('does not use $lookup with let/pipeline (multi-join form unsupported by DocumentDB)', async () => { + const { group } = await seedGroupAndPrompt({ tenantId: TENANT_A }); + const aggregateSpy = jest.spyOn(PromptGroup, 'aggregate'); + + try { + await withTenant(TENANT_A, () => methods.getPromptGroup({ _id: group._id })); + + expect(aggregateSpy).toHaveBeenCalledTimes(1); + const pipeline: PipelineStage[] = aggregateSpy.mock.calls[0][0]; + + for (const stage of pipeline) { + if (!('$lookup' in stage)) continue; + const lookup = stage.$lookup; + expect(lookup.let).toBeUndefined(); + expect(lookup.pipeline).toBeUndefined(); + expect(lookup.localField).toBeDefined(); + expect(lookup.foreignField).toBeDefined(); + } + } finally { + aggregateSpy.mockRestore(); + } + }); + + it('uses indexed _id execution paths for the group match and production prompt lookup', async () => { + const { group } = await seedGroupAndPrompt({ tenantId: TENANT_A }); + const aggregateSpy = jest.spyOn(PromptGroup, 'aggregate'); + + try { + await withTenant(TENANT_A, () => methods.getPromptGroup({ _id: group._id })); + + const pipeline: PipelineStage[] = aggregateSpy.mock.calls[0][0]; + const [matchStage] = pipeline; + expect(matchStage).toEqual({ $match: { _id: group._id } }); + + const lookupStage = pipeline.find((stage) => '$lookup' in stage); + expect(lookupStage).toMatchObject({ + $lookup: { + from: 'prompts', + localField: 'productionId', + foreignField: '_id', + as: 'productionPrompt', + }, + }); + + const explain = await PromptGroup.aggregate(pipeline).explain('executionStats'); + const stages = isRecord(explain) && Array.isArray(explain.stages) ? explain.stages : []; + const lookupExplain = stages.find( + (stage): stage is MongoExplainStage => + isRecord(stage) && isRecord(stage.$lookup) && stage.$lookup.from === 'prompts', + ); + + expect(hasIndexedIdPlan(explain)).toBe(true); + expect(lookupExplain?.indexesUsed).toContain('_id_'); + expect(lookupExplain?.totalDocsExamined).toBeLessThanOrEqual(1); + } finally { + aggregateSpy.mockRestore(); + } + }); + }); +}); diff --git a/packages/data-schemas/src/methods/prompt.ts b/packages/data-schemas/src/methods/prompt.ts index 86d830fecd..3a0bd35262 100644 --- a/packages/data-schemas/src/methods/prompt.ts +++ b/packages/data-schemas/src/methods/prompt.ts @@ -512,37 +512,26 @@ export function createPromptMethods(mongoose: typeof import('mongoose'), deps: P const tenantId = getTenantId(); const useTenantFilter = tenantId && tenantId !== SYSTEM_TENANT_ID; - const lookupStage = useTenantFilter - ? { - $lookup: { - from: 'prompts', - let: { prodId: '$productionId' }, - pipeline: [ - { - $match: { - $expr: { $eq: ['$_id', '$$prodId'] }, - tenantId, - }, - }, - ], - as: 'productionPrompt', - }, - } - : { - $lookup: { - from: 'prompts', - localField: 'productionId', - foreignField: '_id', - as: 'productionPrompt', - }, - }; - const result = await PromptGroup.aggregate([ { $match: matchFilter }, - lookupStage, + { + $lookup: { + from: 'prompts', + localField: 'productionId', + foreignField: '_id', + as: 'productionPrompt', + }, + }, { $unwind: { path: '$productionPrompt', preserveNullAndEmptyArrays: true } }, ]); const group = result[0] || null; + if ( + group?.productionPrompt && + useTenantFilter && + group.productionPrompt.tenantId !== tenantId + ) { + group.productionPrompt = null; + } if (group?.author) { group.author = group.author.toString(); }