From f8a515e89a40431a6c297504897f533f381d0fb5 Mon Sep 17 00:00:00 2001 From: forehalo Date: Wed, 22 Jan 2025 10:38:04 +0000 Subject: [PATCH] feat(server): user feature model (#9843) close CLOUD-108 --- .../src/__tests__/models/feature-user.spec.ts | 95 +++++++++++ .../models/feature-workspace.spec.ts | 127 ++++++++++++++ .../server/src/models/common/feature.ts | 21 +++ packages/backend/server/src/models/feature.ts | 94 ++++++----- packages/backend/server/src/models/index.ts | 6 + .../backend/server/src/models/user-feature.ts | 120 ++++++++++++++ .../server/src/models/workspace-feature.ts | 156 ++++++++++++++++++ 7 files changed, 580 insertions(+), 39 deletions(-) create mode 100644 packages/backend/server/src/__tests__/models/feature-user.spec.ts create mode 100644 packages/backend/server/src/__tests__/models/feature-workspace.spec.ts create mode 100644 packages/backend/server/src/models/user-feature.ts create mode 100644 packages/backend/server/src/models/workspace-feature.ts diff --git a/packages/backend/server/src/__tests__/models/feature-user.spec.ts b/packages/backend/server/src/__tests__/models/feature-user.spec.ts new file mode 100644 index 000000000..f6ebe26ee --- /dev/null +++ b/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; + +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')); +}); diff --git a/packages/backend/server/src/__tests__/models/feature-workspace.spec.ts b/packages/backend/server/src/__tests__/models/feature-workspace.spec.ts new file mode 100644 index 000000000..a3726ee14 --- /dev/null +++ b/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; + +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); +}); diff --git a/packages/backend/server/src/models/common/feature.ts b/packages/backend/server/src/models/common/feature.ts index 43626ff0b..055c1959c 100644 --- a/packages/backend/server/src/models/common/feature.ts +++ b/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 = z.infer< + (typeof Features)[T]['shape']['configs'] +>; diff --git a/packages/backend/server/src/models/feature.ts b/packages/backend/server/src/models/feature.ts index 6ec90cbbd..3fb9a1d2a 100644 --- a/packages/backend/server/src/models/feature.ts +++ b/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 = 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 = z.infer< // This is a huge burden for us and we should remove it. @Injectable() export class FeatureModel extends BaseModel { - async get(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(name: T) { + const feature = await this.get_unchecked(name); return { ...feature, - configs: parseResult.data as FeatureConfigs, + configs: this.check(name, feature.configs), }; } @Transactional() - async upsert(name: T, configs: FeatureConfigs) { - 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(name: T, configs: FeatureConfigs) { + 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 }; } - private async getLatest(name: T) { - return this.tx.feature.findFirst({ + /** + * Get the latest feature from database. + * + * @internal + */ + async try_get_unchecked(name: T) { + const feature = await this.tx.feature.findFirst({ where: { feature: name }, orderBy: { version: 'desc' }, }); + + return feature as Omit & { + configs: Record; + }; } - private getConfigShape(name: FeatureNames): z.ZodObject { + /** + * Get the latest feature from database. + * + * @throws {Error} If the feature is not found in DB. + * @internal + */ + async get_unchecked(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(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; + } + + getConfigShape(name: FeatureName): z.ZodObject { return Features[name]?.shape.configs ?? z.object({}); } } diff --git a/packages/backend/server/src/models/index.ts b/packages/backend/server/src/models/index.ts index bc0d91734..7e7ccb38b 100644 --- a/packages/backend/server/src/models/index.ts +++ b/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'; diff --git a/packages/backend/server/src/models/user-feature.ts b/packages/backend/server/src/models/user-feature.ts new file mode 100644 index 000000000..9308d2382 --- /dev/null +++ b/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(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); + } +} diff --git a/packages/backend/server/src/models/workspace-feature.ts b/packages/backend/server/src/models/workspace-feature.ts new file mode 100644 index 000000000..c509760f6 --- /dev/null +++ b/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(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( + workspaceId: string, + featureName: T, + reason: string, + overrides?: Partial> + ) { + 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( + workspaceId: string, + from: WorkspaceFeatureName, + to: T, + reason: string, + overrides?: Partial> + ) { + await this.remove(workspaceId, from); + return await this.add(workspaceId, to, reason, overrides); + } +}