feat(server): user feature model (#9843)

close CLOUD-108

packages/backend/server/src/__tests__/models/feature-user.spec.ts

@@ -0,0 +1,95 @@
+import { TestingModule } from '@nestjs/testing';
+import { PrismaClient, User } from '@prisma/client';
+import ava, { TestFn } from 'ava';
+
+import { UserFeatureModel, UserModel } from '../../models';
+import { createTestingModule, initTestingDB } from '../utils';
+
+interface Context {
+  module: TestingModule;
+  model: UserFeatureModel;
+  u1: User;
+}
+
+const test = ava as TestFn<Context>;
+
+test.before(async t => {
+  const module = await createTestingModule({});
+
+  t.context.model = module.get(UserFeatureModel);
+  t.context.module = module;
+});
+
+test.beforeEach(async t => {
+  await initTestingDB(t.context.module.get(PrismaClient));
+  t.context.u1 = await t.context.module.get(UserModel).create({
+    email: 'u1@affine.pro',
+    registered: true,
+  });
+});
+
+test.after(async t => {
+  await t.context.module.close();
+});
+
+test('should get null if user feature not found', async t => {
+  const { model, u1 } = t.context;
+  const userFeature = await model.get(u1.id, 'ai_early_access');
+  t.is(userFeature, null);
+});
+
+test('should get user feature', async t => {
+  const { model, u1 } = t.context;
+  const userFeature = await model.get(u1.id, 'free_plan_v1');
+  t.is(userFeature?.feature, 'free_plan_v1');
+});
+
+test('should list user features', async t => {
+  const { model, u1 } = t.context;
+
+  t.like(await model.list(u1.id), ['free_plan_v1']);
+});
+
+test('should directly test user feature existence', async t => {
+  const { model, u1 } = t.context;
+
+  t.true(await model.has(u1.id, 'free_plan_v1'));
+  t.false(await model.has(u1.id, 'ai_early_access'));
+});
+
+test('should add user feature', async t => {
+  const { model, u1 } = t.context;
+
+  await model.add(u1.id, 'unlimited_copilot', 'test');
+  t.true(await model.has(u1.id, 'unlimited_copilot'));
+  t.true((await model.list(u1.id)).includes('unlimited_copilot'));
+});
+
+test('should not add existing user feature', async t => {
+  const { model, u1 } = t.context;
+
+  await model.add(u1.id, 'free_plan_v1', 'test');
+  await model.add(u1.id, 'free_plan_v1', 'test');
+
+  t.like(await model.list(u1.id), ['free_plan_v1']);
+});
+
+test('should remove user feature', async t => {
+  const { model, u1 } = t.context;
+
+  await model.remove(u1.id, 'free_plan_v1');
+  t.false(await model.has(u1.id, 'free_plan_v1'));
+  t.false((await model.list(u1.id)).includes('free_plan_v1'));
+});
+
+test('should switch user feature', async t => {
+  const { model, u1 } = t.context;
+
+  await model.switch(u1.id, 'free_plan_v1', 'pro_plan_v1', 'test');
+
+  t.false(await model.has(u1.id, 'free_plan_v1'));
+  t.true(await model.has(u1.id, 'pro_plan_v1'));
+
+  t.false((await model.list(u1.id)).includes('free_plan_v1'));
+  t.true((await model.list(u1.id)).includes('pro_plan_v1'));
+});

packages/backend/server/src/__tests__/models/feature-workspace.spec.ts

@@ -0,0 +1,127 @@
+import { TestingModule } from '@nestjs/testing';
+import { PrismaClient, Workspace } from '@prisma/client';
+import ava, { TestFn } from 'ava';
+
+import { UserModel, WorkspaceFeatureModel, WorkspaceModel } from '../../models';
+import { createTestingModule, initTestingDB } from '../utils';
+
+interface Context {
+  module: TestingModule;
+  model: WorkspaceFeatureModel;
+  ws: Workspace;
+}
+
+const test = ava as TestFn<Context>;
+
+test.before(async t => {
+  const module = await createTestingModule({});
+
+  t.context.model = module.get(WorkspaceFeatureModel);
+  t.context.module = module;
+});
+
+test.beforeEach(async t => {
+  await initTestingDB(t.context.module.get(PrismaClient));
+  const u1 = await t.context.module.get(UserModel).create({
+    email: 'u1@affine.pro',
+    registered: true,
+  });
+
+  t.context.ws = await t.context.module.get(WorkspaceModel).create(u1.id);
+});
+
+test.after(async t => {
+  await t.context.module.close();
+});
+
+test('should get null if workspace feature not found', async t => {
+  const { model, ws } = t.context;
+  const userFeature = await model.get(ws.id, 'unlimited_workspace');
+  t.is(userFeature, null);
+});
+
+test('should directly test workspace feature existence', async t => {
+  const { model, ws } = t.context;
+
+  t.false(await model.has(ws.id, 'unlimited_workspace'));
+});
+
+test('should list empty workspace features', async t => {
+  const { model, ws } = t.context;
+
+  t.deepEqual(await model.list(ws.id), []);
+});
+
+test('should add workspace feature', async t => {
+  const { model, ws } = t.context;
+
+  await model.add(ws.id, 'unlimited_workspace', 'test');
+  t.is(
+    (await model.get(ws.id, 'unlimited_workspace'))?.feature,
+    'unlimited_workspace'
+  );
+  t.true(await model.has(ws.id, 'unlimited_workspace'));
+  t.true((await model.list(ws.id)).includes('unlimited_workspace'));
+});
+
+test('should add workspace feature with overrides', async t => {
+  const { model, ws } = t.context;
+
+  await model.add(ws.id, 'team_plan_v1', 'test');
+  const f1 = await model.get(ws.id, 'team_plan_v1');
+  await model.add(ws.id, 'team_plan_v1', 'test', { memberLimit: 100 });
+  const f2 = await model.get(ws.id, 'team_plan_v1');
+
+  t.not(f1!.configs.memberLimit, f2!.configs.memberLimit);
+  t.is(f2!.configs.memberLimit, 100);
+});
+
+test('should not add existing workspace feature', async t => {
+  const { model, ws } = t.context;
+
+  await model.add(ws.id, 'team_plan_v1', 'test');
+  await model.add(ws.id, 'team_plan_v1', 'test');
+
+  t.like(await model.list(ws.id), ['team_plan_v1']);
+});
+
+test('should replace existing workspace if overrides updated', async t => {
+  const { model, ws } = t.context;
+
+  await model.add(ws.id, 'team_plan_v1', 'test', { memberLimit: 10 });
+  await model.add(ws.id, 'team_plan_v1', 'test', { memberLimit: 100 });
+  const f2 = await model.get(ws.id, 'team_plan_v1');
+
+  t.is(f2!.configs.memberLimit, 100);
+});
+
+test('should remove workspace feature', async t => {
+  const { model, ws } = t.context;
+
+  await model.add(ws.id, 'team_plan_v1', 'test');
+  await model.remove(ws.id, 'team_plan_v1');
+  t.false(await model.has(ws.id, 'team_plan_v1'));
+  t.false((await model.list(ws.id)).includes('team_plan_v1'));
+});
+
+test('should switch workspace feature', async t => {
+  const { model, ws } = t.context;
+
+  await model.switch(ws.id, 'team_plan_v1', 'unlimited_workspace', 'test');
+
+  t.false(await model.has(ws.id, 'team_plan_v1'));
+  t.true(await model.has(ws.id, 'unlimited_workspace'));
+
+  t.false((await model.list(ws.id)).includes('team_plan_v1'));
+  t.true((await model.list(ws.id)).includes('unlimited_workspace'));
+});
+
+test('should switch workspace feature with overrides', async t => {
+  const { model, ws } = t.context;
+
+  await model.add(ws.id, 'unlimited_workspace', 'test');
+  await model.add(ws.id, 'team_plan_v1', 'test', { memberLimit: 100 });
+  const f2 = await model.get(ws.id, 'team_plan_v1');
+
+  t.is(f2!.configs.memberLimit, 100);
+});

packages/backend/server/src/models/common/feature.ts

@@ -55,3 +55,24 @@ export const Features = {
   restricted_plan_v1: quota(UserPlanQuotaConfig),
   team_plan_v1: quota(WorkspaceQuotaConfig),
 };
+
+export type UserFeatureName = keyof Pick<
+  typeof Features,
+  | 'early_access'
+  | 'ai_early_access'
+  | 'unlimited_copilot'
+  | 'administrator'
+  | 'free_plan_v1'
+  | 'pro_plan_v1'
+  | 'lifetime_pro_plan_v1'
+  | 'restricted_plan_v1'
+>;
+export type WorkspaceFeatureName = keyof Pick<
+  typeof Features,
+  'unlimited_workspace' | 'team_plan_v1'
+>;
+
+export type FeatureName = UserFeatureName | WorkspaceFeatureName;
+export type FeatureConfigs<T extends FeatureName> = z.infer<
+  (typeof Features)[T]['shape']['configs']
+>;

packages/backend/server/src/models/feature.ts

@@ -4,12 +4,12 @@ import { Feature } from '@prisma/client';
 import { z } from 'zod';
 
 import { BaseModel } from './base';
-import { Features, FeatureType } from './common';
-
-type FeatureNames = keyof typeof Features;
-type FeatureConfigs<T extends FeatureNames> = z.infer<
-  (typeof Features)[T]['shape']['configs']
->;
+import {
+  type FeatureConfigs,
+  type FeatureName,
+  Features,
+  FeatureType,
+} from './common';
 
 // TODO(@forehalo):
 //   `version` column in `features` table will deprecated because it's makes the whole system complicated without any benefits.
@@ -19,47 +19,23 @@ type FeatureConfigs<T extends FeatureNames> = z.infer<
 //   This is a huge burden for us and we should remove it.
 @Injectable()
 export class FeatureModel extends BaseModel {
-  async get<T extends FeatureNames>(name: T) {
-    const feature = await this.getLatest(name);
-
-    // All features are hardcoded in the codebase
-    // It would be a fatal error if the feature is not found in DB.
-    if (!feature) {
-      throw new Error(`Feature ${name} not found`);
-    }
-
-    const shape = this.getConfigShape(name);
-    const parseResult = shape.safeParse(feature.configs);
-
-    if (!parseResult.success) {
-      throw new Error(`Invalid feature config for ${name}`, {
-        cause: parseResult.error,
-      });
-    }
+  async get<T extends FeatureName>(name: T) {
+    const feature = await this.get_unchecked(name);
 
     return {
       ...feature,
-      configs: parseResult.data as FeatureConfigs<T>,
+      configs: this.check(name, feature.configs),
     };
   }
 
   @Transactional()
-  async upsert<T extends FeatureNames>(name: T, configs: FeatureConfigs<T>) {
-    const shape = this.getConfigShape(name);
-    const parseResult = shape.safeParse(configs);
-
-    if (!parseResult.success) {
-      throw new Error(`Invalid feature config for ${name}`, {
-        cause: parseResult.error,
-      });
-    }
-
-    const parsedConfigs = parseResult.data;
+  async upsert<T extends FeatureName>(name: T, configs: FeatureConfigs<T>) {
+    const parsedConfigs = this.check(name, configs);
 
     // TODO(@forehalo):
     //   could be a simple upsert operation, but we got useless `version` column in the database
     //   will be fixed when `version` column gets deprecated
-    const latest = await this.getLatest(name);
+    const latest = await this.try_get_unchecked(name);
 
     let feature: Feature;
     if (!latest) {
@@ -84,14 +60,54 @@ export class FeatureModel extends BaseModel {
     return feature as Feature & { configs: FeatureConfigs<T> };
   }
 
-  private async getLatest<T extends FeatureNames>(name: T) {
-    return this.tx.feature.findFirst({
+  /**
+   * Get the latest feature from database.
+   *
+   * @internal
+   */
+  async try_get_unchecked<T extends FeatureName>(name: T) {
+    const feature = await this.tx.feature.findFirst({
       where: { feature: name },
       orderBy: { version: 'desc' },
     });
+
+    return feature as Omit<Feature, 'configs'> & {
+      configs: Record<string, any>;
+    };
   }
 
-  private getConfigShape(name: FeatureNames): z.ZodObject<any> {
+  /**
+   * Get the latest feature from database.
+   *
+   * @throws {Error} If the feature is not found in DB.
+   * @internal
+   */
+  async get_unchecked<T extends FeatureName>(name: T) {
+    const feature = await this.try_get_unchecked(name);
+
+    // All features are hardcoded in the codebase
+    // It would be a fatal error if the feature is not found in DB.
+    if (!feature) {
+      throw new Error(`Feature ${name} not found`);
+    }
+
+    return feature;
+  }
+
+  check<T extends FeatureName>(name: T, config: any) {
+    const shape = this.getConfigShape(name);
+    const parseResult = shape.safeParse(config);
+
+    if (!parseResult.success) {
+      throw new Error(`Invalid feature config for ${name}`, {
+        cause: parseResult.error,
+      });
+    }
+
+    return parseResult.data as FeatureConfigs<T>;
+  }
+
+  getConfigShape(name: FeatureName): z.ZodObject<any> {
     return Features[name]?.shape.configs ?? z.object({});
   }
 }

packages/backend/server/src/models/index.ts

@@ -12,8 +12,10 @@ import { PageModel } from './page';
 import { MODELS_SYMBOL } from './provider';
 import { SessionModel } from './session';
 import { UserModel } from './user';
+import { UserFeatureModel } from './user-feature';
 import { VerificationTokenModel } from './verification-token';
 import { WorkspaceModel } from './workspace';
+import { WorkspaceFeatureModel } from './workspace-feature';
 
 const MODELS = {
   user: UserModel,
@@ -22,6 +24,8 @@ const MODELS = {
   feature: FeatureModel,
   workspace: WorkspaceModel,
   page: PageModel,
+  userFeature: UserFeatureModel,
+  workspaceFeature: WorkspaceFeatureModel,
 };
 
 type ModelsType = {
@@ -77,5 +81,7 @@ export * from './feature';
 export * from './page';
 export * from './session';
 export * from './user';
+export * from './user-feature';
 export * from './verification-token';
 export * from './workspace';
+export * from './workspace-feature';

packages/backend/server/src/models/user-feature.ts

@@ -0,0 +1,120 @@
+import { Injectable } from '@nestjs/common';
+import { Transactional } from '@nestjs-cls/transactional';
+
+import { BaseModel } from './base';
+import type { UserFeatureName } from './common';
+
+@Injectable()
+export class UserFeatureModel extends BaseModel {
+  async get<T extends UserFeatureName>(userId: string, name: T) {
+    // TODO(@forehalo):
+    // all feature query-and-use queries should be simplified like the below when `feature(name)` becomes a unique index
+    //
+    // this.db.userFeature.findFirst({
+    //   include: {
+    //     feature: true
+    //   },
+    //   where: {
+    //     userId,
+    //     activated: true,
+    //     feature: {
+    //       feature: name,
+    //     }
+    //   }
+    // })
+    const feature = await this.models.feature.get(name);
+
+    const count = await this.db.userFeature.count({
+      where: {
+        userId,
+        featureId: feature.id,
+        activated: true,
+      },
+    });
+
+    return count > 0 ? feature : null;
+  }
+
+  async has(userId: string, name: UserFeatureName) {
+    const feature = await this.models.feature.get_unchecked(name);
+
+    const count = await this.db.userFeature.count({
+      where: {
+        userId,
+        featureId: feature.id,
+        activated: true,
+      },
+    });
+
+    return count > 0;
+  }
+
+  async list(userId: string) {
+    const userFeatures = await this.db.userFeature.findMany({
+      include: {
+        feature: true,
+      },
+      where: {
+        userId,
+        activated: true,
+      },
+    });
+
+    return userFeatures.map(
+      userFeature => userFeature.feature.feature
+    ) as UserFeatureName[];
+  }
+
+  async add(userId: string, featureName: UserFeatureName, reason: string) {
+    const feature = await this.models.feature.get_unchecked(featureName);
+
+    const existing = await this.tx.userFeature.findFirst({
+      where: {
+        userId,
+        featureId: feature.id,
+        activated: true,
+      },
+    });
+
+    if (existing) {
+      return existing;
+    }
+
+    const userFeature = await this.tx.userFeature.create({
+      data: {
+        userId,
+        featureId: feature.id,
+        activated: true,
+        reason,
+      },
+    });
+
+    this.logger.verbose(`Feature ${featureName} added to user ${userId}`);
+
+    return userFeature;
+  }
+
+  async remove(userId: string, featureName: UserFeatureName) {
+    const feature = await this.models.feature.get_unchecked(featureName);
+
+    await this.tx.userFeature.deleteMany({
+      where: {
+        userId,
+        featureId: feature.id,
+      },
+    });
+
+    this.logger.verbose(`Feature ${featureName} removed from user ${userId}`);
+  }
+
+  @Transactional()
+  async switch(
+    userId: string,
+    from: UserFeatureName,
+    to: UserFeatureName,
+    reason: string
+  ) {
+    await this.remove(userId, from);
+    await this.add(userId, to, reason);
+  }
+}

packages/backend/server/src/models/workspace-feature.ts

@@ -0,0 +1,156 @@
+import { Injectable } from '@nestjs/common';
+import { Transactional } from '@nestjs-cls/transactional';
+
+import { BaseModel } from './base';
+import type { FeatureConfigs, WorkspaceFeatureName } from './common';
+
+@Injectable()
+export class WorkspaceFeatureModel extends BaseModel {
+  async get<T extends WorkspaceFeatureName>(workspaceId: string, name: T) {
+    const feature = await this.models.feature.get_unchecked(name);
+
+    const workspaceFeature = await this.tx.workspaceFeature.findFirst({
+      where: {
+        workspaceId,
+        featureId: feature.id,
+        activated: true,
+      },
+    });
+
+    if (!workspaceFeature) {
+      return null;
+    }
+
+    return {
+      ...feature,
+      configs: this.models.feature.check(name, {
+        ...feature.configs,
+        ...(workspaceFeature?.configs as {}),
+      }),
+    };
+  }
+
+  async has(workspaceId: string, name: WorkspaceFeatureName) {
+    const feature = await this.models.feature.get_unchecked(name);
+
+    const count = await this.db.workspaceFeature.count({
+      where: {
+        workspaceId,
+        featureId: feature.id,
+        activated: true,
+      },
+    });
+
+    return count > 0;
+  }
+
+  async list(workspaceId: string) {
+    const workspaceFeatures = await this.db.workspaceFeature.findMany({
+      include: {
+        feature: true,
+      },
+      where: {
+        workspaceId,
+        activated: true,
+      },
+    });
+
+    return workspaceFeatures.map(
+      workspaceFeature => workspaceFeature.feature.feature
+    ) as WorkspaceFeatureName[];
+  }
+
+  @Transactional()
+  async add<T extends WorkspaceFeatureName>(
+    workspaceId: string,
+    featureName: T,
+    reason: string,
+    overrides?: Partial<FeatureConfigs<T>>
+  ) {
+    const feature = await this.models.feature.get_unchecked(featureName);
+
+    const existing = await this.tx.workspaceFeature.findFirst({
+      where: {
+        workspaceId,
+        featureId: feature.id,
+        activated: true,
+      },
+    });
+
+    if (existing && !overrides) {
+      return existing;
+    }
+
+    const configs = {
+      ...(existing?.configs as {}),
+      ...overrides,
+    };
+
+    const parseResult = this.models.feature
+      .getConfigShape(featureName)
+      .partial()
+      .safeParse(configs);
+
+    if (!parseResult.success) {
+      throw new Error(`Invalid feature config for ${featureName}`, {
+        cause: parseResult.error,
+      });
+    }
+
+    let workspaceFeature;
+    if (existing) {
+      workspaceFeature = await this.tx.workspaceFeature.update({
+        where: {
+          id: existing.id,
+        },
+        data: {
+          configs: parseResult.data,
+          reason,
+        },
+      });
+    } else {
+      workspaceFeature = await this.tx.workspaceFeature.create({
+        data: {
+          workspaceId,
+          featureId: feature.id,
+          activated: true,
+          reason,
+          configs: parseResult.data,
+        },
+      });
+    }
+
+    this.logger.verbose(
+      `Feature ${featureName} added to workspace ${workspaceId}`
+    );
+
+    return workspaceFeature;
+  }
+
+  async remove(workspaceId: string, featureName: WorkspaceFeatureName) {
+    const feature = await this.models.feature.get_unchecked(featureName);
+
+    await this.tx.workspaceFeature.deleteMany({
+      where: {
+        workspaceId,
+        featureId: feature.id,
+      },
+    });
+
+    this.logger.verbose(
+      `Feature ${featureName} removed from workspace ${workspaceId}`
+    );
+  }
+
+  @Transactional()
+  async switch<T extends WorkspaceFeatureName>(
+    workspaceId: string,
+    from: WorkspaceFeatureName,
+    to: T,
+    reason: string,
+    overrides?: Partial<FeatureConfigs<T>>
+  ) {
+    await this.remove(workspaceId, from);
+    return await this.add(workspaceId, to, reason, overrides);
+  }
+}