feat: struct type feature config (#5142)

packages/backend/server/src/data/migrations/1698652531198-user-features-init.ts

@@ -1,10 +1,12 @@
+import { Prisma } from '@prisma/client';
+
 import {
   CommonFeature,
   FeatureKind,
   Features,
   FeatureType,
 } from '../../modules/features';
-import { Quotas } from '../../modules/quota/types';
+import { Quotas } from '../../modules/quota/schema';
 import { PrismaService } from '../../prisma';
 
 export class UserFeaturesInit1698652531198 {
@@ -48,7 +50,7 @@ async function upsertFeature(
         feature: feature.feature,
         type: feature.type,
         version: feature.version,
-        configs: feature.configs,
+        configs: feature.configs as Prisma.InputJsonValue,
       },
     });
   }

packages/backend/server/src/modules/doc/history.ts

@@ -240,8 +240,7 @@ export class DocHistoryManager {
     }
 
     const quota = await this.quota.getUserQuota(permission.userId);
-
-    return new Date(Date.now() + quota.feature.configs.historyPeriod);
+    return quota.feature.historyPeriodFromNow;
   }
 
   @Cron(CronExpression.EVERY_DAY_AT_MIDNIGHT /* everyday at 12am */)

packages/backend/server/src/modules/features/feature.ts

@@ -1,80 +1,78 @@
-import { Injectable, Logger } from '@nestjs/common';
-
-import { Config } from '../../config';
 import { PrismaService } from '../../prisma';
-import { FeatureService } from './configure';
-import { FeatureType } from './types';
+import { Feature, FeatureSchema, FeatureType } from './types';
 
-enum NewFeaturesKind {
-  EarlyAccess,
-}
+class FeatureConfig {
+  readonly config: Feature;
 
-@Injectable()
-export class FeatureManagementService {
-  protected logger = new Logger(FeatureManagementService.name);
-  constructor(
-    private readonly feature: FeatureService,
-    private readonly prisma: PrismaService,
-    private readonly config: Config
-  ) {}
-
-  isStaff(email: string) {
-    return email.endsWith('@toeverything.info');
-  }
-
-  async addEarlyAccess(userId: string) {
-    return this.feature.addUserFeature(
-      userId,
-      FeatureType.EarlyAccess,
-      1,
-      'Early access user'
-    );
-  }
-
-  async removeEarlyAccess(userId: string) {
-    return this.feature.removeUserFeature(userId, FeatureType.EarlyAccess);
-  }
-
-  async listEarlyAccess() {
-    return this.feature.listFeatureUsers(FeatureType.EarlyAccess);
-  }
-
-  /// check early access by email
-  async canEarlyAccess(email: string) {
-    if (
-      this.config.featureFlags.earlyAccessPreview &&
-      !email.endsWith('@toeverything.info')
-    ) {
-      const user = await this.prisma.user.findFirst({
-        where: {
-          email,
-        },
-      });
-      if (user) {
-        const canEarlyAccess = await this.feature
-          .hasFeature(user.id, FeatureType.EarlyAccess)
-          .catch(() => false);
-        if (canEarlyAccess) {
-          return true;
-        }
-
-        // TODO: Outdated, switch to feature gates
-        const oldCanEarlyAccess = await this.prisma.newFeaturesWaitingList
-          .findUnique({
-            where: { email, type: NewFeaturesKind.EarlyAccess },
-          })
-          .then(x => !!x)
-          .catch(() => false);
-        if (oldCanEarlyAccess) {
-          this.logger.warn(
-            `User ${email} has early access in old table but not in new table`
-          );
-        }
-        return oldCanEarlyAccess;
-      }
-      return false;
+  constructor(data: any) {
+    const config = FeatureSchema.safeParse(data);
+    if (config.success) {
+      this.config = config.data;
     } else {
-      return true;
+      throw new Error(`Invalid quota config: ${config.error.message}`);
     }
   }
+
+  /// feature name of quota
+  get name() {
+    return this.config.feature;
+  }
+}
+
+export class EarlyAccessFeatureConfig extends FeatureConfig {
+  constructor(data: any) {
+    super(data);
+
+    if (this.config.feature !== FeatureType.EarlyAccess) {
+      throw new Error('Invalid feature config: type is not EarlyAccess');
+    }
+  }
+
+  checkWhiteList(email: string) {
+    for (const domain in this.config.configs.whitelist) {
+      if (email.endsWith(domain)) {
+        return true;
+      }
+    }
+    return false;
+  }
+}
+
+const FeatureConfigMap = {
+  [FeatureType.EarlyAccess]: EarlyAccessFeatureConfig,
+};
+
+const FeatureCache = new Map<
+  number,
+  InstanceType<(typeof FeatureConfigMap)[FeatureType]>
+>();
+
+export async function getFeature(prisma: PrismaService, featureId: number) {
+  const cachedQuota = FeatureCache.get(featureId);
+
+  if (cachedQuota) {
+    return cachedQuota;
+  }
+
+  const feature = await prisma.features.findFirst({
+    where: {
+      id: featureId,
+    },
+  });
+  if (!feature) {
+    // this should unreachable
+    throw new Error(`Quota config ${featureId} not found`);
+  }
+  const ConfigClass = FeatureConfigMap[feature.feature as FeatureType];
+
+  if (!ConfigClass) {
+    throw new Error(`Feature config ${featureId} not found`);
+  }
+
+  const config = new ConfigClass(feature);
+  // we always edit quota config as a new quota config
+  // so we can cache it by featureId
+  FeatureCache.set(featureId, config);
+
+  return config;
 }

packages/backend/server/src/modules/features/index.ts

@@ -1,8 +1,8 @@
 import { Module } from '@nestjs/common';
 
 import { PrismaService } from '../../prisma';
-import { FeatureService } from './configure';
-import { FeatureManagementService } from './feature';
+import { FeatureManagementService } from './management';
+import { FeatureService } from './service';
 
 /**
  * Feature module provider pre-user feature flag management.
@@ -16,6 +16,6 @@ import { FeatureManagementService } from './feature';
 })
 export class FeatureModule {}
 
-export type { CommonFeature, Feature } from './types';
+export { type CommonFeature, commonFeatureSchema } from './types';
 export { FeatureKind, Features, FeatureType } from './types';
 export { FeatureManagementService, FeatureService, PrismaService };

packages/backend/server/src/modules/features/management.ts

@@ -0,0 +1,89 @@
+import { Injectable, Logger, OnModuleInit } from '@nestjs/common';
+
+import { Config } from '../../config';
+import { PrismaService } from '../../prisma';
+import { EarlyAccessFeatureConfig } from './feature';
+import { FeatureService } from './service';
+import { FeatureType } from './types';
+
+enum NewFeaturesKind {
+  EarlyAccess,
+}
+
+@Injectable()
+export class FeatureManagementService implements OnModuleInit {
+  protected logger = new Logger(FeatureManagementService.name);
+  private earlyAccessFeature?: EarlyAccessFeatureConfig;
+  constructor(
+    private readonly feature: FeatureService,
+    private readonly prisma: PrismaService,
+    private readonly config: Config
+  ) {}
+  async onModuleInit() {
+    this.earlyAccessFeature = await this.feature.getFeature(
+      FeatureType.EarlyAccess
+    );
+  }
+
+  // ======== Admin ========
+
+  // todo(@darkskygit): replace this with abac
+  isStaff(email: string) {
+    return this.earlyAccessFeature?.checkWhiteList(email) ?? false;
+  }
+
+  // ======== Early Access ========
+
+  async addEarlyAccess(userId: string) {
+    return this.feature.addUserFeature(
+      userId,
+      FeatureType.EarlyAccess,
+      1,
+      'Early access user'
+    );
+  }
+
+  async removeEarlyAccess(userId: string) {
+    return this.feature.removeUserFeature(userId, FeatureType.EarlyAccess);
+  }
+
+  async listEarlyAccess() {
+    return this.feature.listFeatureUsers(FeatureType.EarlyAccess);
+  }
+
+  /// check early access by email
+  async canEarlyAccess(email: string) {
+    if (this.config.featureFlags.earlyAccessPreview && !this.isStaff(email)) {
+      const user = await this.prisma.user.findFirst({
+        where: {
+          email,
+        },
+      });
+      if (user) {
+        const canEarlyAccess = await this.feature
+          .hasFeature(user.id, FeatureType.EarlyAccess)
+          .catch(() => false);
+        if (canEarlyAccess) {
+          return true;
+        }
+
+        // TODO: Outdated, switch to feature gates
+        const oldCanEarlyAccess = await this.prisma.newFeaturesWaitingList
+          .findUnique({
+            where: { email, type: NewFeaturesKind.EarlyAccess },
+          })
+          .then(x => !!x)
+          .catch(() => false);
+        if (oldCanEarlyAccess) {
+          this.logger.warn(
+            `User ${email} has early access in old table but not in new table`
+          );
+        }
+        return oldCanEarlyAccess;
+      }
+      return false;
+    } else {
+      return true;
+    }
+  }
+}

packages/backend/server/src/modules/features/service.ts

@@ -1,7 +1,9 @@
 import { Injectable } from '@nestjs/common';
 
 import { PrismaService } from '../../prisma';
-import { Feature, FeatureKind, FeatureType } from './types';
+import { UserType } from '../users/types';
+import { getFeature } from './feature';
+import { FeatureKind, FeatureType } from './types';
 
 @Injectable()
 export class FeatureService {
@@ -27,15 +29,20 @@ export class FeatureService {
   }
 
   async getFeature(feature: FeatureType) {
-    return this.prisma.features.findFirst({
+    const data = await this.prisma.features.findFirst({
       where: {
         feature,
         type: FeatureKind.Feature,
       },
+      select: { id: true },
       orderBy: {
         version: 'desc',
       },
     });
+    if (data) {
+      return getFeature(this.prisma, data.id);
+    }
+    return undefined;
   }
 
   async addUserFeature(
@@ -120,21 +127,21 @@ export class FeatureService {
         reason: true,
         createdAt: true,
         expiredAt: true,
-        feature: {
-          select: {
-            feature: true,
-            configs: true,
-          },
-        },
+        featureId: true,
       },
     });
-    return features as typeof features &
-      {
-        feature: Pick<Feature, 'feature' | 'configs'>;
-      }[];
+
+    const configs = await Promise.all(
+      features.map(async feature => ({
+        ...feature,
+        feature: await getFeature(this.prisma, feature.featureId),
+      }))
+    );
+
+    return configs.filter(feature => !!feature.feature);
   }
 
-  async listFeatureUsers(feature: FeatureType) {
+  async listFeatureUsers(feature: FeatureType): Promise<UserType[]> {
     return this.prisma.userFeatures
       .findMany({
         where: {

packages/backend/server/src/modules/features/types.ts

@@ -1,26 +1,48 @@
-import type { Prisma } from '@prisma/client';
+import { URL } from 'node:url';
+
+import { z } from 'zod';
+
+/// ======== common schema ========
 
 export enum FeatureKind {
   Feature,
   Quota,
 }
 
-export type CommonFeature = {
-  feature: string;
-  type: FeatureKind;
-  version: number;
-  configs: Prisma.InputJsonValue;
-};
+export const commonFeatureSchema = z.object({
+  feature: z.string(),
+  type: z.nativeEnum(FeatureKind),
+  version: z.number(),
+  configs: z.unknown(),
+});
 
-export type Feature = CommonFeature & {
-  type: FeatureKind.Feature;
-  feature: FeatureType;
-};
+export type CommonFeature = z.infer<typeof commonFeatureSchema>;
+
+/// ======== feature define ========
 
 export enum FeatureType {
   EarlyAccess = 'early_access',
 }
 
+function checkHostname(host: string) {
+  try {
+    return new URL(`https://${host}`).hostname === host;
+  } catch (_) {
+    return false;
+  }
+}
+
+const featureEarlyAccess = z.object({
+  feature: z.literal(FeatureType.EarlyAccess),
+  configs: z.object({
+    whitelist: z
+      .string()
+      .startsWith('@')
+      .refine(domain => checkHostname(domain.slice(1)))
+      .array(),
+  }),
+});
+
 export const Features: Feature[] = [
   {
     feature: FeatureType.EarlyAccess,
@@ -31,3 +53,13 @@ export const Features: Feature[] = [
     },
   },
 ];
+
+/// ======== schema infer ========
+
+export const FeatureSchema = commonFeatureSchema
+  .extend({
+    type: z.literal(FeatureKind.Feature),
+  })
+  .and(z.discriminatedUnion('feature', [featureEarlyAccess]));
+
+export type Feature = z.infer<typeof FeatureSchema>;

packages/backend/server/src/modules/payment/service.ts

@@ -480,9 +480,9 @@ export class SubscriptionService {
 
   private getPlanQuota(plan: SubscriptionPlan) {
     if (plan === SubscriptionPlan.Free) {
-      return QuotaType.Quota_FreePlanV1;
+      return QuotaType.FreePlanV1;
     } else if (plan === SubscriptionPlan.Pro) {
-      return QuotaType.Quota_ProPlanV1;
+      return QuotaType.ProPlanV1;
     } else {
       throw new Error(`Unknown plan: ${plan}`);
     }
@@ -520,7 +520,7 @@ export class SubscriptionService {
       }
     } else {
       // switch to free plan if subscription is canceled
-      await this.quota.switchUserQuota(user.id, QuotaType.Quota_FreePlanV1);
+      await this.quota.switchUserQuota(user.id, QuotaType.FreePlanV1);
     }
 
     const commonData = {

packages/backend/server/src/modules/quota/constant.ts

@@ -0,0 +1,5 @@
+export const OneKB = 1024;
+export const OneMB = OneKB * OneKB;
+export const OneGB = OneKB * OneMB;
+export const OneDay = 1000 * 60 * 60 * 24;
+export const ByteUnit = ['B', 'KB', 'MB', 'GB', 'TB', 'PB', 'EB', 'ZB', 'YB'];

packages/backend/server/src/modules/quota/index.ts

@@ -1,7 +1,7 @@
 import { Module } from '@nestjs/common';
 
 import { PermissionService } from '../workspaces/permission';
-import { QuotaService } from './quota';
+import { QuotaService } from './service';
 import { QuotaManagementService } from './storage';
 
 /**
@@ -17,4 +17,5 @@ import { QuotaManagementService } from './storage';
 export class QuotaModule {}
 
 export { QuotaManagementService, QuotaService };
-export { Quota_FreePlanV1, Quota_ProPlanV1, Quotas, QuotaType } from './types';
+export { Quota_FreePlanV1, Quota_ProPlanV1, Quotas } from './schema';
+export { QuotaType } from './types';

packages/backend/server/src/modules/quota/quota.ts

@@ -1,157 +1,81 @@
-import { Injectable } from '@nestjs/common';
-
 import { PrismaService } from '../../prisma';
-import { FeatureKind } from '../features';
-import {
-  formatDate,
-  formatSize,
-  getQuotaName,
-  Quota,
-  QuotaType,
-} from './types';
+import { formatDate, formatSize, Quota, QuotaSchema } from './types';
 
-@Injectable()
-export class QuotaService {
-  constructor(private readonly prisma: PrismaService) {}
+const QuotaCache = new Map<number, QuotaConfig>();
 
-  // get activated user quota
-  async getUserQuota(userId: string) {
-    const quota = await this.prisma.userFeatures.findFirst({
+export class QuotaConfig {
+  readonly config: Quota;
+
+  static async get(prisma: PrismaService, featureId: number) {
+    const cachedQuota = QuotaCache.get(featureId);
+
+    if (cachedQuota) {
+      return cachedQuota;
+    }
+
+    const quota = await prisma.features.findFirst({
       where: {
-        user: {
-          id: userId,
-        },
-        feature: {
-          type: FeatureKind.Quota,
-        },
-        activated: true,
-      },
-      select: {
-        reason: true,
-        createdAt: true,
-        expiredAt: true,
-        feature: {
-          select: {
-            feature: true,
-            configs: true,
-          },
-        },
+        id: featureId,
       },
     });
-    console.error(userId, quota);
-    return quota as typeof quota & {
-      feature: Pick<Quota, 'feature' | 'configs'>;
-    };
+
+    if (!quota) {
+      throw new Error(`Quota config ${featureId} not found`);
+    }
+
+    const config = new QuotaConfig(quota);
+    // we always edit quota config as a new quota config
+    // so we can cache it by featureId
+    QuotaCache.set(featureId, config);
+
+    return config;
   }
 
-  getHumanReadableQuota(feature: QuotaType, configs: Quota['configs']) {
+  private constructor(data: any) {
+    const config = QuotaSchema.safeParse(data);
+    if (config.success) {
+      this.config = config.data;
+    } else {
+      throw new Error(
+        `Invalid quota config: ${config.error.message}, ${JSON.stringify(
+          data
+        )})}`
+      );
+    }
+  }
+
+  /// feature name of quota
+  get name() {
+    return this.config.feature;
+  }
+
+  get blobLimit() {
+    return this.config.configs.blobLimit;
+  }
+
+  get storageQuota() {
+    return this.config.configs.storageQuota;
+  }
+
+  get historyPeriod() {
+    return this.config.configs.historyPeriod;
+  }
+
+  get historyPeriodFromNow() {
+    return new Date(Date.now() + this.historyPeriod);
+  }
+
+  get memberLimit() {
+    return this.config.configs.memberLimit;
+  }
+
+  get humanReadable() {
     return {
-      name: getQuotaName(feature),
-      blobLimit: formatSize(configs.blobLimit),
-      storageQuota: formatSize(configs.storageQuota),
-      historyPeriod: formatDate(configs.historyPeriod),
-      memberLimit: configs.memberLimit.toString(),
+      name: this.config.configs.name,
+      blobLimit: formatSize(this.blobLimit),
+      storageQuota: formatSize(this.storageQuota),
+      historyPeriod: formatDate(this.historyPeriod),
+      memberLimit: this.memberLimit.toString(),
     };
   }
-
-  // get all user quota records
-  async getUserQuotas(userId: string) {
-    const quotas = await this.prisma.userFeatures.findMany({
-      where: {
-        user: {
-          id: userId,
-        },
-        feature: {
-          type: FeatureKind.Quota,
-        },
-      },
-      select: {
-        activated: true,
-        reason: true,
-        createdAt: true,
-        expiredAt: true,
-        feature: {
-          select: {
-            feature: true,
-            configs: true,
-          },
-        },
-      },
-    });
-    return quotas as typeof quotas &
-      {
-        feature: Pick<Quota, 'feature' | 'configs'>;
-      }[];
-  }
-
-  // switch user to a new quota
-  // currently each user can only have one quota
-  async switchUserQuota(
-    userId: string,
-    quota: QuotaType,
-    reason?: string,
-    expiredAt?: Date
-  ) {
-    await this.prisma.$transaction(async tx => {
-      const latestFreePlan = await tx.features.aggregate({
-        where: {
-          feature: QuotaType.Quota_FreePlanV1,
-        },
-        _max: {
-          version: true,
-        },
-      });
-
-      // we will deactivate all exists quota for this user
-      await tx.userFeatures.updateMany({
-        where: {
-          id: undefined,
-          userId,
-          feature: {
-            type: FeatureKind.Quota,
-          },
-        },
-        data: {
-          activated: false,
-        },
-      });
-
-      await tx.userFeatures.create({
-        data: {
-          user: {
-            connect: {
-              id: userId,
-            },
-          },
-          feature: {
-            connect: {
-              feature_version: {
-                feature: quota,
-                version: latestFreePlan._max.version || 1,
-              },
-              type: FeatureKind.Quota,
-            },
-          },
-          reason: reason ?? 'switch quota',
-          activated: true,
-          expiredAt,
-        },
-      });
-    });
-  }
-
-  async hasQuota(userId: string, quota: QuotaType) {
-    return this.prisma.userFeatures
-      .count({
-        where: {
-          userId,
-          feature: {
-            feature: quota,
-            type: FeatureKind.Quota,
-          },
-          activated: true,
-        },
-      })
-      .then(count => count > 0);
-  }
 }

packages/backend/server/src/modules/quota/schema.ts

@@ -0,0 +1,50 @@
+import { FeatureKind } from '../features';
+import { OneDay, OneGB, OneMB } from './constant';
+import { Quota, QuotaType } from './types';
+
+export const Quotas: Quota[] = [
+  {
+    feature: QuotaType.FreePlanV1,
+    type: FeatureKind.Quota,
+    version: 1,
+    configs: {
+      // quota name
+      name: 'Free',
+      // single blob limit 10MB
+      blobLimit: 10 * OneMB,
+      // total blob limit 10GB
+      storageQuota: 10 * OneGB,
+      // history period of validity 7 days
+      historyPeriod: 7 * OneDay,
+      // member limit 3
+      memberLimit: 3,
+    },
+  },
+  {
+    feature: QuotaType.ProPlanV1,
+    type: FeatureKind.Quota,
+    version: 1,
+    configs: {
+      // quota name
+      name: 'Pro',
+      // single blob limit 100MB
+      blobLimit: 100 * OneMB,
+      // total blob limit 100GB
+      storageQuota: 100 * OneGB,
+      // history period of validity 30 days
+      historyPeriod: 30 * OneDay,
+      // member limit 10
+      memberLimit: 10,
+    },
+  },
+];
+
+export const Quota_FreePlanV1 = {
+  feature: Quotas[0].feature,
+  version: Quotas[0].version,
+};
+
+export const Quota_ProPlanV1 = {
+  feature: Quotas[1].feature,
+  version: Quotas[1].version,
+};

packages/backend/server/src/modules/quota/service.ts

@@ -0,0 +1,147 @@
+import { Injectable } from '@nestjs/common';
+
+import { PrismaService } from '../../prisma';
+import { FeatureKind } from '../features';
+import { QuotaConfig } from './quota';
+import { QuotaType } from './types';
+
+@Injectable()
+export class QuotaService {
+  constructor(private readonly prisma: PrismaService) {}
+
+  // get activated user quota
+  async getUserQuota(userId: string) {
+    const quota = await this.prisma.userFeatures.findFirst({
+      where: {
+        user: {
+          id: userId,
+        },
+        feature: {
+          type: FeatureKind.Quota,
+        },
+        activated: true,
+      },
+      select: {
+        reason: true,
+        createdAt: true,
+        expiredAt: true,
+        featureId: true,
+      },
+    });
+
+    if (!quota) {
+      // this should unreachable
+      throw new Error(`User ${userId} has no quota`);
+    }
+
+    const feature = await QuotaConfig.get(this.prisma, quota.featureId);
+    return { ...quota, feature };
+  }
+
+  // get user all quota records
+  async getUserQuotas(userId: string) {
+    const quotas = await this.prisma.userFeatures.findMany({
+      where: {
+        user: {
+          id: userId,
+        },
+        feature: {
+          type: FeatureKind.Quota,
+        },
+      },
+      select: {
+        activated: true,
+        reason: true,
+        createdAt: true,
+        expiredAt: true,
+        featureId: true,
+      },
+    });
+    const configs = await Promise.all(
+      quotas.map(async quota => {
+        try {
+          return {
+            ...quota,
+            feature: await QuotaConfig.get(this.prisma, quota.featureId),
+          };
+        } catch (_) {}
+        return null as unknown as typeof quota & {
+          feature: QuotaConfig;
+        };
+      })
+    );
+
+    return configs.filter(quota => !!quota);
+  }
+
+  // switch user to a new quota
+  // currently each user can only have one quota
+  async switchUserQuota(
+    userId: string,
+    quota: QuotaType,
+    reason?: string,
+    expiredAt?: Date
+  ) {
+    await this.prisma.$transaction(async tx => {
+      const latestFreePlan = await tx.features.aggregate({
+        where: {
+          feature: QuotaType.FreePlanV1,
+        },
+        _max: {
+          version: true,
+        },
+      });
+
+      // we will deactivate all exists quota for this user
+      await tx.userFeatures.updateMany({
+        where: {
+          id: undefined,
+          userId,
+          feature: {
+            type: FeatureKind.Quota,
+          },
+        },
+        data: {
+          activated: false,
+        },
+      });
+
+      await tx.userFeatures.create({
+        data: {
+          user: {
+            connect: {
+              id: userId,
+            },
+          },
+          feature: {
+            connect: {
+              feature_version: {
+                feature: quota,
+                version: latestFreePlan._max.version || 1,
+              },
+              type: FeatureKind.Quota,
+            },
+          },
+          reason: reason ?? 'switch quota',
+          activated: true,
+          expiredAt,
+        },
+      });
+    });
+  }
+
+  async hasQuota(userId: string, quota: QuotaType) {
+    return this.prisma.userFeatures
+      .count({
+        where: {
+          userId,
+          feature: {
+            feature: quota,
+            type: FeatureKind.Quota,
+          },
+          activated: true,
+        },
+      })
+      .then(count => count > 0);
+  }
+}

packages/backend/server/src/modules/quota/storage.ts

@@ -1,14 +1,9 @@
 import type { Storage } from '@affine/storage';
-import {
-  ForbiddenException,
-  Inject,
-  Injectable,
-  NotFoundException,
-} from '@nestjs/common';
+import { Inject, Injectable, NotFoundException } from '@nestjs/common';
 
 import { StorageProvide } from '../../storage';
 import { PermissionService } from '../workspaces/permission';
-import { QuotaService } from './quota';
+import { QuotaService } from './service';
 
 @Injectable()
 export class QuotaManagementService {
@@ -20,17 +15,15 @@ export class QuotaManagementService {
 
   async getUserQuota(userId: string) {
     const quota = await this.quota.getUserQuota(userId);
-    if (quota) {
-      return {
-        name: quota.feature.feature,
-        reason: quota.reason,
-        createAt: quota.createdAt,
-        expiredAt: quota.expiredAt,
-        blobLimit: quota.feature.configs.blobLimit,
-        storageQuota: quota.feature.configs.storageQuota,
-      };
-    }
-    return null;
+
+    return {
+      name: quota.feature.name,
+      reason: quota.reason,
+      createAt: quota.createdAt,
+      expiredAt: quota.expiredAt,
+      blobLimit: quota.feature.blobLimit,
+      storageQuota: quota.feature.storageQuota,
+    };
   }
 
   // TODO: lazy calc, need to be optimized with cache
@@ -45,7 +38,7 @@ export class QuotaManagementService {
     const { user: owner } =
       await this.permissions.getWorkspaceOwner(workspaceId);
     if (!owner) throw new NotFoundException('Workspace owner not found');
-    const { storageQuota } = (await this.getUserQuota(owner.id)) || {};
+    const { storageQuota } = await this.getUserQuota(owner.id);
     // get all workspaces size of owner used
     const usageSize = await this.getUserUsage(owner.id);
 
@@ -55,9 +48,6 @@ export class QuotaManagementService {
   async checkBlobQuota(workspaceId: string, size: number) {
     const { quota, size: usageSize } =
       await this.getWorkspaceUsage(workspaceId);
-    if (typeof quota !== 'number') {
-      throw new ForbiddenException(`user's quota not exists`);
-    }
 
     return quota - (size + usageSize);
   }

packages/backend/server/src/modules/quota/types.ts

@@ -1,31 +1,37 @@
-import { CommonFeature, FeatureKind } from '../features';
+import { z } from 'zod';
+
+import { commonFeatureSchema, FeatureKind } from '../features';
+import { ByteUnit, OneDay, OneKB } from './constant';
+
+/// ======== quota define ========
 
 export enum QuotaType {
-  Quota_FreePlanV1 = 'free_plan_v1',
-  Quota_ProPlanV1 = 'pro_plan_v1',
+  FreePlanV1 = 'free_plan_v1',
+  ProPlanV1 = 'pro_plan_v1',
 }
 
-export enum QuotaName {
-  free_plan_v1 = 'Free Plan',
-  pro_plan_v1 = 'Pro Plan',
-}
+const quotaPlan = z.object({
+  feature: z.enum([QuotaType.FreePlanV1, QuotaType.ProPlanV1]),
+  configs: z.object({
+    name: z.string(),
+    blobLimit: z.number().positive().int(),
+    storageQuota: z.number().positive().int(),
+    historyPeriod: z.number().positive().int(),
+    memberLimit: z.number().positive().int(),
+  }),
+});
 
-export type Quota = CommonFeature & {
-  type: FeatureKind.Quota;
-  feature: QuotaType;
-  configs: {
-    blobLimit: number;
-    storageQuota: number;
-    historyPeriod: number;
-    memberLimit: number;
-  };
-};
+/// ======== schema infer ========
 
-const OneKB = 1024;
-const OneMB = OneKB * OneKB;
-const OneGB = OneKB * OneMB;
+export const QuotaSchema = commonFeatureSchema
+  .extend({
+    type: z.literal(FeatureKind.Quota),
+  })
+  .and(z.discriminatedUnion('feature', [quotaPlan]));
 
-const sizes = ['B', 'KB', 'MB', 'GB', 'TB', 'PB', 'EB', 'ZB', 'YB'];
+export type Quota = z.infer<typeof QuotaSchema>;
+
+/// ======== utils ========
 
 export function formatSize(bytes: number, decimals: number = 2): string {
   if (bytes === 0) return '0 B';
@@ -34,60 +40,11 @@ export function formatSize(bytes: number, decimals: number = 2): string {
 
   const i = Math.floor(Math.log(bytes) / Math.log(OneKB));
 
-  return parseFloat((bytes / Math.pow(OneKB, i)).toFixed(dm)) + ' ' + sizes[i];
+  return (
+    parseFloat((bytes / Math.pow(OneKB, i)).toFixed(dm)) + ' ' + ByteUnit[i]
+  );
 }
 
-const OneDay = 1000 * 60 * 60 * 24;
-
 export function formatDate(ms: number): string {
   return `${(ms / OneDay).toFixed(0)} days`;
 }
-
-export function getQuotaName(quota: QuotaType): string {
-  return QuotaName[quota];
-}
-
-export const Quotas: Quota[] = [
-  {
-    feature: QuotaType.Quota_FreePlanV1,
-    type: FeatureKind.Quota,
-    version: 1,
-    configs: {
-      // single blob limit 10MB
-      blobLimit: 10 * OneMB,
-      // total blob limit 10GB
-      storageQuota: 10 * OneGB,
-      // history period of validity 7 days
-      historyPeriod: 7 * OneDay,
-      // member limit 3
-      memberLimit: 3,
-    },
-  },
-  {
-    feature: QuotaType.Quota_ProPlanV1,
-    type: FeatureKind.Quota,
-    version: 1,
-    configs: {
-      // single blob limit 100MB
-      blobLimit: 100 * OneMB,
-      // total blob limit 100GB
-      storageQuota: 100 * OneGB,
-      // history period of validity 30 days
-      historyPeriod: 30 * OneDay,
-      // member limit 10
-      memberLimit: 10,
-    },
-  },
-];
-
-// ======== payload ========
-
-export const Quota_FreePlanV1 = {
-  feature: Quotas[0].feature,
-  version: Quotas[0].version,
-};
-
-export const Quota_ProPlanV1 = {
-  feature: Quotas[1].feature,
-  version: Quotas[1].version,
-};

packages/backend/server/src/modules/users/resolver.ts

@@ -114,18 +114,8 @@ export class UserResolver {
   @ResolveField(() => UserQuotaType, { name: 'quota', nullable: true })
   async getQuota(@CurrentUser() me: User) {
     const quota = await this.quota.getUserQuota(me.id);
-    const configs = quota.feature.configs;
 
-    return Object.assign(
-      {
-        name: quota.feature.feature,
-        humanReadable: this.quota.getHumanReadableQuota(
-          quota.feature.feature,
-          configs
-        ),
-      },
-      configs
-    );
+    return quota.feature;
   }
 
   @Throttle({ default: { limit: 10, ttl: 60 } })