Add batch assign endpoint and functionality

backend/packages/Upgrade/src/api/controllers/BatchAssignController.ts

@@ -0,0 +1,146 @@
+import { JsonController, Post, Authorized, Req, Body } from 'routing-controllers';
+import { AppRequest } from '../../types';
+import { BatchAssignValidator } from '../controllers/validators/BatchAssignValidator';
+import { ExperimentAssignmentService } from '../services/ExperimentAssignmentService';
+import { ExperimentUserService } from '../services/ExperimentUserService';
+import { IExperimentAssignmentv5 } from 'upgrade_types';
+
+@Authorized()
+@JsonController('/batch-assign')
+export class BatchAssignController {
+  constructor(
+    public experimentAssignmentService: ExperimentAssignmentService,
+    public experimentUserService: ExperimentUserService
+  ) {}
+  /**
+   * @swagger
+   * /batch-assign:
+   *    post:
+   *       description: Get Batch Assignments for students
+   *       consumes:
+   *         - application/json
+   *       parameters:
+   *         - in: body
+   *           name: body
+   *           description: Batch Assign Request Body
+   *           required: true
+   *           schema:
+   *             type: object
+   *             properties:
+   *               userIds:
+   *                 type: array
+   *                 items:
+   *                   type: string
+   *               context:
+   *                 type: string
+   *               site:
+   *                 type: string
+   *               target:
+   *                 type: string
+   *       produces:
+   *         - application/json
+   *       tags:
+   *         - BatchAssign
+   *       responses:
+   *          '200':
+   *            description: Get Batch Assignments for students
+   *            schema:
+   *              type: object
+   *              additionalProperties:
+   *                type: object
+   *                required:
+   *                  - site
+   *                  - target
+   *                  - condition
+   *                properties:
+   *                  site:
+   *                    type: string
+   *                    minLength: 1
+   *                  target:
+   *                    type: string
+   *                    minLength: 1
+   *                  experimentType:
+   *                    type: string
+   *                    enum: [Simple, Factorial]
+   *                  assignedCondition:
+   *                    type: array
+   *                    items:
+   *                      type: object
+   *                      properties:
+   *                         conditionCode:
+   *                          type: string
+   *                          minLength: 1
+   *                         payload:
+   *                          type: object
+   *                          properties:
+   *                            type:
+   *                              type: string
+   *                            value:
+   *                              type: string
+   *                         id:
+   *                          type: string
+   *                         experimentId:
+   *                          type: string
+   *              example:
+   *                {
+   *                  "user1": {
+   *                    "site": "site1",
+   *                    "target": "target1",
+   *                    "condition": "condition1",
+   *                    "experimentType": "Simple",
+   *                    "assignedCondition": [
+   *                      {
+   *                        "conditionCode": "condition1",
+   *                        "payload": {
+   *                          "type": "type1",
+   *                          "value": "value1"
+   *                        },
+   *                        "id": "conditionId1",
+   *                        "experimentId": "experimentId1"
+   *                      }
+   *                    ]
+   *                  },
+   *                  "user2": null
+   *                }
+   *          '400':
+   *            description: BadRequestError - InvalidParameterValue
+   *          '401':
+   *            description: AuthorizationRequiredError
+   *          '500':
+   *            description: Internal Server Error
+   */
+  @Post('/')
+  public async getBatchAssignments(
+    @Body({ validate: true }) requestBody: BatchAssignValidator,
+    @Req() request: AppRequest
+  ): Promise<Record<string, IExperimentAssignmentv5 | null>> {
+    request.logger.info({ message: 'Request received for batch assignments' });
+    const { context, site, target, userIds } = requestBody;
+    request.logger.info({
+      message: `Context: ${context}, Site: ${site}, Target: ${target}, User IDs: ${userIds.join(', ')}`,
+    });
+
+    const userDocs = await this.experimentUserService.getUserDocs(userIds, request.logger);
+
+    // Initialize response for all userIds
+    const initializedResponse = userIds.reduce((acc, userId) => {
+      if (userDocs.find((doc) => doc.id === userId)) {
+        acc[userId] = null; // initialize with null if userDoc exists
+      } else {
+        acc[userId] = 'NOT_FOUND'; // 'NOT_FOUND' if userDoc does not exist
+      }
+      return acc;
+    }, {});
+    request.logger.info({ message: `User Docs: ${JSON.stringify(userDocs)}` });
+
+    const assignments = await this.experimentAssignmentService.getBatchExperimentConditions(
+      userDocs,
+      context,
+      site,
+      target,
+      request.logger
+    );
+
+    return { ...initializedResponse, ...assignments };
+  }
+}

backend/packages/Upgrade/src/api/controllers/ExperimentClientController.v5.ts

@@ -15,15 +15,7 @@ import { ExperimentAssignmentValidator } from './validators/ExperimentAssignment
 import { ExperimentUser } from '../models/ExperimentUser';
 import { ExperimentUserService } from '../services/ExperimentUserService';
 import { UpdateWorkingGroupValidator } from './validators/UpdateWorkingGroupValidator';
-import {
-  IExperimentAssignmentv5,
-  SERVER_ERROR,
-  IGroupMembership,
-  IUserAliases,
-  IWorkingGroup,
-  PAYLOAD_TYPE,
-  IPayload,
-} from 'upgrade_types';
+import { IExperimentAssignmentv5, SERVER_ERROR, IGroupMembership, IUserAliases, IWorkingGroup } from 'upgrade_types';
 import { FeatureFlagService } from '../services/FeatureFlagService';
 import { ClientLibMiddleware } from '../middlewares/ClientLibMiddleware';
 import { LogValidator } from './validators/LogValidator';
@@ -588,38 +580,7 @@ export class ExperimentClientController {
       request.logger
     );
 
-    return assignedData.map(({ assignedFactor, assignedCondition, ...rest }) => {
-      const finalFactorData = assignedFactor?.map((factor) => {
-        const updatedAssignedFactor: Record<string, { level: string; payload: IPayload }> = {};
-        Object.keys(factor).forEach((key) => {
-          updatedAssignedFactor[key] = {
-            level: factor[key].level,
-            payload:
-              factor[key].payload && factor[key].payload.value
-                ? { type: PAYLOAD_TYPE.STRING, value: factor[key].payload.value }
-                : null,
-          };
-        });
-        return updatedAssignedFactor;
-      });
-
-      const finalConditionData = assignedCondition.map((condition) => {
-        return {
-          id: condition.id,
-          conditionCode: condition.conditionCode,
-          payload:
-            condition.payload && condition.payload.value
-              ? { type: condition.payload.type, value: condition.payload.value }
-              : null,
-          experimentId: condition.experimentId,
-        };
-      });
-      return {
-        ...rest,
-        assignedCondition: finalConditionData,
-        assignedFactor: assignedFactor ? finalFactorData : undefined,
-      };
-    });
+    return this.experimentAssignmentService.formatAssignments(assignedData);
   }
 
   /**

backend/packages/Upgrade/src/api/controllers/ExperimentClientController.v6.ts

@@ -16,14 +16,7 @@ import { FeatureFlagRequestValidator } from './validators/FeatureFlagRequestVali
 import { ExperimentUser } from '../models/ExperimentUser';
 import { ExperimentUserService } from '../services/ExperimentUserService';
 import { UpdateWorkingGroupValidatorv6 } from './validators/UpdateWorkingGroupValidator';
-import {
-  IExperimentAssignmentv5,
-  IGroupMembership,
-  IUserAliases,
-  IWorkingGroup,
-  PAYLOAD_TYPE,
-  IPayload,
-} from 'upgrade_types';
+import { IExperimentAssignmentv5, IGroupMembership, IUserAliases, IWorkingGroup } from 'upgrade_types';
 import { FeatureFlagService } from '../services/FeatureFlagService';
 import { ClientLibMiddleware } from '../middlewares/ClientLibMiddleware';
 import { LogValidatorv6 } from './validators/LogValidator';
@@ -549,38 +542,7 @@ export class ExperimentClientController {
       request.logger
     );
 
-    return assignedData.map(({ assignedFactor, assignedCondition, ...rest }) => {
-      const finalFactorData = assignedFactor?.map((factor) => {
-        const updatedAssignedFactor: Record<string, { level: string; payload: IPayload }> = {};
-        Object.keys(factor).forEach((key) => {
-          updatedAssignedFactor[key] = {
-            level: factor[key].level,
-            payload:
-              factor[key].payload && factor[key].payload.value
-                ? { type: PAYLOAD_TYPE.STRING, value: factor[key].payload.value }
-                : null,
-          };
-        });
-        return updatedAssignedFactor;
-      });
-
-      const finalConditionData = assignedCondition.map((condition) => {
-        return {
-          id: condition.id,
-          conditionCode: condition.conditionCode,
-          payload:
-            condition.payload && condition.payload.value
-              ? { type: condition.payload.type, value: condition.payload.value }
-              : null,
-          experimentId: condition.experimentId,
-        };
-      });
-      return {
-        ...rest,
-        assignedCondition: finalConditionData,
-        assignedFactor: assignedFactor ? finalFactorData : undefined,
-      };
-    });
+    return this.experimentAssignmentService.formatAssignments(assignedData);
   }
 
   /**

backend/packages/Upgrade/src/api/controllers/validators/BatchAssignValidator.ts

@@ -0,0 +1,20 @@
+import { IsNotEmpty } from 'class-validator';
+import { IsString } from 'class-validator';
+
+export class BatchAssignValidator {
+  @IsNotEmpty()
+  @IsString()
+  public context: string;
+
+  @IsNotEmpty()
+  @IsString()
+  public site: string;
+
+  @IsNotEmpty()
+  @IsString()
+  public target: string;
+
+  @IsNotEmpty()
+  @IsString({ each: true })
+  public userIds: string[];
+}

backend/packages/Upgrade/src/api/repositories/ExperimentRepository.ts

@@ -215,6 +215,25 @@ export class ExperimentRepository extends Repository<Experiment> {
     return mergedData;
   }
 
+  public async getValidExperimentsForContextAndDecisionPoint(
+    context: string,
+    site: string,
+    target: string
+  ): Promise<Experiment[]> {
+    const whereExperimentsClause =
+      '(experiment.state = :enrolling OR experiment.state = :enrollmentComplete) AND NOT (experiment.state = :enrollmentComplete AND experiment.postExperimentRule = :assign AND experiment.revertTo IS NULL) AND :context ILIKE ANY (ARRAY[experiment.context]) AND partitions.site = :site AND partitions.target = :target';
+    const whereClauseParams = {
+      enrolling: 'enrolling',
+      enrollmentComplete: 'enrollmentComplete',
+      assign: 'assign',
+      context,
+      site,
+      target,
+    };
+    const experiment = await this.createBaseQueryBuilder().where(whereExperimentsClause, whereClauseParams).getMany();
+    return experiment;
+  }
+
   public async getValidExperimentsWithPreview(context: string): Promise<Experiment[]> {
     const whereExperimentsClause =
       '(experiment.state = :enrolling OR experiment.state = :enrollmentComplete OR experiment.state = :preview) AND NOT (experiment.state = :enrollmentComplete AND experiment.postExperimentRule = :assign AND experiment.revertTo IS NULL) AND :context ILIKE ANY (ARRAY[experiment.context])';

backend/packages/Upgrade/src/api/repositories/IndividualEnrollmentRepository.ts

@@ -11,6 +11,13 @@ export class IndividualEnrollmentRepository extends Repository<IndividualEnrollm
     });
   }
 
+  public findEnrollmentsForUsers(userIds: string[], experimentIds: string[]): Promise<IndividualEnrollment[]> {
+    return this.find({
+      where: { experimentId: In(experimentIds), userId: In(userIds) },
+      select: ['id', 'userId', 'experimentId', 'enrollmentCode', 'conditionId'],
+    });
+  }
+
   public async deleteEnrollmentsOfUserInExperiments(userId: string, experimentIds: string[]): Promise<DeleteResult> {
     return this.delete({
       user: { id: userId },

backend/packages/Upgrade/src/api/repositories/IndividualExclusionRepository.ts

@@ -52,6 +52,24 @@ export class IndividualExclusionRepository extends Repository<IndividualExclusio
       });
   }
 
+  public async findExcludedForUsers(userIds: string[], experimentIds: string[]): Promise<IndividualExclusion[]> {
+    const primaryKeys = experimentIds.flatMap((experimentId) => {
+      return userIds.map((userId) => `${experimentId}_${userId}`);
+    });
+    return await this.createQueryBuilder('individualExclusion')
+      .whereInIds(primaryKeys)
+      .getMany()
+      .catch((errorMsg: any) => {
+        const errorMsgString = repositoryError(
+          this.constructor.name,
+          'findExcludedForUsers',
+          { userIds, experimentIds },
+          errorMsg
+        );
+        throw errorMsgString;
+      });
+  }
+
   public async findExcludedByExperimentId(experimentId: string): Promise<IndividualExclusion[]> {
     return await this.createQueryBuilder('individualExclusion')
       .leftJoinAndSelect('individualExclusion.experiment', 'experiment')