feat(abba): manually save user assignments

Author: frodi-karlsson

packages/abba/src/abba.test.ts

@@ -6,7 +6,7 @@ import { bucketDao } from './dao/bucket.dao.js'
 import { experimentDao } from './dao/experiment.dao.js'
 import { userAssignmentDao } from './dao/userAssignment.dao.js'
 import { mockBucket, mockExperiment, mockUserAssignment, mockUserId1 } from './test/mocks.js'
-import { AssignmentStatus } from './types.js'
+import { AssignmentStatus, SegmentationRuleOperator } from './types.js'
 import type { DecoratedUserAssignment } from './types.js'
 
 const db = new InMemoryDB({ logger: undefined })
@@ -530,3 +530,120 @@ describe('softDeleteExperiment', () => {
     expect(updatedExperiment2.exclusions).toEqual([])
   })
 })
+
+describe('saveManualUserAssignments', () => {
+  test('should create new assignments and overwrite existing ones in a single batch', async () => {
+    const experiment1 = await experimentsDAO.save(mockExperiment({ key: 'EXP_1' }))
+    const control1 = await bucketsDAO.save(mockBucket(experiment1.id, 'control', 50))
+    const test1 = await bucketsDAO.save(mockBucket(experiment1.id, 'test', 50))
+    const experiment2 = await experimentsDAO.save(mockExperiment({ key: 'EXP_2' }))
+    const test2 = await bucketsDAO.save(mockBucket(experiment2.id, 'test', 100))
+    const existing = await userAssignmentsDAO.save({
+      userId: 'userA',
+      experimentId: experiment1.id,
+      bucketId: control1.id,
+    })
+
+    const results = await abba.saveManualUserAssignments([
+      { userId: 'userA', experimentKey: 'EXP_1', bucketKey: 'test' },
+      { userId: 'userB', experimentKey: 'EXP_1', bucketKey: 'control' },
+      { userId: 'userA', experimentKey: 'EXP_2', bucketKey: 'test' },
+    ])
+
+    expect(results).toHaveLength(3)
+    expect(results[0]).toMatchObject({ id: existing.id, userId: 'userA', bucketId: test1.id })
+    expect(results[1]).toMatchObject({ userId: 'userB', bucketId: control1.id })
+    expect(results[2]).toMatchObject({ userId: 'userA', bucketId: test2.id })
+    const all = await userAssignmentsDAO.getBy('userId', 'userA')
+    expect(all).toHaveLength(2)
+  })
+
+  test('should return an empty array for an empty input', async () => {
+    const result = await abba.saveManualUserAssignments([])
+    expect(result).toEqual([])
+  })
+
+  test('should apply last-wins dedupe for duplicate (userId, experimentKey) pairs', async () => {
+    const experiment = await experimentsDAO.save(mockExperiment())
+    await bucketsDAO.save(mockBucket(experiment.id, 'control', 50))
+    const testBucket = await bucketsDAO.save(mockBucket(experiment.id, 'test', 50))
+
+    const result = await abba.saveManualUserAssignments([
+      { userId: mockUserId1, experimentKey: experiment.key, bucketKey: 'control' },
+      { userId: mockUserId1, experimentKey: experiment.key, bucketKey: 'test' },
+    ])
+
+    expect(result).toHaveLength(1)
+    expect(result[0]!.bucketId).toBe(testBucket.id)
+    const persisted = await userAssignmentsDAO.getBy('userId', mockUserId1)
+    expect(persisted).toHaveLength(1)
+    expect(persisted[0]!.bucketId).toBe(testBucket.id)
+  })
+
+  test('should throw when any row references an unknown experiment', async () => {
+    const experiment = await experimentsDAO.save(mockExperiment())
+    await bucketsDAO.save(mockBucket(experiment.id, 'test', 100))
+
+    await expect(
+      abba.saveManualUserAssignments([
+        { userId: 'userA', experimentKey: experiment.key, bucketKey: 'test' },
+        { userId: 'userB', experimentKey: 'NOPE', bucketKey: 'test' },
+      ]),
+    ).rejects.toThrow('Experiment does not exist: NOPE')
+    const written = await userAssignmentsDAO.getBy('userId', 'userA')
+    expect(written).toEqual([])
+  })
+
+  test('should throw when the experiment is soft-deleted', async () => {
+    const experiment = await experimentsDAO.save(mockExperiment({ deleted: true }))
+    await bucketsDAO.save(mockBucket(experiment.id, 'test', 100))
+
+    await expect(
+      abba.saveManualUserAssignments([
+        { userId: mockUserId1, experimentKey: experiment.key, bucketKey: 'test' },
+      ]),
+    ).rejects.toThrow(`Experiment does not exist: ${experiment.key}`)
+  })
+
+  test('should throw when a bucket key does not belong to its experiment', async () => {
+    const experiment = await experimentsDAO.save(mockExperiment())
+    await bucketsDAO.save(mockBucket(experiment.id, 'test', 100))
+
+    await expect(
+      abba.saveManualUserAssignments([
+        { userId: mockUserId1, experimentKey: experiment.key, bucketKey: 'control' },
+      ]),
+    ).rejects.toThrow(`Bucket does not exist on experiment ${experiment.key}: control`)
+  })
+
+  test.each([
+    ['Active', AssignmentStatus.Active],
+    ['Paused', AssignmentStatus.Paused],
+    ['Inactive', AssignmentStatus.Inactive],
+  ])('should succeed when the experiment status is %s', async (_name, status) => {
+    const experiment = await experimentsDAO.save(mockExperiment({ status }))
+    const bucket = await bucketsDAO.save(mockBucket(experiment.id, 'test', 100))
+
+    const [result] = await abba.saveManualUserAssignments([
+      { userId: mockUserId1, experimentKey: experiment.key, bucketKey: bucket.key },
+    ])
+
+    expect(result!.bucketId).toBe(bucket.id)
+  })
+
+  test('should ignore segmentation rules and sampling', async () => {
+    const experiment = await experimentsDAO.save(
+      mockExperiment({
+        sampling: 0,
+        rules: [{ key: 'country', operator: SegmentationRuleOperator.EqualsText, value: 'SE' }],
+      }),
+    )
+    const bucket = await bucketsDAO.save(mockBucket(experiment.id, 'test', 100))
+
+    const [result] = await abba.saveManualUserAssignments([
+      { userId: mockUserId1, experimentKey: experiment.key, bucketKey: bucket.key },
+    ])
+
+    expect(result!.bucketId).toBe(bucket.id)
+  })
+})

packages/abba/src/abba.ts

@@ -1,4 +1,4 @@
-import { _shuffle } from '@naturalcycles/js-lib/array/array.util.js'
+import { _mapBy, _shuffle, _uniq } from '@naturalcycles/js-lib/array/array.util.js'
 import { localTime } from '@naturalcycles/js-lib/datetime/localTime.js'
 import { _Memo } from '@naturalcycles/js-lib/decorators/memo.decorator.js'
 import { _assert } from '@naturalcycles/js-lib/error/assert.js'
@@ -19,6 +19,7 @@ import type {
   ExperimentAssignmentStatistics,
   ExperimentInput,
   ExperimentWithBuckets,
+  ManualUserAssignmentInput,
   SegmentationData,
   UserAssignment,
   UserExperiment,
@@ -328,6 +329,84 @@ export class Abba {
     }
   }
 
+  /**
+   * Manually assigns users to specific buckets, overwriting any existing assignments.
+   * Bypasses sampling, segmentation, exclusions, and experiment status. Intended for
+   * QA / internal testing where deterministic bucket placement is required.
+   *
+   * An empty input returns an empty array. If the input contains duplicate
+   * (userId, experimentKey) pairs, the last occurrence wins.
+   *
+   * Throws AbbaErrorCode.ExperimentNotFound or AbbaErrorCode.BucketNotFound on the
+   * first invalid row; no rows are written if validation fails.
+   *
+   * Cold method.
+   */
+  async saveManualUserAssignments(
+    inputs: readonly ManualUserAssignmentInput[],
+  ): Promise<DecoratedUserAssignment[]> {
+    if (!inputs.length) return []
+
+    const dedupedInputs = Array.from(
+      _mapBy(inputs, input => `${input.userId}|${input.experimentKey}`).values(),
+    )
+
+    const experimentKeys = _uniq(dedupedInputs.map(input => input.experimentKey))
+    const experiments = await pMap(experimentKeys, async experimentKey => {
+      const experiment = await this.experimentDao.getByKey(experimentKey)
+      _assert(experiment && !experiment.deleted, `Experiment does not exist: ${experimentKey}`, {
+        code: AbbaErrorCode.ExperimentNotFound,
+      })
+      const buckets = await this.bucketDao.getByExperimentId(experiment.id)
+      return { ...experiment, buckets }
+    })
+    const experimentByKey = _mapBy(experiments, experiment => experiment.key)
+
+    const resolvedInputs = dedupedInputs.map(input => {
+      const experiment = experimentByKey.get(input.experimentKey)!
+      const bucket = experiment.buckets.find(bucket => bucket.key === input.bucketKey)
+      _assert(
+        bucket,
+        `Bucket does not exist on experiment ${input.experimentKey}: ${input.bucketKey}`,
+        { code: AbbaErrorCode.BucketNotFound },
+      )
+      return { input, experiment, bucket }
+    })
+
+    const userIds = _uniq(dedupedInputs.map(input => input.userId))
+    const experimentIds = experiments.map(experiment => experiment.id)
+    const existingAssignments = await this.userAssignmentDao.getByUserIdsAndExperimentIds(
+      userIds,
+      experimentIds,
+    )
+    const existingByKey = _mapBy(
+      existingAssignments,
+      assignment => `${assignment.userId}|${assignment.experimentId}`,
+    )
+
+    const toSave: Unsaved<UserAssignment>[] = resolvedInputs.map(
+      ({ input, experiment, bucket }) => ({
+        ...existingByKey.get(`${input.userId}|${experiment.id}`),
+        userId: input.userId,
+        experimentId: experiment.id,
+        bucketId: bucket.id,
+      }),
+    )
+
+    const savedAssignments = await this.userAssignmentDao.saveBatch(toSave)
+
+    return savedAssignments.map((savedAssignment, i) => {
+      const resolvedInput = resolvedInputs[i]!
+      return {
+        ...savedAssignment,
+        experimentKey: resolvedInput.experiment.key,
+        experimentData: resolvedInput.experiment.data,
+        bucketKey: resolvedInput.bucket.key,
+        bucketData: resolvedInput.bucket.data,
+      }
+    })
+  }
+
   /**
    * Get all existing user assignments.
    * Hot method.

packages/abba/src/dao/userAssignment.dao.ts

@@ -21,6 +21,20 @@ export class UserAssignmentDao extends CommonDao<UserAssignment> {
     return await this.runQuery(query)
   }
 
+  /**
+   * Returns every UserAssignment whose `userId` is in `userIds` AND whose `experimentId`
+   * is in `experimentIds`. This is the cross-product, not paired lookup: callers must
+   * filter the result by the specific (userId, experimentId) pairs they care about.
+   */
+  async getByUserIdsAndExperimentIds(
+    userIds: string[],
+    experimentIds: string[],
+  ): Promise<UserAssignment[]> {
+    if (!userIds.length || !experimentIds.length) return []
+    const query = this.query().filterIn('userId', userIds).filterIn('experimentId', experimentIds)
+    return await this.runQuery(query)
+  }
+
   async deleteByExperimentId(experimentId: string): Promise<void> {
     await this.query().filterEq('experimentId', experimentId).deleteByQuery()
   }

packages/abba/src/types.ts

@@ -7,6 +7,7 @@ export const AbbaErrorCode = {
   SegmentationDataRequired: 'abba/segmentationDataRequired',
   InvalidBucketRatio: 'abba/invalidBucketRatio',
   BucketDeterminationFailed: 'abba/bucketDeterminationFailed',
+  BucketNotFound: 'abba/bucketNotFound',
 } as const
 
 export interface AbbaConfig {
@@ -92,6 +93,12 @@ export type UserAssignment = BaseDBEntity & {
   bucketId: string | null
 }
 
+export interface ManualUserAssignmentInput {
+  userId: string
+  experimentKey: string
+  bucketKey: string
+}
+
 export type DecoratedUserAssignment = UserAssignment & {
   experimentKey: Experiment['key']
   experimentData: Experiment['data']