@naturalcycles/abba 2.11.0 → 2.12.0

package/dist/abba.d.ts

@@ -1,6 +1,6 @@
 import type { Unsaved } from '@naturalcycles/js-lib/types';
 import type { GetAllExperimentsOpts } from './dao/experiment.dao.js';
-import type { AbbaConfig, Bucket, BucketInput, DecoratedUserAssignment, Experiment, ExperimentAssignmentStatistics, ExperimentInput, ExperimentWithBuckets, SegmentationData, UserExperiment } from './types.js';
+import type { AbbaConfig, Bucket, BucketInput, DecoratedUserAssignment, Experiment, ExperimentAssignmentStatistics, ExperimentInput, ExperimentWithBuckets, ManualUserAssignmentInput, SegmentationData, UserExperiment } from './types.js';
 export declare class Abba {
     cfg: AbbaConfig;
     private experimentDao;
@@ -52,6 +52,20 @@ export declare class Abba {
      * @param segmentationData Required if existingOnly is false
      */
     getUserAssignment(experimentKey: string, userId: string, existingOnly: boolean, segmentationData?: SegmentationData): Promise<DecoratedUserAssignment | null>;
+    /**
+     * 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.
+     */
+    saveManualUserAssignments(inputs: readonly ManualUserAssignmentInput[]): Promise<DecoratedUserAssignment[]>;
     /**
      * Get all existing user assignments.
      * Hot method.

package/dist/abba.js

@@ -1,5 +1,5 @@
 import { __decorate } from "tslib";
-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';
@@ -223,6 +223,61 @@ export class Abba {
             bucketData: bucket?.data || null,
         };
     }
+    /**
+     * 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) {
+        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 = 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.

package/dist/dao/userAssignment.dao.d.ts

@@ -4,6 +4,12 @@ import type { UserAssignment } from '../types.js';
 export declare class UserAssignmentDao extends CommonDao<UserAssignment> {
     getUserAssignmentByExperimentId(userId: string, experimentId: string): Promise<UserAssignment | null>;
     getUserAssigmentsByExperimentIds(userId: string, experimentIds: string[]): Promise<UserAssignment[]>;
+    /**
+     * 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.
+     */
+    getByUserIdsAndExperimentIds(userIds: string[], experimentIds: string[]): Promise<UserAssignment[]>;
     deleteByExperimentId(experimentId: string): Promise<void>;
     getCountByExperimentId(experimentId: string): Promise<number>;
     getCountByBucketId(bucketId: string): Promise<number>;

package/dist/dao/userAssignment.dao.js

@@ -10,6 +10,17 @@ export class UserAssignmentDao extends CommonDao {
         const query = this.query().filterEq('userId', userId).filterIn('experimentId', experimentIds);
         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, experimentIds) {
+        if (!userIds.length || !experimentIds.length)
+            return [];
+        const query = this.query().filterIn('userId', userIds).filterIn('experimentId', experimentIds);
+        return await this.runQuery(query);
+    }
     async deleteByExperimentId(experimentId) {
         await this.query().filterEq('experimentId', experimentId).deleteByQuery();
     }

package/dist/types.d.ts

@@ -6,6 +6,7 @@ export declare const AbbaErrorCode: {
     readonly SegmentationDataRequired: 'abba/segmentationDataRequired';
     readonly InvalidBucketRatio: 'abba/invalidBucketRatio';
     readonly BucketDeterminationFailed: 'abba/bucketDeterminationFailed';
+    readonly BucketNotFound: 'abba/bucketNotFound';
 };
 export interface AbbaConfig {
     db: CommonDB;
@@ -81,6 +82,11 @@ export type UserAssignment = BaseDBEntity & {
     experimentId: string;
     bucketId: string | null;
 };
+export interface ManualUserAssignmentInput {
+    userId: string;
+    experimentKey: string;
+    bucketKey: string;
+}
 export type DecoratedUserAssignment = UserAssignment & {
     experimentKey: Experiment['key'];
     experimentData: Experiment['data'];

package/dist/types.js

@@ -4,6 +4,7 @@ export const AbbaErrorCode = {
     SegmentationDataRequired: 'abba/segmentationDataRequired',
     InvalidBucketRatio: 'abba/invalidBucketRatio',
     BucketDeterminationFailed: 'abba/bucketDeterminationFailed',
+    BucketNotFound: 'abba/bucketNotFound',
 };
 export var AssignmentStatus;
 (function (AssignmentStatus) {

package/package.json

@@ -1,7 +1,7 @@
 {
   "name": "@naturalcycles/abba",
   "type": "module",
-  "version": "2.11.0",
+  "version": "2.12.0",
   "dependencies": {
     "@naturalcycles/db-lib": "^10",
     "@naturalcycles/js-lib": "^15",

package/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.

package/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()
   }

package/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']