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>;
+    };
+  }
+
+  /**
+   * 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>;
   }
 
-  private getConfigShape(name: FeatureNames): z.ZodObject<any> {
+  getConfigShape(name: FeatureName): z.ZodObject<any> {
     return Features[name]?.shape.configs ?? z.object({});
   }
 }