feat(server): feature model (#9709)

close CLOUD-101

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,
+    }

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

@@ -0,0 +1,121 @@
+import { TestingModule } from '@nestjs/testing';
+import { PrismaClient } from '@prisma/client';
+import ava, { TestFn } from 'ava';
+
+import { FeatureModel } from '../../models/feature';
+import { createTestingModule, initTestingDB } from '../utils';
+
+interface Context {
+  module: TestingModule;
+  feature: FeatureModel;
+}
+
+const test = ava as TestFn<Context>;
+
+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',
+  });
+});

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<T extends z.ZodRawShape>(configs: z.ZodObject<T>) {
+  return z.object({
+    type: z.literal(FeatureType.Feature),
+    configs: configs,
+  });
+}
+
+function quota<T extends z.ZodRawShape>(configs: z.ZodObject<T>) {
+  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),
+};

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

@@ -0,0 +1 @@
+export * from './feature';

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<T extends FeatureNames> = 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<T extends FeatureNames>(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<T>,
+    };
+  }
+
+  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;
+
+    // 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<T> };
+  }
+
+  private async getLatest<T extends FeatureNames>(
+    client: PrismaTransaction,
+    name: T
+  ) {
+    return client.feature.findFirst({
+      where: { feature: name },
+      orderBy: { version: 'desc' },
+    });
+  }
+
+  private getConfigShape(name: FeatureNames): z.ZodObject<any> {
+    return Features[name]?.shape.configs ?? z.object({});
+  }
+}

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';