feat: user usage init (#5074)

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

@@ -0,0 +1,177 @@
+import { Injectable } from '@nestjs/common';
+
+import { PrismaService } from '../../prisma';
+import { Feature, FeatureKind, FeatureType } from './types';
+
+@Injectable()
+export class FeatureService {
+  constructor(private readonly prisma: PrismaService) {}
+
+  async getFeaturesVersion() {
+    const features = await this.prisma.features.findMany({
+      where: {
+        type: FeatureKind.Feature,
+      },
+      select: {
+        feature: true,
+        version: true,
+      },
+    });
+    return features.reduce(
+      (acc, feature) => {
+        acc[feature.feature] = feature.version;
+        return acc;
+      },
+      {} as Record<string, number>
+    );
+  }
+
+  async getFeature(feature: FeatureType) {
+    return this.prisma.features.findFirst({
+      where: {
+        feature,
+        type: FeatureKind.Feature,
+      },
+      orderBy: {
+        version: 'desc',
+      },
+    });
+  }
+
+  async addUserFeature(
+    userId: string,
+    feature: FeatureType,
+    version: number,
+    reason: string,
+    expiredAt?: Date | string
+  ) {
+    return this.prisma.$transaction(async tx => {
+      const latestFlag = await tx.userFeatures.findFirst({
+        where: {
+          userId,
+          feature: {
+            feature,
+            type: FeatureKind.Feature,
+          },
+          activated: true,
+        },
+        orderBy: {
+          createdAt: 'desc',
+        },
+      });
+      if (latestFlag) {
+        return latestFlag.id;
+      } else {
+        return tx.userFeatures
+          .create({
+            data: {
+              reason,
+              expiredAt,
+              activated: true,
+              user: {
+                connect: {
+                  id: userId,
+                },
+              },
+              feature: {
+                connect: {
+                  feature_version: {
+                    feature,
+                    version,
+                  },
+                  type: FeatureKind.Feature,
+                },
+              },
+            },
+          })
+          .then(r => r.id);
+      }
+    });
+  }
+
+  async removeUserFeature(userId: string, feature: FeatureType) {
+    return this.prisma.userFeatures
+      .updateMany({
+        where: {
+          userId,
+          feature: {
+            feature,
+            type: FeatureKind.Feature,
+          },
+          activated: true,
+        },
+        data: {
+          activated: false,
+        },
+      })
+      .then(r => r.count);
+  }
+
+  async getUserFeatures(userId: string) {
+    const features = await this.prisma.userFeatures.findMany({
+      where: {
+        user: { id: userId },
+        feature: {
+          type: FeatureKind.Feature,
+        },
+      },
+      select: {
+        activated: true,
+        reason: true,
+        createdAt: true,
+        expiredAt: true,
+        feature: {
+          select: {
+            feature: true,
+            configs: true,
+          },
+        },
+      },
+    });
+    return features as typeof features &
+      {
+        feature: Pick<Feature, 'feature' | 'configs'>;
+      }[];
+  }
+
+  async listFeatureUsers(feature: FeatureType) {
+    return this.prisma.userFeatures
+      .findMany({
+        where: {
+          activated: true,
+          feature: {
+            feature: feature,
+            type: FeatureKind.Feature,
+          },
+        },
+        select: {
+          user: {
+            select: {
+              id: true,
+              name: true,
+              avatarUrl: true,
+              email: true,
+              emailVerified: true,
+              createdAt: true,
+            },
+          },
+        },
+      })
+      .then(users => users.map(user => user.user));
+  }
+
+  async hasFeature(userId: string, feature: FeatureType) {
+    return this.prisma.userFeatures
+      .count({
+        where: {
+          userId,
+          activated: true,
+          feature: {
+            feature,
+            type: FeatureKind.Feature,
+          },
+        },
+      })
+      .then(count => count > 0);
+  }
+}

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

@@ -0,0 +1,80 @@
+import { Injectable, Logger } from '@nestjs/common';
+
+import { Config } from '../../config';
+import { PrismaService } from '../../prisma';
+import { FeatureService } from './configure';
+import { FeatureType } from './types';
+
+export enum NewFeaturesKind {
+  EarlyAccess,
+}
+
+@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;
+    } else {
+      return true;
+    }
+  }
+}

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

@@ -0,0 +1,54 @@
+import { Module } from '@nestjs/common';
+
+import { PrismaService } from '../../prisma';
+import { FeatureService } from './configure';
+import { FeatureManagementService } from './feature';
+import type { CommonFeature } from './types';
+
+// upgrade features from lower version to higher version
+async function upsertFeature(
+  db: PrismaService,
+  feature: CommonFeature
+): Promise<void> {
+  const hasEqualOrGreaterVersion =
+    (await db.features.count({
+      where: {
+        feature: feature.feature,
+        version: {
+          gte: feature.version,
+        },
+      },
+    })) > 0;
+  // will not update exists version
+  if (!hasEqualOrGreaterVersion) {
+    await db.features.create({
+      data: {
+        feature: feature.feature,
+        type: feature.type,
+        version: feature.version,
+        configs: feature.configs,
+      },
+    });
+  }
+}
+
+/**
+ * Feature module provider pre-user feature flag management.
+ * includes:
+ * - feature query/update/permit
+ * - feature statistics
+ */
+@Module({
+  providers: [FeatureService, FeatureManagementService],
+  exports: [FeatureService, FeatureManagementService],
+})
+export class FeatureModule {}
+
+export type { CommonFeature, Feature } from './types';
+export { FeatureKind, Features, FeatureType } from './types';
+export {
+  FeatureManagementService,
+  FeatureService,
+  PrismaService,
+  upsertFeature,
+};

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

@@ -0,0 +1,33 @@
+import type { Prisma } from '@prisma/client';
+
+export enum FeatureKind {
+  Feature,
+  Quota,
+}
+
+export type CommonFeature = {
+  feature: string;
+  type: FeatureKind;
+  version: number;
+  configs: Prisma.InputJsonValue;
+};
+
+export type Feature = CommonFeature & {
+  type: FeatureKind.Feature;
+  feature: FeatureType;
+};
+
+export enum FeatureType {
+  EarlyAccess = 'early_access',
+}
+
+export const Features: Feature[] = [
+  {
+    feature: FeatureType.EarlyAccess,
+    type: FeatureKind.Feature,
+    version: 1,
+    configs: {
+      whitelist: ['@toeverything.info'],
+    },
+  },
+];

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

@@ -0,0 +1,21 @@
+import { Module } from '@nestjs/common';
+
+import { PermissionService } from '../workspaces/permission';
+import { QuotaService } from './quota';
+import { QuotaManagementService } from './storage';
+
+/**
+ * Quota module provider pre-user quota management.
+ * includes:
+ * - quota query/update/permit
+ * - quota statistics
+ */
+@Module({
+  providers: [PermissionService, QuotaService, QuotaManagementService],
+  exports: [QuotaService, QuotaManagementService],
+})
+export class QuotaModule {}
+
+export { QuotaManagementService, QuotaService };
+export { PrismaService } from '../../prisma';
+export { Quota_FreePlanV1, Quota_ProPlanV1, Quotas, QuotaType } from './types';

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

@@ -0,0 +1,140 @@
+import { Injectable } from '@nestjs/common';
+
+import { PrismaService } from '../../prisma';
+import { FeatureKind } from '../features';
+import { Quota, 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,
+        feature: {
+          select: {
+            feature: true,
+            configs: true,
+          },
+        },
+      },
+    });
+    return quota as typeof quota & {
+      feature: Pick<Quota, 'feature' | 'configs'>;
+    };
+  }
+
+  // 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/storage.ts

@@ -0,0 +1,64 @@
+import type { Storage } from '@affine/storage';
+import {
+  ForbiddenException,
+  Inject,
+  Injectable,
+  NotFoundException,
+} from '@nestjs/common';
+
+import { StorageProvide } from '../../storage';
+import { PermissionService } from '../workspaces/permission';
+import { QuotaService } from './quota';
+
+@Injectable()
+export class QuotaManagementService {
+  constructor(
+    private readonly quota: QuotaService,
+    private readonly permissions: PermissionService,
+    @Inject(StorageProvide) private readonly storage: Storage
+  ) {}
+
+  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;
+  }
+
+  // TODO: lazy calc, need to be optimized with cache
+  async getUserUsage(userId: string) {
+    const workspaces = await this.permissions.getOwnedWorkspaces(userId);
+    return this.storage.blobsSize(workspaces);
+  }
+
+  // get workspace's owner quota and total size of used
+  // quota was apply to owner's account
+  async getWorkspaceUsage(workspaceId: string) {
+    const { user: owner } =
+      await this.permissions.getWorkspaceOwner(workspaceId);
+    if (!owner) throw new NotFoundException('Workspace owner not found');
+    const { storageQuota } = (await this.getUserQuota(owner.id)) || {};
+    // get all workspaces size of owner used
+    const usageSize = await this.getUserUsage(owner.id);
+
+    return { quota: storageQuota, size: usageSize };
+  }
+
+  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

@@ -0,0 +1,52 @@
+import { CommonFeature, FeatureKind } from '../features';
+
+export enum QuotaType {
+  Quota_FreePlanV1 = 'free_plan_v1',
+  Quota_ProPlanV1 = 'pro_plan_v1',
+}
+
+export type Quota = CommonFeature & {
+  type: FeatureKind.Quota;
+  feature: QuotaType;
+  configs: {
+    blobLimit: number;
+    storageQuota: number;
+  };
+};
+
+export const Quotas: Quota[] = [
+  {
+    feature: QuotaType.Quota_FreePlanV1,
+    type: FeatureKind.Quota,
+    version: 1,
+    configs: {
+      // single blob limit 10MB
+      blobLimit: 10 * 1024 * 1024,
+      // total blob limit 10GB
+      storageQuota: 10 * 1024 * 1024 * 1024,
+    },
+  },
+  {
+    feature: QuotaType.Quota_ProPlanV1,
+    type: FeatureKind.Quota,
+    version: 1,
+    configs: {
+      // single blob limit 100MB
+      blobLimit: 100 * 1024 * 1024,
+      // total blob limit 100GB
+      storageQuota: 100 * 1024 * 1024 * 1024,
+    },
+  },
+];
+
+// ======== 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/users.ts

@@ -44,7 +44,10 @@ export class UsersService {
       })
       .then(user => user?.features.map(f => f.feature) ?? []);
 
-    return getStorageQuota(features) || this.config.objectStorage.quota;
+    return (
+      getStorageQuota(features.map(f => f.feature)) ||
+      this.config.objectStorage.quota
+    );
   }
 
   async findUserByEmail(email: string) {

packages/backend/server/src/modules/workspaces/permission.ts

@@ -26,6 +26,18 @@ export class PermissionService {
     return data?.type as Permission;
   }
 
+  async getOwnedWorkspaces(userId: string) {
+    return this.prisma.workspaceUserPermission
+      .findMany({
+        where: {
+          userId,
+          accepted: true,
+          type: Permission.Owner,
+        },
+      })
+      .then(data => data.map(({ workspaceId }) => workspaceId));
+  }
+
   async getWorkspaceOwner(workspaceId: string) {
     return this.prisma.workspaceUserPermission.findFirstOrThrow({
       where: {