From 16871848911e11ddf6949cb58afd93438ac8e380 Mon Sep 17 00:00:00 2001 From: forehalo Date: Thu, 16 Jan 2025 09:36:16 +0000 Subject: [PATCH] feat(server): feature model (#9709) close CLOUD-101 --- .../models/__snapshots__/feature.spec.ts.md | 31 +++++ .../models/__snapshots__/feature.spec.ts.snap | Bin 0 -> 330 bytes .../src/__tests__/models/feature.spec.ts | 121 ++++++++++++++++++ .../server/src/models/common/feature.ts | 57 +++++++++ .../backend/server/src/models/common/index.ts | 1 + packages/backend/server/src/models/feature.ts | 103 +++++++++++++++ packages/backend/server/src/models/index.ts | 20 ++- 7 files changed, 327 insertions(+), 6 deletions(-) create mode 100644 packages/backend/server/src/__tests__/models/__snapshots__/feature.spec.ts.md create mode 100644 packages/backend/server/src/__tests__/models/__snapshots__/feature.spec.ts.snap create mode 100644 packages/backend/server/src/__tests__/models/feature.spec.ts create mode 100644 packages/backend/server/src/models/common/feature.ts create mode 100644 packages/backend/server/src/models/common/index.ts create mode 100644 packages/backend/server/src/models/feature.ts diff --git a/packages/backend/server/src/__tests__/models/__snapshots__/feature.spec.ts.md b/packages/backend/server/src/__tests__/models/__snapshots__/feature.spec.ts.md new file mode 100644 index 0000000000..aa90298228 --- /dev/null +++ b/packages/backend/server/src/__tests__/models/__snapshots__/feature.spec.ts.md @@ -0,0 +1,31 @@ +# Snapshot report for `src/__tests__/models/feature.spec.ts` + +The actual snapshot is saved in `feature.spec.ts.snap`. + +Generated by [AVA](https://avajs.dev). + +## should get feature + +> Snapshot 1 + + { + blobLimit: 10485760, + copilotActionLimit: 10, + historyPeriod: 604800000, + memberLimit: 3, + name: 'Free', + storageQuota: 10737418240, + } + +## should get feature if extra fields exist in feature config + +> Snapshot 1 + + { + blobLimit: 10485760, + copilotActionLimit: 10, + historyPeriod: 604800000, + memberLimit: 3, + name: 'Free', + storageQuota: 10737418240, + } diff --git a/packages/backend/server/src/__tests__/models/__snapshots__/feature.spec.ts.snap b/packages/backend/server/src/__tests__/models/__snapshots__/feature.spec.ts.snap new file mode 100644 index 0000000000000000000000000000000000000000..296805ddc2c193f185ed1d1d4522222287771ae8 GIT binary patch literal 330 zcmV-Q0k!@?RzVycvQ1=+XDF6rH3BW6WHvlaHM${Qn=RCXST!~ieU@!pnK^@Qz zM%3D)1k@hVX6Af$CY9)?lBU2j8q99)r{cUG; + +test.before(async t => { + const module = await createTestingModule({ + providers: [FeatureModel], + }); + + t.context.feature = module.get(FeatureModel); + t.context.module = module; +}); + +test.beforeEach(async t => { + await initTestingDB(t.context.module.get(PrismaClient)); +}); + +test.after(async t => { + await t.context.module.close(); +}); + +test('should get feature', async t => { + const { feature } = t.context; + const freePlanFeature = await feature.get('free_plan_v1'); + + t.snapshot(freePlanFeature.configs); +}); + +test('should throw if feature not found', async t => { + const { feature } = t.context; + await t.throwsAsync(feature.get('not_found_feature' as any), { + message: 'Feature not_found_feature not found', + }); +}); + +test('should throw if feature config in invalid', async t => { + const { feature } = t.context; + const freePlanFeature = await feature.get('free_plan_v1'); + + // @ts-expect-error internal + await feature.db.feature.update({ + where: { + id: freePlanFeature.id, + }, + data: { + configs: { + ...freePlanFeature.configs, + memberLimit: 'invalid' as any, + }, + }, + }); + + await t.throwsAsync(feature.get('free_plan_v1'), { + message: 'Invalid feature config for free_plan_v1', + }); +}); + +// NOTE(@forehalo): backward compatibility +// new version of feature config may introduce new field +// this test means to ensure that the older version of AFFiNE Server can still read it +test('should get feature if extra fields exist in feature config', async t => { + const { feature } = t.context; + const freePlanFeature = await feature.get('free_plan_v1'); + + // @ts-expect-error internal + await feature.db.feature.update({ + where: { + id: freePlanFeature.id, + }, + data: { + configs: { + ...freePlanFeature.configs, + extraField: 'extraValue', + }, + }, + }); + + const freePlanFeature2 = await feature.get('free_plan_v1'); + + t.snapshot(freePlanFeature2.configs); +}); + +test('should create feature', async t => { + const { feature } = t.context; + + const newFeature = await feature.upsert('new_feature' as any, {}); + + t.deepEqual(newFeature.configs, {}); +}); + +test('should update feature', async t => { + const { feature } = t.context; + const freePlanFeature = await feature.get('free_plan_v1'); + + const newFreePlanFeature = await feature.upsert('free_plan_v1', { + ...freePlanFeature.configs, + memberLimit: 10, + }); + + t.deepEqual(newFreePlanFeature.configs, { + ...freePlanFeature.configs, + memberLimit: 10, + }); +}); + +test('should throw if feature config is invalid when updating', async t => { + const { feature } = t.context; + await t.throwsAsync(feature.upsert('free_plan_v1', {} as any), { + message: 'Invalid feature config for free_plan_v1', + }); +}); diff --git a/packages/backend/server/src/models/common/feature.ts b/packages/backend/server/src/models/common/feature.ts new file mode 100644 index 0000000000..43626ff0b4 --- /dev/null +++ b/packages/backend/server/src/models/common/feature.ts @@ -0,0 +1,57 @@ +import { z } from 'zod'; + +export enum FeatureType { + Feature = 0, + Quota = 1, +} + +// TODO(@forehalo): quota is a useless extra concept, merge it with feature +export const UserPlanQuotaConfig = z.object({ + // quota name + name: z.string(), + // single blob limit + blobLimit: z.number(), + // total blob limit + storageQuota: z.number(), + // history period of validity + historyPeriod: z.number(), + // member limit + memberLimit: z.number(), + // copilot action limit 10 + copilotActionLimit: z.number(), +}); + +export const WorkspaceQuotaConfig = UserPlanQuotaConfig.extend({ + // seat quota + seatQuota: z.number(), +}).omit({ + copilotActionLimit: true, +}); + +function feature(configs: z.ZodObject) { + return z.object({ + type: z.literal(FeatureType.Feature), + configs: configs, + }); +} + +function quota(configs: z.ZodObject) { + return z.object({ + type: z.literal(FeatureType.Quota), + configs: configs, + }); +} + +export const Features = { + copilot: feature(z.object({})), + early_access: feature(z.object({ whitelist: z.array(z.string()) })), + unlimited_workspace: feature(z.object({})), + unlimited_copilot: feature(z.object({})), + ai_early_access: feature(z.object({})), + administrator: feature(z.object({})), + free_plan_v1: quota(UserPlanQuotaConfig), + pro_plan_v1: quota(UserPlanQuotaConfig), + lifetime_pro_plan_v1: quota(UserPlanQuotaConfig), + restricted_plan_v1: quota(UserPlanQuotaConfig), + team_plan_v1: quota(WorkspaceQuotaConfig), +}; diff --git a/packages/backend/server/src/models/common/index.ts b/packages/backend/server/src/models/common/index.ts new file mode 100644 index 0000000000..42a9f556aa --- /dev/null +++ b/packages/backend/server/src/models/common/index.ts @@ -0,0 +1 @@ +export * from './feature'; diff --git a/packages/backend/server/src/models/feature.ts b/packages/backend/server/src/models/feature.ts new file mode 100644 index 0000000000..c32b713177 --- /dev/null +++ b/packages/backend/server/src/models/feature.ts @@ -0,0 +1,103 @@ +import { Injectable, Logger } from '@nestjs/common'; +import { Feature, PrismaClient } from '@prisma/client'; +import { z } from 'zod'; + +import { PrismaTransaction } from '../base'; +import { Features, FeatureType } from './common'; + +type FeatureNames = keyof typeof Features; +type FeatureConfigs = z.infer< + (typeof Features)[T]['shape']['configs'] +>; + +// TODO(@forehalo): +// `version` column in `features` table will deprecated because it's makes the whole system complicated without any benefits. +// It was brought to introduce a version control for features, but the version controlling is not and will not actually needed. +// It even makes things harder when a new version of an existing feature is released. +// We have to manually update all the users and workspaces binding to the latest version, which are thousands of handreds. +// This is a huge burden for us and we should remove it. +@Injectable() +export class FeatureModel { + private readonly logger = new Logger(FeatureModel.name); + + constructor(private readonly db: PrismaClient) {} + + async get(name: T) { + const feature = await this.getLatest(this.db, 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, + }); + } + + return { + ...feature, + configs: parseResult.data as FeatureConfigs, + }; + } + + 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; + + // 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 feature = await this.db.$transaction(async tx => { + const latest = await this.getLatest(tx, name); + + if (!latest) { + return await tx.feature.create({ + data: { + type: FeatureType.Feature, + feature: name, + configs: parsedConfigs, + }, + }); + } else { + return await tx.feature.update({ + where: { id: latest.id }, + data: { + configs: parsedConfigs, + }, + }); + } + }); + + this.logger.verbose(`Feature ${name} upserted`); + + return feature as Feature & { configs: FeatureConfigs }; + } + + private async getLatest( + client: PrismaTransaction, + name: T + ) { + return client.feature.findFirst({ + where: { feature: name }, + orderBy: { version: 'desc' }, + }); + } + + private getConfigShape(name: FeatureNames): 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 0e7ec89f63..37ba5f8e6d 100644 --- a/packages/backend/server/src/models/index.ts +++ b/packages/backend/server/src/models/index.ts @@ -1,21 +1,24 @@ import { Global, Injectable, Module } from '@nestjs/common'; +import { FeatureModel } from './feature'; import { SessionModel } from './session'; import { UserModel } from './user'; import { VerificationTokenModel } from './verification-token'; -export * from './session'; -export * from './user'; -export * from './verification-token'; - -const models = [UserModel, SessionModel, VerificationTokenModel] as const; +const models = [ + UserModel, + SessionModel, + VerificationTokenModel, + FeatureModel, +] as const; @Injectable() export class Models { constructor( public readonly user: UserModel, public readonly session: SessionModel, - public readonly verificationToken: VerificationTokenModel + public readonly verificationToken: VerificationTokenModel, + public readonly feature: FeatureModel ) {} } @@ -25,3 +28,8 @@ export class Models { exports: [Models], }) export class ModelModules {} + +export * from './feature'; +export * from './session'; +export * from './user'; +export * from './verification-token';