@naturalcycles/abba 1.7.0 → 1.8.0

package/src/abba.ts

@@ -1,41 +1,55 @@
-import { Bucket, Experiment, Prisma, PrismaClient, UserAssignment } from './prisma-output'
-import { AssignmentStatus } from './types'
-import { determineAssignment, validateSegmentationRules, validateTotalBucketRatio } from './util'
+import { pMap, Saved } from '@naturalcycles/js-lib'
 import {
-  BucketInput,
+  AbbaConfig,
+  AssignmentStatus,
+  Bucket,
+  Experiment,
   ExperimentWithBuckets,
-  ExperimentInput,
-  SegmentationData,
-  SegmentationRule,
-  AssignmentStatistics,
-} from '.'
-
-// Note: Schema currently contains an output dir which generates all the files to the prisma dir
-// it would be tidier not to include it here when possible later on:
-// Explanation is here: https://github.com/prisma/prisma/issues/9435#issuecomment-960290681
+  UserAssignment,
+} from './types'
+import { determineAssignment, validateSegmentationRules, validateTotalBucketRatio } from './util'
+import { ExperimentDao, experimentDao } from './dao/experiment.dao'
+import { UserAssignmentDao, userAssignmentDao } from './dao/userAssignment.dao'
+import { BucketDao, bucketDao } from './dao/bucket.dao'
+import { SegmentationData, SegmentationRule, AssignmentStatistics } from '.'
 
 export class Abba {
-  private client: PrismaClient
+  private experimentDao: ExperimentDao
+  private bucketDao: BucketDao
+  private userAssignmentDao: UserAssignmentDao
 
-  constructor(dbUrl?: string) {
-    this.client = new PrismaClient({
-      datasources: {
-        db: {
-          url: dbUrl,
-        },
-      },
-    })
+  constructor({ db }: AbbaConfig) {
+    this.experimentDao = experimentDao(db)
+    this.bucketDao = bucketDao(db)
+    this.userAssignmentDao = userAssignmentDao(db)
   }
+
+  // TODO: Cache me
   /**
    * Returns all experiments
    *
    * @returns
    */
   async getAllExperiments(excludeInactive: boolean = false): Promise<ExperimentWithBuckets[]> {
-    return await this.client.experiment.findMany({
-      where: excludeInactive ? { NOT: { status: AssignmentStatus.Inactive } } : undefined,
-      include: { buckets: true },
-    })
+    const query = this.experimentDao.query()
+    if (excludeInactive) {
+      query.filter('status', '!=', AssignmentStatus.Inactive)
+    }
+
+    const experiments = await this.experimentDao.runQuery(query)
+    const buckets = await this.bucketDao
+      .query()
+      .filter(
+        'experimentId',
+        'in',
+        experiments.map(e => e.id),
+      )
+      .runQuery()
+
+    return experiments.map(experiment => ({
+      ...experiment,
+      buckets: buckets.filter(bucket => bucket.experimentId === experiment.id),
+    }))
   }
 
   /**
@@ -46,28 +60,21 @@ export class Abba {
    * @returns
    */
   async createExperiment(
-    experiment: ExperimentInput,
-    buckets: BucketInput[],
+    experiment: Experiment,
+    buckets: Bucket[],
   ): Promise<ExperimentWithBuckets> {
     if (experiment.status === AssignmentStatus.Active) {
       validateTotalBucketRatio(buckets)
     }
 
-    const created = await this.client.experiment.create({
-      data: {
-        ...experiment,
-        rules: experiment.rules as Prisma.InputJsonArray,
-        buckets: {
-          createMany: {
-            data: buckets,
-          },
-        },
-      },
-      include: {
-        buckets: true,
-      },
-    })
-    return created
+    const created = await this.experimentDao.save(experiment)
+
+    return {
+      ...created,
+      buckets: await this.bucketDao.saveBatch(
+        buckets.map(b => ({ ...b, experimentId: created.id })),
+      ),
+    }
   }
 
   /**
@@ -81,19 +88,21 @@ export class Abba {
    */
   async saveExperiment(
     id: number,
-    experiment: ExperimentInput,
-    buckets: BucketInput[],
+    experiment: Experiment,
+    buckets: Bucket[],
   ): Promise<ExperimentWithBuckets> {
     if (experiment.status === AssignmentStatus.Active) {
       validateTotalBucketRatio(buckets)
     }
 
-    const updatedExperiment = await this.updateExperiment(id, experiment)
-    const updatedBuckets = await this.saveBuckets(buckets)
+    const updated = await this.experimentDao.save({
+      ...experiment,
+      id,
+    })
 
     return {
-      ...updatedExperiment,
-      buckets: updatedBuckets,
+      ...updated,
+      buckets: await this.bucketDao.saveBatch(buckets.map(b => ({ ...b, experimentId: id }))),
     }
   }
 
@@ -103,7 +112,7 @@ export class Abba {
    * @param id
    */
   async deleteExperiment(id: number): Promise<void> {
-    await this.client.experiment.delete({ where: { id } })
+    await this.experimentDao.deleteById(id)
   }
 
   /**
@@ -111,8 +120,8 @@ export class Abba {
    *
    * @param experimentId
    * @param userId
-   * @param createNew
-   * @param segmentationData
+   * @param existingOnly Do not generate any new assignments for this experiment
+   * @param segmentationData Required if existingOnly is false
    * @returns
    */
   async getUserAssignment(
@@ -120,13 +129,12 @@ export class Abba {
     userId: string,
     existingOnly: boolean,
     segmentationData?: SegmentationData,
-  ): Promise<UserAssignment | null> {
-    const experiment = await this.client.experiment.findUnique({
-      where: { id: experimentId },
-      include: { buckets: true },
-    })
+  ): Promise<Saved<UserAssignment> | null> {
+    const experiment = await this.experimentDao.requireById(experimentId)
     if (!experiment) throw new Error('Experiment not found')
 
+    const buckets = await this.bucketDao.getBy('experimentId', experimentId)
+
     const existing = await this.getExistingUserAssignment(experimentId, userId)
     if (existing) return existing
 
@@ -135,17 +143,17 @@ export class Abba {
     if (!segmentationData)
       throw new Error('Segmentation data required when creating a new assignment')
 
-    return await this.generateUserAssignment(experiment, userId, segmentationData)
+    return await this.generateUserAssignment({ ...experiment, buckets }, userId, segmentationData)
   }
 
   /**
    * Get all existing user assignments
    *
-   * @param userId G
+   * @param userId
    * @returns
    */
-  async getAllExistingUserAssignments(userId: string): Promise<UserAssignment[]> {
-    return await this.client.userAssignment.findMany({ where: { userId } })
+  async getAllExistingUserAssignments(userId: string): Promise<Saved<UserAssignment>[]> {
+    return await this.userAssignmentDao.getBy('userId', userId)
   }
 
   /**
@@ -158,12 +166,12 @@ export class Abba {
   async generateUserAssignments(
     userId: string,
     segmentationData: SegmentationData,
-  ): Promise<UserAssignment[]> {
+  ): Promise<Saved<UserAssignment>[]> {
     const experiments = await this.getAllExperiments(true)
     const existingAssignments = await this.getAllExistingUserAssignments(userId)
 
-    const assignments: UserAssignment[] = []
-    const generatedAssignments: Promise<UserAssignment | null>[] = []
+    const assignments: Saved<UserAssignment>[] = []
+    const generatedAssignments: Promise<Saved<UserAssignment> | null>[] = []
 
     for (const experiment of experiments) {
       const existing = existingAssignments.find(ua => ua.experimentId === experiment.id)
@@ -176,7 +184,7 @@ export class Abba {
     }
 
     const generated = await Promise.all(generatedAssignments)
-    const filtered = generated.filter((ua): ua is UserAssignment => ua !== null)
+    const filtered = generated.filter((ua): ua is Saved<UserAssignment> => ua !== null)
     return [...assignments, ...filtered]
   }
 
@@ -188,21 +196,17 @@ export class Abba {
    */
   async getExperimentAssignmentStatistics(experimentId: number): Promise<AssignmentStatistics> {
     const statistics = {
-      sampled: await this.client.userAssignment.count({ where: { experimentId } }),
+      sampled: await this.userAssignmentDao
+        .query()
+        .filterEq('experimentId', experimentId)
+        .runQueryCount(),
       buckets: {},
     }
 
-    const buckets = await this.client.bucket.findMany({ where: { experimentId } })
-    const assignmentCounts = await this.client.userAssignment.groupBy({
-      where: { experimentId },
-      by: ['bucketId'],
-      _count: {
-        _all: true,
-      },
-    })
-
-    buckets.forEach(({ id }) => {
-      statistics.buckets[`${id}`] = assignmentCounts.find(i => i.bucketId === id)?._count?._all || 0
+    const buckets = await this.bucketDao.getBy('experimentId', experimentId)
+    await pMap(buckets, async bucket => {
+      const count = this.userAssignmentDao.query().filterEq('bucketId', bucket.id)
+      statistics[bucket.id] = await this.userAssignmentDao.runQueryCount(count)
     })
 
     return statistics
@@ -220,16 +224,17 @@ export class Abba {
     experiment: ExperimentWithBuckets,
     userId: string,
     segmentationData: SegmentationData,
-  ): Promise<UserAssignment | null> {
+  ): Promise<Saved<UserAssignment> | null> {
     const segmentationMatch = validateSegmentationRules(
       experiment.rules as unknown as SegmentationRule[],
       segmentationData,
     )
     if (!segmentationMatch) return null
 
-    const bucketId = determineAssignment(experiment.sampling, experiment.buckets)
-    return await this.client.userAssignment.create({
-      data: { userId, experimentId: experiment.id, bucketId },
+    return await this.userAssignmentDao.save({
+      userId,
+      experimentId: experiment.id,
+      bucketId: determineAssignment(experiment.sampling, experiment.buckets),
     })
   }
 
@@ -243,58 +248,9 @@ export class Abba {
   private async getExistingUserAssignment(
     experimentId: number,
     userId: string,
-  ): Promise<UserAssignment | null> {
-    return await this.client.userAssignment.findFirst({ where: { userId, experimentId } })
-  }
-
-  /**
-   * Update experiment information
-   *
-   * @param id
-   * @param experiment
-   * @param rules
-   * @returns
-   */
-  private async updateExperiment(id: number, experiment: Partial<Experiment>): Promise<Experiment> {
-    return await this.client.experiment.update({
-      where: { id },
-      data: {
-        ...experiment,
-        rules: experiment.rules as Prisma.InputJsonArray,
-      },
-    })
-  }
-
-  /**
-   * Upserts bucket info
-   *
-   * @param id
-   * @param experiment
-   * @returns
-   */
-  private async saveBuckets(buckets: BucketInput[]): Promise<Bucket[]> {
-    const savedBuckets: Promise<Bucket>[] = []
-    for (const bucket of buckets) {
-      const { id, ...data } = bucket
-      if (id) {
-        savedBuckets.push(this.client.bucket.update({ where: { id }, data }))
-      } else {
-        savedBuckets.push(
-          this.client.bucket.create({
-            data: {
-              key: data.key,
-              ratio: data.ratio,
-              experiment: {
-                connect: {
-                  id: data.experimentId,
-                },
-              },
-            },
-          }),
-        )
-      }
-    }
-
-    return await Promise.all(savedBuckets)
+  ): Promise<Saved<UserAssignment> | null> {
+    const assignments = await this.userAssignmentDao.getBy('userId', userId)
+    const assignment = assignments.find(assignment => assignment.experimentId === experimentId)
+    return assignment || null
   }
 }

package/src/dao/bucket.dao.ts

@@ -0,0 +1,13 @@
+import { CommonDao, CommonDB } from '@naturalcycles/db-lib'
+import { Bucket } from '../types'
+
+export class BucketDao extends CommonDao<Bucket> {}
+
+export const bucketDao = (db: CommonDB): BucketDao =>
+  new BucketDao({
+    db,
+    table: 'Bucket',
+    createId: false, // mysql auto_increment is used instead
+    idType: 'number',
+    assignGeneratedIds: true,
+  })

package/src/dao/experiment.dao.ts

@@ -0,0 +1,22 @@
+import { CommonDao, CommonDB } from '@naturalcycles/db-lib'
+import { Saved } from '@naturalcycles/js-lib'
+import { BaseExperiment, Experiment } from '../types'
+
+type ExperimentDBM = Saved<BaseExperiment> & {
+  rules: string | null
+}
+
+export class ExperimentDao extends CommonDao<Experiment, ExperimentDBM> {}
+
+export const experimentDao = (db: CommonDB): ExperimentDao =>
+  new ExperimentDao({
+    db,
+    table: 'Experiment',
+    createId: false, // mysql auto_increment is used instead
+    idType: 'number',
+    assignGeneratedIds: true,
+    hooks: {
+      beforeBMToDBM: bm => ({ ...bm, rules: bm.rules.length ? JSON.stringify(bm.rules) : null }),
+      beforeDBMToBM: dbm => ({ ...dbm, rules: dbm.rules && JSON.parse(dbm.rules) }),
+    },
+  })

package/src/dao/userAssignment.dao.ts

@@ -0,0 +1,13 @@
+import { CommonDao, CommonDB } from '@naturalcycles/db-lib'
+import { UserAssignment } from '../types'
+
+export class UserAssignmentDao extends CommonDao<UserAssignment> {}
+
+export const userAssignmentDao = (db: CommonDB): UserAssignmentDao =>
+  new UserAssignmentDao({
+    db,
+    table: 'UserAssignment',
+    createId: false, // mysql auto_increment is used instead
+    idType: 'number',
+    assignGeneratedIds: true,
+  })

package/src/index.ts

@@ -1,5 +1,2 @@
-export type { Bucket, Experiment, UserAssignment } from './prisma-output'
-
 export * from './types'
-
 export { Abba } from './abba'

package/src/migrations/init.sql

@@ -0,0 +1,47 @@
+-- CreateTable
+CREATE TABLE IF NOT EXISTS `Bucket` (
+    `id` INTEGER NOT NULL AUTO_INCREMENT,
+    `experimentId` INTEGER NOT NULL,
+    `key` VARCHAR(10) NOT NULL,
+    `ratio` INTEGER NOT NULL,
+    `created` INTEGER(11) NOT NULL,
+    `updated` INTEGER(11) NOT NULL,
+
+    PRIMARY KEY (`id`)
+) DEFAULT CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci;
+
+-- CreateTable
+CREATE TABLE IF NOT EXISTS `Experiment` (
+    `id` INTEGER NOT NULL AUTO_INCREMENT,
+    `name` VARCHAR(191) NOT NULL,
+    `status` INTEGER NOT NULL,
+    `sampling` INTEGER NOT NULL,
+    `description` VARCHAR(240) NULL,
+    `created` INTEGER(11) NOT NULL,
+    `updated` INTEGER(11) NOT NULL,
+    `rules` JSON NULL,
+
+    PRIMARY KEY (`id`)
+) DEFAULT CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci;
+
+-- CreateTable
+CREATE TABLE IF NOT EXISTS `UserAssignment` (
+    `id` INTEGER NOT NULL AUTO_INCREMENT,
+    `userId` VARCHAR(191) NOT NULL,
+    `experimentId` INTEGER NOT NULL,
+    `bucketId` INTEGER NULL,
+    `created` INTEGER(11) NOT NULL,
+    `updated` INTEGER(11) NOT NULL,
+
+    UNIQUE INDEX `UserAssignment_userId_experimentId_key`(`userId`, `experimentId`),
+    PRIMARY KEY (`id`)
+) DEFAULT CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci;
+
+-- AddForeignKey
+ALTER TABLE `Bucket` ADD CONSTRAINT `Bucket_experimentId_fkey` FOREIGN KEY (`experimentId`) REFERENCES `Experiment`(`id`) ON DELETE CASCADE ON UPDATE CASCADE;
+
+-- AddForeignKey
+ALTER TABLE `UserAssignment` ADD CONSTRAINT `UserAssignment_bucketId_fkey` FOREIGN KEY (`bucketId`) REFERENCES `Bucket`(`id`) ON DELETE CASCADE ON UPDATE CASCADE;
+
+-- AddForeignKey
+ALTER TABLE `UserAssignment` ADD CONSTRAINT `UserAssignment_experimentId_fkey` FOREIGN KEY (`experimentId`) REFERENCES `Experiment`(`id`) ON DELETE CASCADE ON UPDATE CASCADE;

package/src/types.ts

@@ -1,16 +1,42 @@
-import { Bucket, Experiment } from './prisma-output'
+import { CommonDB } from '@naturalcycles/db-lib'
+import { BaseDBEntity, Saved } from '@naturalcycles/js-lib'
 
-export type Unsaved<T> = Omit<T, 'createdAt' | 'updatedAt'> & {
-  id?: number
+export interface AbbaConfig {
+  db: CommonDB
 }
 
-export type ExperimentWithBuckets = Experiment & {
-  buckets: Bucket[]
+export type BaseExperiment = BaseDBEntity<number> & {
+  name: string
+  status: number
+  sampling: number
+  description: string | null
 }
 
-export type ExperimentInput = Unsaved<Experiment>
+export type Experiment = BaseExperiment & {
+  rules: SegmentationRule[]
+}
+
+export type ExperimentWithBuckets = Saved<Experiment> & {
+  buckets: Saved<Bucket>[]
+}
+
+export interface BucketInput {
+  experimentId: number
+  key: string
+  ratio: number
+}
 
-export type BucketInput = Unsaved<Bucket>
+export type Bucket = BaseDBEntity<number> & {
+  experimentId: number
+  key: string
+  ratio: number
+}
+
+export type UserAssignment = BaseDBEntity<number> & {
+  userId: string
+  experimentId: number
+  bucketId: number | null
+}
 
 export type SegmentationData = Record<string, string | boolean | number>
 

package/src/util.ts

@@ -1,6 +1,6 @@
+import { Saved } from '@naturalcycles/js-lib'
 import { satisfies } from 'semver'
-import { Bucket } from './prisma-output'
-import { BucketInput, SegmentationData, SegmentationRule } from '.'
+import { Bucket, SegmentationData, SegmentationRule } from './types'
 
 /**
  * Generate a random number between 0 and 100
@@ -18,7 +18,7 @@ export const rollDie = (): number => {
  * @param buckets
  * @returns
  */
-export const determineAssignment = (sampling: number, buckets: Bucket[]): number | null => {
+export const determineAssignment = (sampling: number, buckets: Saved<Bucket>[]): number | null => {
   // Should this person be considered for the experiment?
   if (rollDie() > sampling) {
     return null
@@ -34,7 +34,7 @@ export const determineAssignment = (sampling: number, buckets: Bucket[]): number
  * @param buckets
  * @returns
  */
-export const determineBucket = (buckets: Bucket[]): number => {
+export const determineBucket = (buckets: Saved<Bucket>[]): number => {
   const bucketRoll = rollDie()
   let range: [number, number] | undefined
   const bucket = buckets.find(b => {
@@ -62,7 +62,7 @@ export const determineBucket = (buckets: Bucket[]): number => {
  * @param buckets
  * @returns
  */
-export const validateTotalBucketRatio = (buckets: BucketInput[]): void => {
+export const validateTotalBucketRatio = (buckets: Bucket[]): void => {
   const bucketSum = buckets.reduce((sum, current) => sum + current.ratio, 0)
   if (bucketSum !== 100) {
     throw new Error('Total bucket ratio must be 100 before you can activate an experiment')

package/dist/prisma-output/index-browser.js

@@ -1,141 +0,0 @@
-Object.defineProperty(exports, '__esModule', { value: true })
-
-const { Decimal } = require('./runtime/index-browser')
-
-const Prisma = {}
-
-exports.Prisma = Prisma
-
-/**
- * Prisma Client JS version: 3.13.0
- * Query Engine version: efdf9b1183dddfd4258cd181a72125755215ab7b
- */
-Prisma.prismaVersion = {
-  client: '3.13.0',
-  engine: 'efdf9b1183dddfd4258cd181a72125755215ab7b',
-}
-
-Prisma.PrismaClientKnownRequestError = () => {
-  throw new Error(`PrismaClientKnownRequestError is unable to be run in the browser.
-In case this error is unexpected for you, please report it in https://github.com/prisma/prisma/issues`)
-}
-Prisma.PrismaClientUnknownRequestError = () => {
-  throw new Error(`PrismaClientUnknownRequestError is unable to be run in the browser.
-In case this error is unexpected for you, please report it in https://github.com/prisma/prisma/issues`)
-}
-Prisma.PrismaClientRustPanicError = () => {
-  throw new Error(`PrismaClientRustPanicError is unable to be run in the browser.
-In case this error is unexpected for you, please report it in https://github.com/prisma/prisma/issues`)
-}
-Prisma.PrismaClientInitializationError = () => {
-  throw new Error(`PrismaClientInitializationError is unable to be run in the browser.
-In case this error is unexpected for you, please report it in https://github.com/prisma/prisma/issues`)
-}
-Prisma.PrismaClientValidationError = () => {
-  throw new Error(`PrismaClientValidationError is unable to be run in the browser.
-In case this error is unexpected for you, please report it in https://github.com/prisma/prisma/issues`)
-}
-Prisma.Decimal = Decimal
-
-/**
- * Re-export of sql-template-tag
- */
-Prisma.sql = () => {
-  throw new Error(`sqltag is unable to be run in the browser.
-In case this error is unexpected for you, please report it in https://github.com/prisma/prisma/issues`)
-}
-Prisma.empty = () => {
-  throw new Error(`empty is unable to be run in the browser.
-In case this error is unexpected for you, please report it in https://github.com/prisma/prisma/issues`)
-}
-Prisma.join = () => {
-  throw new Error(`join is unable to be run in the browser.
-In case this error is unexpected for you, please report it in https://github.com/prisma/prisma/issues`)
-}
-Prisma.raw = () => {
-  throw new Error(`raw is unable to be run in the browser.
-In case this error is unexpected for you, please report it in https://github.com/prisma/prisma/issues`)
-}
-Prisma.validator = () => val => val
-
-/**
- * Shorthand utilities for JSON filtering
- */
-Prisma.DbNull = 'DbNull'
-Prisma.JsonNull = 'JsonNull'
-Prisma.AnyNull = 'AnyNull'
-
-/**
- * Enums
- */
-// Based on
-// https://github.com/microsoft/TypeScript/issues/3192#issuecomment-261720275
-function makeEnum(x) {
-  return x
-}
-
-exports.Prisma.BucketScalarFieldEnum = makeEnum({
-  id: 'id',
-  experimentId: 'experimentId',
-  key: 'key',
-  ratio: 'ratio',
-  createdAt: 'createdAt',
-  updatedAt: 'updatedAt',
-})
-
-exports.Prisma.ExperimentScalarFieldEnum = makeEnum({
-  id: 'id',
-  name: 'name',
-  status: 'status',
-  sampling: 'sampling',
-  createdAt: 'createdAt',
-  updatedAt: 'updatedAt',
-  description: 'description',
-  rules: 'rules',
-})
-
-exports.Prisma.UserAssignmentScalarFieldEnum = makeEnum({
-  id: 'id',
-  userId: 'userId',
-  experimentId: 'experimentId',
-  bucketId: 'bucketId',
-  createdAt: 'createdAt',
-  updatedAt: 'updatedAt',
-})
-
-exports.Prisma.SortOrder = makeEnum({
-  asc: 'asc',
-  desc: 'desc',
-})
-
-exports.Prisma.NullableJsonNullValueInput = makeEnum({
-  DbNull: 'DbNull',
-  JsonNull: 'JsonNull',
-})
-
-exports.Prisma.JsonNullValueFilter = makeEnum({
-  DbNull: 'DbNull',
-  JsonNull: 'JsonNull',
-  AnyNull: 'AnyNull',
-})
-
-exports.Prisma.ModelName = makeEnum({
-  Bucket: 'Bucket',
-  Experiment: 'Experiment',
-  UserAssignment: 'UserAssignment',
-})
-
-/**
- * Create the Client
- */
-class PrismaClient {
-  constructor() {
-    throw new Error(
-      `PrismaClient is unable to be run in the browser.
-In case this error is unexpected for you, please report it in https://github.com/prisma/prisma/issues`,
-    )
-  }
-}
-exports.PrismaClient = PrismaClient
-
-Object.assign(exports, Prisma)