From f69223230ad6fb06e7db71b1c048f2aaa1228fb5 Mon Sep 17 00:00:00 2001 From: Benjamin Blanchard Date: Thu, 27 Aug 2026 17:28:14 -0400 Subject: [PATCH 1/2] remove segments joins from assign path --- .../api/repositories/ExperimentRepository.ts | 275 +++++++----------- .../services/ExperimentAssignmentService.ts | 75 ++--- packages/backend/test/unit/mockdata/raw.ts | 4 + .../repositories/ExperimentRepository.test.ts | 160 +++++----- .../ExperimentAssignmentService.test.ts | 26 ++ 5 files changed, 263 insertions(+), 277 deletions(-) diff --git a/packages/backend/src/api/repositories/ExperimentRepository.ts b/packages/backend/src/api/repositories/ExperimentRepository.ts index 61cf65fa3..f1cfd1dcc 100644 --- a/packages/backend/src/api/repositories/ExperimentRepository.ts +++ b/packages/backend/src/api/repositories/ExperimentRepository.ts @@ -1,5 +1,5 @@ import { EXPERIMENT_STATE, SERVER_ERROR } from 'upgrade_types'; -import { Repository, EntityManager, Brackets } from 'typeorm'; +import { Repository, EntityManager, Brackets, SelectQueryBuilder } from 'typeorm'; import { EntityRepository } from '../../typeorm-typedi-extensions'; import { Experiment } from '../models/Experiment'; import repositoryError from './utils/repositoryError'; @@ -155,24 +155,7 @@ export class ExperimentRepository extends Repository { }) ); - const experimentInclusionSegmentQuery = this.buildInclusionSegmentQuery().where( - new Brackets((qb) => { - qb.where(whereExperimentsClause, whereClauseParams); - }) - ); - - const experimentExclusionSegmentQuery = this.buildExclusionSegmentQuery().where( - new Brackets((qb) => { - qb.where(whereExperimentsClause, whereClauseParams); - }) - ); - - const [ - experimentConditionLevelPayloadData, - experimentFactorDecisionPointLevelPayloadData, - experimentInclusionSegmentData, - experimentExclusionSegmentData, - ] = await Promise.all([ + const [experimentConditionLevelPayloadData, experimentFactorDecisionPointLevelPayloadData] = await Promise.all([ experimentConditionLevelPayloadQuery.getMany().catch((errorMsg: any) => { const errorMsgString = repositoryError( 'ExperimentRepository', @@ -191,42 +174,19 @@ export class ExperimentRepository extends Repository { ); throw errorMsgString; }), - experimentInclusionSegmentQuery.getMany().catch((errorMsg: any) => { - const errorMsgString = repositoryError( - 'ExperimentRepository', - 'getValidExperiments-experimentInclusionSegmentQuery', - {}, - errorMsg - ); - throw errorMsgString; - }), - experimentExclusionSegmentQuery.getMany().catch((errorMsg: any) => { - const errorMsgString = repositoryError( - 'ExperimentRepository', - 'getValidExperiments-experimentExclusionSegmentQuery', - {}, - errorMsg - ); - throw errorMsgString; - }), ]); - const experimentSegmentData = this.mergeSegmentData(experimentInclusionSegmentData, experimentExclusionSegmentData); - - const experimentData = experimentConditionLevelPayloadData.map((data) => { + // Inclusion/exclusion segment ids are intentionally NOT joined in here: they're only ever needed + // by ExperimentAssignmentService for an experiment that has no row in the experiment_precomputed_segment + // table, which is the uncommon case (not an in-memory cache miss — the in-memory cache transparently + // re-reads Postgres for experiments that DO have a row, regardless of cache state; only a genuine + // absence of the row itself, whether freshly confirmed or served from a cached negative result, counts). + // See ExperimentRepository.getSegmentIdsForExperiments, which fetches them lazily and scoped to just + // the experiment ids that actually need them. + return experimentConditionLevelPayloadData.map((data) => { const data2 = experimentFactorDecisionPointLevelPayloadData.find((i) => i.id === data.id); return { ...data, ...data2 }; }); - - const mergedData = experimentData.map((data) => { - const { id } = data; - const segmentData = experimentSegmentData.find((segmentData) => { - return segmentData.id === id; - }); - return segmentData ? { ...data, ...segmentData } : data; - }); - - return mergedData; } public async getValidExperimentsForContextAndDecisionPoint( @@ -262,73 +222,33 @@ export class ExperimentRepository extends Repository { }) ); - const inclusionSegmentQuery = this.buildInclusionSegmentQuery() - .leftJoin('experiment.partitions', 'partitions') - .where( - new Brackets((qb) => { - qb.where(decisionPointWhereClause, whereClauseParams); - }) - ); - - const exclusionSegmentQuery = this.buildExclusionSegmentQuery() - .leftJoin('experiment.partitions', 'partitions') - .where( - new Brackets((qb) => { - qb.where(decisionPointWhereClause, whereClauseParams); - }) - ); - - const [conditionLevelPayloadData, factorDecisionPointPayloadData, inclusionSegmentData, exclusionSegmentData] = - await Promise.all([ - conditionLevelPayloadQuery.getMany().catch((errorMsg: any) => { - const errorMsgString = repositoryError( - 'ExperimentRepository', - 'getValidExperimentsForContextAndDecisionPoint-conditionLevelPayloadData', - {}, - errorMsg - ); - throw errorMsgString; - }), - factorDecisionPointPayloadQuery.getMany().catch((errorMsg: any) => { - const errorMsgString = repositoryError( - 'ExperimentRepository', - 'getValidExperimentsForContextAndDecisionPoint-factorDecisionPointPayloadData', - {}, - errorMsg - ); - throw errorMsgString; - }), - inclusionSegmentQuery.getMany().catch((errorMsg: any) => { - const errorMsgString = repositoryError( - 'ExperimentRepository', - 'getValidExperimentsForContextAndDecisionPoint-inclusionSegmentData', - {}, - errorMsg - ); - throw errorMsgString; - }), - exclusionSegmentQuery.getMany().catch((errorMsg: any) => { - const errorMsgString = repositoryError( - 'ExperimentRepository', - 'getValidExperimentsForContextAndDecisionPoint-exclusionSegmentData', - {}, - errorMsg - ); - throw errorMsgString; - }), - ]); - - const segmentData = this.mergeSegmentData(inclusionSegmentData, exclusionSegmentData); + const [conditionLevelPayloadData, factorDecisionPointPayloadData] = await Promise.all([ + conditionLevelPayloadQuery.getMany().catch((errorMsg: any) => { + const errorMsgString = repositoryError( + 'ExperimentRepository', + 'getValidExperimentsForContextAndDecisionPoint-conditionLevelPayloadData', + {}, + errorMsg + ); + throw errorMsgString; + }), + factorDecisionPointPayloadQuery.getMany().catch((errorMsg: any) => { + const errorMsgString = repositoryError( + 'ExperimentRepository', + 'getValidExperimentsForContextAndDecisionPoint-factorDecisionPointPayloadData', + {}, + errorMsg + ); + throw errorMsgString; + }), + ]); - const experimentData = factorDecisionPointPayloadData.map((data) => { + // See the comment in getValidExperiments: inclusion/exclusion segment ids are fetched lazily via + // getSegmentIdsForExperiments instead of being joined in here. + return factorDecisionPointPayloadData.map((data) => { const condData = conditionLevelPayloadData.find((i) => i.id === data.id); return { ...condData, ...data }; }); - - return experimentData.map((data) => { - const seg = segmentData.find((s) => s.id === data.id); - return seg ? { ...data, ...seg } : data; - }); } public async findFirstValidContextByDecisionPoint(site: string, target: string): Promise { @@ -389,24 +309,7 @@ export class ExperimentRepository extends Repository { }) ); - const experimentInclusionSegmentQuery = this.buildInclusionSegmentQuery().where( - new Brackets((qb) => { - qb.where(whereExperimentsClause, whereClauseParams); - }) - ); - - const experimentExclusionSegmentQuery = this.buildExclusionSegmentQuery().where( - new Brackets((qb) => { - qb.where(whereExperimentsClause, whereClauseParams); - }) - ); - - const [ - experimentConditionLevelPayloadData, - experimentFactorDecisionPointLevelPayloadData, - experimentInclusionSegmentData, - experimentExclusionSegmentData, - ] = await Promise.all([ + const [experimentConditionLevelPayloadData, experimentFactorDecisionPointLevelPayloadData] = await Promise.all([ experimentConditionLevelPayloadQuery.getMany().catch((errorMsg: any) => { const errorMsgString = repositoryError( 'ExperimentRepository', @@ -425,42 +328,14 @@ export class ExperimentRepository extends Repository { ); throw errorMsgString; }), - experimentInclusionSegmentQuery.getMany().catch((errorMsg: any) => { - const errorMsgString = repositoryError( - 'ExperimentRepository', - 'getValidExperimentsWithPreview-experimentInclusionSegmentQuery', - {}, - errorMsg - ); - throw errorMsgString; - }), - experimentExclusionSegmentQuery.getMany().catch((errorMsg: any) => { - const errorMsgString = repositoryError( - 'ExperimentRepository', - 'getValidExperimentsWithPreview-experimentExclusionSegmentQuery', - {}, - errorMsg - ); - throw errorMsgString; - }), ]); - const experimentSegmentData = this.mergeSegmentData(experimentInclusionSegmentData, experimentExclusionSegmentData); - - const experimentData = experimentConditionLevelPayloadData.map((data) => { + // See the comment in getValidExperiments: inclusion/exclusion segment ids are fetched lazily via + // getSegmentIdsForExperiments instead of being joined in here. + return experimentConditionLevelPayloadData.map((data) => { const data2 = experimentFactorDecisionPointLevelPayloadData.find((i) => i.id === data.id); return { ...data, ...data2 }; }); - - const mergedData = experimentData.map((data) => { - const { id } = data; - const segmentData = experimentSegmentData.find((segmentData) => { - return segmentData.id === id; - }); - return segmentData ? { ...data, ...segmentData } : data; - }); - - return mergedData; } public async updateState( @@ -623,6 +498,82 @@ export class ExperimentRepository extends Repository { .leftJoinAndSelect('segmentExclusion.subSegments', 'subSegmentExclusion'); } + // Lightweight variants for the assignment read path: only the referenced segment `id`s are ever + // consumed downstream (ExperimentAssignmentService reads `.segmentId` off the junction row); + // actual membership resolution comes either from the experiment_precomputed_segment table (keyed + // by experiment id, not these relations) or, when an experiment has no row there, from a fresh + // independent lookup via SegmentService.getSegmentByIds. So the deep + // individualForSegment/groupForSegment/subSegments hydration above is never read on this path. No + // join to the `segment` table is needed either: `segmentId` is already a plain column on the + // junction entity itself (ExperimentSegmentInclusion/ExperimentSegmentExclusion), populated via + // their OneToOne @JoinColumn. + // + // These are only ever invoked scoped to a specific set of experiment ids (see + // getSegmentIdsForExperiments below), not eagerly for every valid experiment: since the data is + // usually unused (the experiment already has a precomputed row), joining it in for every experiment + // on every getValidExperiments*/-ForContextAndDecisionPoint/-WithPreview call would be wasted work. + // Return types are asserted to the narrowed fragment shape (rather than the inferred + // `SelectQueryBuilder`) because only `id` and the joined relation are actually + // selected/hydrated here; the rest of `Experiment`'s fields are never populated by this query. + private buildInclusionSegmentIdQuery(): SelectQueryBuilder { + return this.createQueryBuilder('experiment') + .select('experiment.id') + .leftJoinAndSelect( + 'experiment.experimentSegmentInclusion', + 'experimentSegmentInclusion' + ) as SelectQueryBuilder; + } + + private buildExclusionSegmentIdQuery(): SelectQueryBuilder { + return this.createQueryBuilder('experiment') + .select('experiment.id') + .leftJoinAndSelect( + 'experiment.experimentSegmentExclusion', + 'experimentSegmentExclusion' + ) as SelectQueryBuilder; + } + + /** + * Fetch inclusion/exclusion segment ids for a specific, known set of experiment ids. Intended to be + * called lazily by ExperimentAssignmentService only for experiments that have no row in the + * experiment_precomputed_segment table, rather than eagerly for the full valid-experiment set. + */ + public async getSegmentIdsForExperiments(experimentIds: string[]): Promise { + if (experimentIds.length === 0) { + return []; + } + + const inclusionSegmentQuery = this.buildInclusionSegmentIdQuery().where('experiment.id IN (:...experimentIds)', { + experimentIds, + }); + const exclusionSegmentQuery = this.buildExclusionSegmentIdQuery().where('experiment.id IN (:...experimentIds)', { + experimentIds, + }); + + const [inclusionSegmentData, exclusionSegmentData] = await Promise.all([ + inclusionSegmentQuery.getMany().catch((errorMsg: any) => { + const errorMsgString = repositoryError( + 'ExperimentRepository', + 'getSegmentIdsForExperiments-inclusionSegmentQuery', + { experimentIds }, + errorMsg + ); + throw errorMsgString; + }), + exclusionSegmentQuery.getMany().catch((errorMsg: any) => { + const errorMsgString = repositoryError( + 'ExperimentRepository', + 'getSegmentIdsForExperiments-exclusionSegmentQuery', + { experimentIds }, + errorMsg + ); + throw errorMsgString; + }), + ]); + + return this.mergeSegmentData(inclusionSegmentData, exclusionSegmentData); + } + private mergeSegmentData( inclusionData: SegmentInclusionFragment[], exclusionData: SegmentExclusionFragment[] diff --git a/packages/backend/src/api/services/ExperimentAssignmentService.ts b/packages/backend/src/api/services/ExperimentAssignmentService.ts index 3ebb824a9..25b2ef55d 100644 --- a/packages/backend/src/api/services/ExperimentAssignmentService.ts +++ b/packages/backend/src/api/services/ExperimentAssignmentService.ts @@ -2192,13 +2192,18 @@ export class ExperimentAssignmentService { }, Promise.resolve(resolveData)); } - private async getSegmentObject( + /** + * Determines which experiments need segment inclusion/exclusion resolution for these users (those + * already individually enrolled under a GROUP/INDIVIDUAL-consistency experiment are skipped). Only + * returns the experiment ids: the segment ids themselves are resolved lazily and only for whichever + * of these ids have no row in the experiment_precomputed_segment table (see + * fetchExperimentPrecomputedWithFallback / ExperimentRepository.getSegmentIdsForExperiments) rather + * than eagerly for every experiment here, since that data goes unused when a precomputed row exists. + */ + private async getExperimentIdsForSegmentResolution( experiments: Experiment[], experimentUsers: ExperimentUser[] - ): Promise { - const segmentObj: EntitySegmentResolutionInput = {}; - - // Creates a segment object for all experiments and users + ): Promise { const experimentIdsForIndividualConsistency = experiments .filter( (experiment) => @@ -2214,23 +2219,7 @@ export class ExperimentAssignmentService { }); const experimentsEnrolledIds = experimentsEnrolled.map((enrollment) => enrollment.experimentId); - // creates segment Object for all experiments - experiments.forEach((exp) => { - if (!experimentsEnrolledIds.includes(exp.id)) { - const includeIds = exp.experimentSegmentInclusion?.map((segmentInclusion) => segmentInclusion.segment.id) || []; - const excludeIds = exp.experimentSegmentExclusion?.map((segmentExclusion) => segmentExclusion.segment.id) || []; - - segmentObj[exp.id] = { - segmentIdsQueue: [...includeIds, ...excludeIds], - currentIncludedSegmentIds: includeIds, - currentExcludedSegmentIds: excludeIds, - allIncludedSegmentIds: includeIds, - allExcludedSegmentIds: excludeIds, - }; - } - }); - - return segmentObj; + return experiments.filter((exp) => !experimentsEnrolledIds.includes(exp.id)).map((exp) => exp.id); } public async resolveSegmentsForEntities( @@ -2284,10 +2273,9 @@ export class ExperimentAssignmentService { experimentUser: ExperimentUser, logger: UpgradeLogger ): Promise<[Experiment[], { experiment: Experiment; reason: string; matchedGroup: boolean }[]]> { - const segmentObj = await this.getSegmentObject(experiments, [experimentUser]); - const expIds = Object.keys(segmentObj); + const expIds = await this.getExperimentIdsForSegmentResolution(experiments, [experimentUser]); const { precomputedMap, fallbackIncludeData, fallbackExcludeData } = - await this.fetchExperimentPrecomputedWithFallback(segmentObj, logger); + await this.fetchExperimentPrecomputedWithFallback(expIds, logger); const [includeData, excludeData] = this.buildExperimentIncludeExcludeData( expIds, experimentUser, @@ -2313,12 +2301,11 @@ export class ExperimentAssignmentService { experimentUsers: ExperimentUser[], logger: UpgradeLogger ): Promise<{ userId: string; experiments: Experiment[] }[]> { - const segmentObj = await this.getSegmentObject(experiments, experimentUsers); - const expIds = Object.keys(segmentObj); + const expIds = await this.getExperimentIdsForSegmentResolution(experiments, experimentUsers); // The precomputed rows (and any on-the-fly fallback) are user-independent, so fetch once and // build each user's include/exclude view from them. const { precomputedMap, fallbackIncludeData, fallbackExcludeData } = - await this.fetchExperimentPrecomputedWithFallback(segmentObj, logger); + await this.fetchExperimentPrecomputedWithFallback(expIds, logger); const experimentIdsWithFilter: { id: string; filterMode: FILTER_MODE }[] = experiments.map( ({ id, filterMode, group }) => ({ id, filterMode, group }) ); @@ -2347,22 +2334,22 @@ export class ExperimentAssignmentService { } /** - * Read the precomputed experiment segment rows for the experiments in `segmentObj`. Mirrors the - * feature-flag read path: a read failure (e.g. the table hasn't been migrated yet) is swallowed and - * treated as "every row missing", and any experiment without a precomputed row falls back to - * on-the-fly recursive segment resolution so a missing row never silently produces a wrong decision. - * The fallback data is user-independent (full member lists), matching `resolveSegment`'s output. + * Read the precomputed experiment segment rows for `expIds`. Mirrors the feature-flag read path: a + * read failure (e.g. the table hasn't been migrated yet) is swallowed and treated as "every row + * missing", and any experiment without a precomputed row falls back to on-the-fly recursive segment + * resolution so a missing row never silently produces a wrong decision. Segment ids for the + * fallback set are fetched lazily here, scoped to just `missingExpIds`, rather than eagerly for + * every experiment in `expIds` up front. The fallback data is user-independent (full member lists), + * matching `resolveSegment`'s output. */ private async fetchExperimentPrecomputedWithFallback( - segmentObj: EntitySegmentResolutionInput, + expIds: string[], logger: UpgradeLogger ): Promise<{ precomputedMap: Map; fallbackIncludeData: EntitySegmentMembers; fallbackExcludeData: EntitySegmentMembers; }> { - const expIds = Object.keys(segmentObj); - let precomputedMap: Map; try { precomputedMap = await this.experimentPrecomputedSegmentService.getPrecomputedSets(expIds); @@ -2381,8 +2368,22 @@ export class ExperimentAssignmentService { message: `experimentLevelExclusionInclusion: ${missingExpIds.length} experiment(s) missing an experiment_precomputed_segment row; resolving on-the-fly`, details: { missingExpIds }, }); + const missingSegmentIdData = await this.experimentRepository.getSegmentIdsForExperiments(missingExpIds); const missingSegmentObj: EntitySegmentResolutionInput = {}; - missingExpIds.forEach((id) => (missingSegmentObj[id] = segmentObj[id])); + missingExpIds.forEach((id) => { + const segmentIdData = missingSegmentIdData.find((data) => data.id === id); + const includeIds = + segmentIdData?.experimentSegmentInclusion?.map((segmentInclusion) => segmentInclusion.segmentId) || []; + const excludeIds = + segmentIdData?.experimentSegmentExclusion?.map((segmentExclusion) => segmentExclusion.segmentId) || []; + missingSegmentObj[id] = { + segmentIdsQueue: [...includeIds, ...excludeIds], + currentIncludedSegmentIds: includeIds, + currentExcludedSegmentIds: excludeIds, + allIncludedSegmentIds: includeIds, + allExcludedSegmentIds: excludeIds, + }; + }); [fallbackIncludeData, fallbackExcludeData] = await this.resolveSegmentsForEntities(missingSegmentObj); } diff --git a/packages/backend/test/unit/mockdata/raw.ts b/packages/backend/test/unit/mockdata/raw.ts index a26e6ba69..0ea7f1619 100644 --- a/packages/backend/test/unit/mockdata/raw.ts +++ b/packages/backend/test/unit/mockdata/raw.ts @@ -135,6 +135,7 @@ export const simpleIndividualExperiment = { createdAt: '2023-06-01T18:44:41.153Z', updatedAt: '2023-06-01T18:44:41.153Z', versionNumber: 1, + segmentId: '89246cff-c81f-4515-91f3-c033341e45b9', segment: { createdAt: '2023-06-01T18:44:41.245Z', updatedAt: '2023-06-01T18:44:41.245Z', @@ -163,6 +164,7 @@ export const simpleIndividualExperiment = { createdAt: '2023-06-01T18:44:41.153Z', updatedAt: '2023-06-01T18:44:41.153Z', versionNumber: 1, + segmentId: 'd958bf52-7066-4594-ad8a-baf2e75324cf', segment: { createdAt: '2023-06-01T18:44:41.267Z', updatedAt: '2023-06-01T18:44:41.267Z', @@ -261,6 +263,7 @@ export const simpleGroupExperiment = { createdAt: '2023-06-02T15:03:46.960Z', updatedAt: '2023-06-02T15:03:46.960Z', versionNumber: 1, + segmentId: '89246cff-c81f-4515-91f3-c033341e45b9', segment: { createdAt: '2023-06-02T15:03:47.066Z', updatedAt: '2023-06-02T15:03:47.066Z', @@ -289,6 +292,7 @@ export const simpleGroupExperiment = { createdAt: '2023-06-02T15:03:46.960Z', updatedAt: '2023-06-02T15:03:46.960Z', versionNumber: 1, + segmentId: 'd958bf52-7066-4594-ad8a-baf2e75324cf', segment: { createdAt: '2023-06-02T15:03:47.095Z', updatedAt: '2023-06-02T15:03:47.095Z', diff --git a/packages/backend/test/unit/repositories/ExperimentRepository.test.ts b/packages/backend/test/unit/repositories/ExperimentRepository.test.ts index 7e6e8781e..1f98b85c8 100644 --- a/packages/backend/test/unit/repositories/ExperimentRepository.test.ts +++ b/packages/backend/test/unit/repositories/ExperimentRepository.test.ts @@ -270,12 +270,17 @@ describe('ExperimentRepository Testing', () => { const res = await repo.getValidExperiments('context'); - expect(repo.createQueryBuilder).toHaveBeenCalledTimes(4); - - expect(mock.leftJoinAndSelect).toHaveBeenCalledTimes(20); - expect(mock.where).toHaveBeenCalledTimes(4); - expect(mock.select).toHaveBeenCalledTimes(2); - expect(mock.getMany).toHaveBeenCalledTimes(4); + expect(repo.createQueryBuilder).toHaveBeenCalledTimes(2); + + // 4 (conditionLevel) + 6 (factorDecisionPoint) = 10. Inclusion/exclusion segment ids are no + // longer joined here: they're fetched lazily via getSegmentIdsForExperiments only when + // ExperimentAssignmentService actually needs them (an experiment_precomputed_segment cache miss). + expect(mock.leftJoinAndSelect).toHaveBeenCalledTimes(10); + expect(mock.leftJoin).not.toHaveBeenCalled(); + expect(mock.addSelect).not.toHaveBeenCalled(); + expect(mock.where).toHaveBeenCalledTimes(2); + expect(mock.select).not.toHaveBeenCalled(); + expect(mock.getMany).toHaveBeenCalledTimes(2); expect(res).toEqual(result); }); @@ -287,12 +292,12 @@ describe('ExperimentRepository Testing', () => { await repo.getValidExperiments('context'); }).rejects.toThrow(err); - expect(repo.createQueryBuilder).toHaveBeenCalledTimes(4); + expect(repo.createQueryBuilder).toHaveBeenCalledTimes(2); - expect(mock.leftJoinAndSelect).toHaveBeenCalledTimes(20); - expect(mock.where).toHaveBeenCalledTimes(4); - expect(mock.select).toHaveBeenCalledTimes(2); - expect(mock.getMany).toHaveBeenCalledTimes(4); + expect(mock.leftJoinAndSelect).toHaveBeenCalledTimes(10); + expect(mock.where).toHaveBeenCalledTimes(2); + expect(mock.select).not.toHaveBeenCalled(); + expect(mock.getMany).toHaveBeenCalledTimes(2); }); it('should get valid experiments with preview', async () => { @@ -301,12 +306,14 @@ describe('ExperimentRepository Testing', () => { const res = await repo.getValidExperimentsWithPreview('context'); - expect(repo.createQueryBuilder).toHaveBeenCalledTimes(4); + expect(repo.createQueryBuilder).toHaveBeenCalledTimes(2); - expect(mock.leftJoinAndSelect).toHaveBeenCalledTimes(20); - expect(mock.where).toHaveBeenCalledTimes(4); - expect(mock.select).toHaveBeenCalledTimes(2); - expect(mock.getMany).toHaveBeenCalledTimes(4); + expect(mock.leftJoinAndSelect).toHaveBeenCalledTimes(10); + expect(mock.leftJoin).not.toHaveBeenCalled(); + expect(mock.addSelect).not.toHaveBeenCalled(); + expect(mock.where).toHaveBeenCalledTimes(2); + expect(mock.select).not.toHaveBeenCalled(); + expect(mock.getMany).toHaveBeenCalledTimes(2); expect(res).toEqual(result); }); @@ -318,12 +325,12 @@ describe('ExperimentRepository Testing', () => { await repo.getValidExperimentsWithPreview('context'); }).rejects.toThrow(err); - expect(repo.createQueryBuilder).toHaveBeenCalledTimes(4); + expect(repo.createQueryBuilder).toHaveBeenCalledTimes(2); - expect(mock.leftJoinAndSelect).toHaveBeenCalledTimes(20); - expect(mock.where).toHaveBeenCalledTimes(4); - expect(mock.select).toHaveBeenCalledTimes(2); - expect(mock.getMany).toHaveBeenCalledTimes(4); + expect(mock.leftJoinAndSelect).toHaveBeenCalledTimes(10); + expect(mock.where).toHaveBeenCalledTimes(2); + expect(mock.select).not.toHaveBeenCalled(); + expect(mock.getMany).toHaveBeenCalledTimes(2); }); it('should update experiment state', async () => { @@ -500,33 +507,30 @@ describe('ExperimentRepository Testing', () => { }); describe('getValidExperimentsForContextAndDecisionPoint', () => { - it('should build four queries and add a leftJoin on partitions for the condition and segment queries', async () => { + it('should build two queries and add a leftJoin on partitions for the condition query', async () => { const result = [experiment]; mock.getMany.mockResolvedValue(result); const res = await repo.getValidExperimentsForContextAndDecisionPoint('context', 'site1', 'target1'); - expect(repo.createQueryBuilder).toHaveBeenCalledTimes(4); - // 4 (conditionLevel) + 6 (factorDecisionPoint) + 5 (inclusion) + 5 (exclusion) = 20 - expect(mock.leftJoinAndSelect).toHaveBeenCalledTimes(20); - // conditionLevelPayloadQuery and both segment queries add a non-selecting leftJoin for partition filtering - expect(mock.leftJoin).toHaveBeenCalledTimes(3); + expect(repo.createQueryBuilder).toHaveBeenCalledTimes(2); + // 4 (conditionLevel) + 6 (factorDecisionPoint) = 10. Inclusion/exclusion segment ids are no + // longer joined here; see getSegmentIdsForExperiments. + expect(mock.leftJoinAndSelect).toHaveBeenCalledTimes(10); + // conditionLevelPayloadQuery adds a non-selecting leftJoin for partition filtering. + expect(mock.leftJoin).toHaveBeenCalledTimes(1); expect(mock.leftJoin).toHaveBeenCalledWith('experiment.partitions', 'partitions'); - // Both segment queries call .select('experiment.id') - expect(mock.select).toHaveBeenCalledTimes(2); - expect(mock.where).toHaveBeenCalledTimes(4); - expect(mock.getMany).toHaveBeenCalledTimes(4); + expect(mock.addSelect).not.toHaveBeenCalled(); + expect(mock.select).not.toHaveBeenCalled(); + expect(mock.where).toHaveBeenCalledTimes(2); + expect(mock.getMany).toHaveBeenCalledTimes(2); expect(res).toEqual(result); }); it('should return empty array when no experiments match the site/target', async () => { - // conditionLevel and segment find experiments, but factorDecisionPoint finds none at this site/target - mock.getMany - .mockResolvedValueOnce([experiment]) - .mockResolvedValueOnce([]) - .mockResolvedValueOnce([experiment]) - .mockResolvedValueOnce([experiment]); + // conditionLevel finds an experiment, but factorDecisionPoint finds none at this site/target + mock.getMany.mockResolvedValueOnce([experiment]).mockResolvedValueOnce([]); const res = await repo.getValidExperimentsForContextAndDecisionPoint('context', 'site1', 'target1'); @@ -539,12 +543,8 @@ describe('ExperimentRepository Testing', () => { const expB = new Experiment(); expB.id = 'exp-b'; - // conditionLevel and segment over-fetch; only expA matches the site/target in factorDecisionPoint - mock.getMany - .mockResolvedValueOnce([expA, expB]) - .mockResolvedValueOnce([expA]) - .mockResolvedValueOnce([expA, expB]) - .mockResolvedValueOnce([expA, expB]); + // conditionLevel over-fetches; only expA matches the site/target in factorDecisionPoint + mock.getMany.mockResolvedValueOnce([expA, expB]).mockResolvedValueOnce([expA]); const res = await repo.getValidExperimentsForContextAndDecisionPoint('context', 'site1', 'target1'); @@ -552,17 +552,11 @@ describe('ExperimentRepository Testing', () => { expect(res[0].id).toBe('exp-a'); }); - it('should merge condition, partition, and segment data onto each result experiment', async () => { + it('should merge condition and partition data onto each result experiment', async () => { const condData = { id: 'exp-a', conditions: ['cond1'] } as any; const factorData = { id: 'exp-a', partitions: ['part1'] } as any; - const inclusionData = { id: 'exp-a', experimentSegmentInclusion: ['seg1'] } as any; - const exclusionData = { id: 'exp-a', experimentSegmentExclusion: ['seg2'] } as any; - mock.getMany - .mockResolvedValueOnce([condData]) - .mockResolvedValueOnce([factorData]) - .mockResolvedValueOnce([inclusionData]) - .mockResolvedValueOnce([exclusionData]); + mock.getMany.mockResolvedValueOnce([condData]).mockResolvedValueOnce([factorData]); const [result] = await repo.getValidExperimentsForContextAndDecisionPoint('context', 'site1', 'target1'); @@ -570,44 +564,19 @@ describe('ExperimentRepository Testing', () => { id: 'exp-a', conditions: ['cond1'], partitions: ['part1'], - experimentSegmentInclusion: ['seg1'], - experimentSegmentExclusion: ['seg2'], }); }); - it('should return experiment without segment data when segment query returns no match', async () => { - const condData = { id: 'exp-a', conditions: ['cond1'] } as any; - const factorData = { id: 'exp-a', partitions: ['part1'] } as any; - - mock.getMany - .mockResolvedValueOnce([condData]) - .mockResolvedValueOnce([factorData]) - .mockResolvedValueOnce([]) - .mockResolvedValueOnce([]); - - const [result] = await repo.getValidExperimentsForContextAndDecisionPoint('context', 'site1', 'target1'); - - expect(result).toMatchObject({ id: 'exp-a', conditions: ['cond1'], partitions: ['part1'] }); - }); - it('should return factorDecisionPoint data even when conditionLevel query returns no match', async () => { const factorData = { id: 'exp-a', partitions: ['part1'] } as any; - const inclusionData = { id: 'exp-a', experimentSegmentInclusion: ['seg1'] } as any; - const exclusionData = { id: 'exp-a', experimentSegmentExclusion: ['seg2'] } as any; - mock.getMany - .mockResolvedValueOnce([]) - .mockResolvedValueOnce([factorData]) - .mockResolvedValueOnce([inclusionData]) - .mockResolvedValueOnce([exclusionData]); + mock.getMany.mockResolvedValueOnce([]).mockResolvedValueOnce([factorData]); const [result] = await repo.getValidExperimentsForContextAndDecisionPoint('context', 'site1', 'target1'); expect(result).toMatchObject({ id: 'exp-a', partitions: ['part1'], - experimentSegmentInclusion: ['seg1'], - experimentSegmentExclusion: ['seg2'], }); }); @@ -616,7 +585,42 @@ describe('ExperimentRepository Testing', () => { await expect(repo.getValidExperimentsForContextAndDecisionPoint('context', 'site1', 'target1')).rejects.toThrow(); - expect(repo.createQueryBuilder).toHaveBeenCalledTimes(4); + expect(repo.createQueryBuilder).toHaveBeenCalledTimes(2); + }); + }); + + describe('getSegmentIdsForExperiments', () => { + it('should return an empty array without querying when given no experiment ids', async () => { + const res = await repo.getSegmentIdsForExperiments([]); + + expect(repo.createQueryBuilder).not.toHaveBeenCalled(); + expect(res).toEqual([]); + }); + + it('should fetch and merge inclusion/exclusion segment ids scoped to the given experiment ids', async () => { + const inclusionData = { id: 'exp-a', experimentSegmentInclusion: ['inclusion'] } as any; + const exclusionData = { id: 'exp-a', experimentSegmentExclusion: ['exclusion'] } as any; + + mock.getMany.mockResolvedValueOnce([inclusionData]).mockResolvedValueOnce([exclusionData]); + + const res = await repo.getSegmentIdsForExperiments(['exp-a']); + + expect(repo.createQueryBuilder).toHaveBeenCalledTimes(2); + expect(mock.leftJoinAndSelect).toHaveBeenCalledTimes(2); + expect(mock.select).toHaveBeenCalledTimes(2); + expect(mock.where).toHaveBeenCalledTimes(2); + expect(mock.where).toHaveBeenCalledWith('experiment.id IN (:...experimentIds)', { experimentIds: ['exp-a'] }); + expect(mock.getMany).toHaveBeenCalledTimes(2); + + expect(res).toMatchObject([ + { id: 'exp-a', experimentSegmentInclusion: ['inclusion'], experimentSegmentExclusion: ['exclusion'] }, + ]); + }); + + it('should throw an error when a sub-query fails', async () => { + mock.getMany.mockRejectedValue(err); + + await expect(repo.getSegmentIdsForExperiments(['exp-a'])).rejects.toThrow(err); }); }); diff --git a/packages/backend/test/unit/services/ExperimentAssignmentService.test.ts b/packages/backend/test/unit/services/ExperimentAssignmentService.test.ts index 25f6543a3..79eb590c4 100644 --- a/packages/backend/test/unit/services/ExperimentAssignmentService.test.ts +++ b/packages/backend/test/unit/services/ExperimentAssignmentService.test.ts @@ -77,6 +77,32 @@ describe('Experiment Assignment Service Test', () => { // Default to "no precomputed rows" so the assignment read path exercises the on-the-fly fallback // (recursive segment resolution) these tests were written against. experimentPrecomputedSegmentServiceMock.getPrecomputedSets.resolves(new Map()); + // ExperimentAssignmentService now fetches inclusion/exclusion segment ids lazily, scoped to just the + // experiment ids that miss the precomputed cache (see ExperimentRepository.getSegmentIdsForExperiments), + // instead of reading them off the experiments returned by getValidExperiments*. Mirror that lookup + // here against the known mock fixtures so existing tests don't each need their own stub for it. + const allMockExperimentsForSegmentLookup: any[] = [ + simpleIndividualAssignmentExperiment, + simpleIndividualAssignmentExperiment2, + simpleGroupAssignmentExperiment, + factorialGroupAssignmentExperiment, + factorialIndividualAssignmentExperiment, + simpleDPExperiment, + simpleWithinSubjectOrderedRoundRobinExperiment, + simpleWithinSubjectRandomRoundRobinExperiment, + withinSubjectDPExperiment, + factorialGroupExperiment, + factorialIndividualExperiment, + ]; + experimentRepositoryMock.getSegmentIdsForExperiments.callsFake(async (experimentIds: string[]) => + allMockExperimentsForSegmentLookup + .filter((exp) => experimentIds.includes(exp.id)) + .map((exp) => ({ + id: exp.id, + experimentSegmentInclusion: exp.experimentSegmentInclusion, + experimentSegmentExclusion: exp.experimentSegmentExclusion, + })) + ); experimentServiceMock.formattingConditionPayload.restore(); experimentServiceMock.formattingPayload.restore(); From b94b1d6d2c778c7575a7b5a8fbcd4fdaf527c688 Mon Sep 17 00:00:00 2001 From: Benjamin Blanchard Date: Fri, 28 Aug 2026 14:08:58 -0400 Subject: [PATCH 2/2] build a lookup map for the fallback case --- .../backend/src/api/services/ExperimentAssignmentService.ts | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/packages/backend/src/api/services/ExperimentAssignmentService.ts b/packages/backend/src/api/services/ExperimentAssignmentService.ts index 25b2ef55d..4c5fae8ae 100644 --- a/packages/backend/src/api/services/ExperimentAssignmentService.ts +++ b/packages/backend/src/api/services/ExperimentAssignmentService.ts @@ -2369,9 +2369,10 @@ export class ExperimentAssignmentService { details: { missingExpIds }, }); const missingSegmentIdData = await this.experimentRepository.getSegmentIdsForExperiments(missingExpIds); + const missingSegmentIdDataMap = new Map(missingSegmentIdData.map((data) => [data.id, data])); const missingSegmentObj: EntitySegmentResolutionInput = {}; missingExpIds.forEach((id) => { - const segmentIdData = missingSegmentIdData.find((data) => data.id === id); + const segmentIdData = missingSegmentIdDataMap.get(id); const includeIds = segmentIdData?.experimentSegmentInclusion?.map((segmentInclusion) => segmentInclusion.segmentId) || []; const excludeIds =