remove select queries in experiment update transaction

Author: bcb37

packages/backend/src/api/models/ConditionPayload.ts

@@ -22,6 +22,12 @@ export class ConditionPayload extends BaseModel {
   @ManyToOne(() => ExperimentCondition, (condition) => condition.conditionPayloads, { onDelete: 'CASCADE' })
   public parentCondition: ExperimentCondition;
 
+  @Column({ name: 'parentConditionId', nullable: true })
+  parentConditionId?: string;
+
   @ManyToOne(() => DecisionPoint, (decisionPoint) => decisionPoint.conditionPayloads, { onDelete: 'CASCADE' })
   public decisionPoint: DecisionPoint;
+
+  @Column({ name: 'decisionPointId', nullable: true })
+  decisionPointId?: string;
 }

packages/backend/src/api/services/ExperimentService.ts

@@ -83,7 +83,6 @@ import {
   ValidatedExperimentError,
   ConditionPayloadValidator,
 } from '../DTO/ExperimentDTO';
-import { ConditionPayloadDTO } from '../DTO/ConditionPayloadDTO';
 import { FactorDTO } from '../DTO/FactorDTO';
 import { LevelDTO } from '../DTO/LevelDTO';
 import { CacheService } from './CacheService';
@@ -760,10 +759,10 @@ export class ExperimentService {
       this.experimentSchedulerService.updateExperimentSchedules(experiment as any, logger);
     }
 
+    let uniqueIdentifiers = await this.getAllUniqueIdentifiers(logger);
     return entityManager
       .transaction(async (transactionalEntityManager) => {
         experiment.context = experiment.context.map((context) => context.toLocaleLowerCase());
-        let uniqueIdentifiers = await this.getAllUniqueIdentifiers(logger);
         if (experiment.conditions.length) {
           const response = this.setConditionOrPartitionIdentifiers(experiment.conditions, uniqueIdentifiers);
           experiment.conditions = response[0];
@@ -861,16 +860,16 @@ export class ExperimentService {
           });
         });
 
-        const conditionPayloadDocToSave: Array<Partial<Omit<ConditionPayload, 'parentCondition' | 'decisionPoint'>>> =
+        const conditionPayloadDocToSave: Array<Partial<ConditionPayload>> =
           (newPayloads &&
             newPayloads.length > 0 &&
             newPayloads.map((conditionPayload) => {
               const conditionPayloadToReturn = {
                 id: conditionPayload.id,
                 payloadType: conditionPayload.payload.type,
                 payloadValue: conditionPayload.payload.value,
-                parentCondition: conditionPayload.parentCondition,
-                decisionPoint: conditionPayload.decisionPoint,
+                parentConditionId: conditionPayload.parentCondition,
+                decisionPointId: conditionPayload.decisionPoint,
               };
               return conditionPayloadToReturn;
             })) ||
@@ -993,7 +992,6 @@ export class ExperimentService {
         let conditionDocs: ExperimentCondition[];
         let decisionPointDocs: DecisionPoint[];
         let queryDocs: Query[];
-        let conditionPayloadDocs: ConditionPayloadDTO[];
         try {
           [conditionDocs, decisionPointDocs, queryDocs] = await Promise.all([
             Promise.all(
@@ -1026,16 +1024,14 @@ export class ExperimentService {
         }
 
         try {
-          [conditionPayloadDocs] = await Promise.all([
-            Promise.all(
-              conditionPayloadDocToSave.map(async (conditionPayload) => {
-                return this.conditionPayloadRepository.upsertConditionPayload(
-                  conditionPayload,
-                  transactionalEntityManager
-                );
-              })
-            ) as any,
-          ]);
+          await Promise.all(
+            conditionPayloadDocToSave.map(async (conditionPayload) => {
+              return this.conditionPayloadRepository.upsertConditionPayload(
+                conditionPayload,
+                transactionalEntityManager
+              );
+            })
+          );
         } catch (err) {
           const error = err as Error;
           error.message = `Error in creating conditionPayloads "updateExperimentInDB"`;
@@ -1051,11 +1047,15 @@ export class ExperimentService {
           return { ...decisionPointDoc, experiment: decisionPointDoc.experiment };
         });
 
-        const conditionPayloadDocToReturn = await transactionalEntityManager.getRepository(ConditionPayload).find({
-          relations: ['parentCondition', 'decisionPoint'],
-          where: { id: In(conditionPayloadDocs.map((conditionPayload) => conditionPayload.id)) },
+        const conditionPayloadDocToReturn = conditionPayloadDocToSave.map((conditionPayload) => {
+          const parentCondition = conditionDocs.find((c) => c.id === conditionPayload.parentConditionId);
+          const decisionPoint = decisionPointDocs.find((dp) => dp.id === conditionPayload.decisionPointId);
+          return {
+            ...conditionPayload,
+            parentCondition: parentCondition,
+            decisionPoint: decisionPoint,
+          };
         });
-
         // sort the payloads by decision point order and condition order
         conditionPayloadDocToReturn.sort((a, b) => {
           if (a.decisionPoint.order === b.decisionPoint.order) {

packages/backend/test/unit/services/ExperimentService.test.ts

@@ -83,6 +83,13 @@ describe('ExperimentService Testing', () => {
     order: 1,
   };
 
+  const mockCondition3: Partial<ExperimentCondition> = {
+    id: 'condition-3',
+    conditionCode: 'new-treatment',
+    assignmentWeight: 40,
+    order: 2,
+  };
+
   const createMockDecisionPoint1 = () => ({
     id: 'partition-1',
     site: 'test-site',
@@ -115,7 +122,9 @@ describe('ExperimentService Testing', () => {
       payloadType: PAYLOAD_TYPE.STRING,
       payloadValue: 'control',
       parentCondition: mockCondition1 as ExperimentCondition,
+      parentConditionId: mockCondition1.id,
       decisionPoint: mockDecisionPoint1 as DecisionPoint,
+      decisionPointId: mockDecisionPoint1.id,
     };
 
     mockExperiment = {
@@ -170,14 +179,14 @@ describe('ExperimentService Testing', () => {
       conditionPayloads: [
         {
           id: 'payload-1',
-          parentCondition: 'condition-1',
-          decisionPoint: 'partition-1',
+          parentCondition: mockCondition1,
+          decisionPoint: mockDecisionPoint1,
           payload: { type: PAYLOAD_TYPE.STRING, value: 'control' },
         } as any,
         {
           id: 'payload-2',
-          parentCondition: 'condition-3',
-          decisionPoint: 'partition-1',
+          parentCondition: mockCondition3,
+          decisionPoint: mockDecisionPoint1,
           payload: { type: PAYLOAD_TYPE.STRING, value: 'new-treatment' },
         } as any,
       ],
@@ -267,7 +276,7 @@ describe('ExperimentService Testing', () => {
           useValue: {
             find: jest.fn().mockResolvedValue([mockCondition1]),
             save: jest.fn().mockResolvedValue(mockCondition1),
-            upsertExperimentCondition: jest.fn().mockResolvedValue(mockCondition1),
+            upsertExperimentCondition: jest.fn().mockImplementation((value) => Promise.resolve(value)),
             deleteCondition: jest.fn().mockResolvedValue({ affected: 1 }),
             getAllUniqueIdentifier: jest.fn().mockResolvedValue(['C1', 'C2']),
             insertConditions: jest.fn().mockResolvedValue([mockCondition1]),
@@ -279,7 +288,7 @@ describe('ExperimentService Testing', () => {
             find: jest.fn().mockResolvedValue([mockDecisionPoint1]),
             save: jest.fn().mockResolvedValue(mockDecisionPoint1),
             findOne: jest.fn().mockResolvedValue(null),
-            upsertDecisionPoint: jest.fn().mockResolvedValue(mockDecisionPoint1),
+            upsertDecisionPoint: jest.fn().mockImplementation((value) => Promise.resolve(value)),
             deleteDecisionPoint: jest.fn().mockResolvedValue({ affected: 1 }),
             deleteByIds: jest.fn().mockResolvedValue({ affected: 1 }),
             getAllUniqueIdentifier: jest.fn().mockResolvedValue(['C1', 'C2']),
@@ -693,8 +702,8 @@ describe('ExperimentService Testing', () => {
       // Should create payloads for partition-2 with both conditions
       expect(conditionPayloadRepo.upsertConditionPayload).toHaveBeenCalledWith(
         expect.objectContaining({
-          parentCondition: 'condition-1',
-          decisionPoint: 'partition-2',
+          parentConditionId: 'condition-1',
+          decisionPointId: 'partition-2',
           payloadValue: 'control',
         }),
         entityManager