feat: user usage init (#5074)

packages/backend/server/migrations/20231030104009_user_features/migration.sql

@@ -0,0 +1,45 @@
+/*
+  Warnings:
+
+  - You are about to drop the `user_feature_gates` table. If the table is not empty, all the data it contains will be lost.
+
+*/
+-- DropForeignKey
+ALTER TABLE "user_feature_gates" DROP CONSTRAINT "user_feature_gates_user_id_fkey";
+
+-- DropTable
+DROP TABLE "user_feature_gates";
+
+-- CreateTable
+CREATE TABLE "user_features" (
+    "id" SERIAL NOT NULL,
+    "user_id" VARCHAR(36) NOT NULL,
+    "feature_id" INTEGER NOT NULL,
+    "reason" VARCHAR NOT NULL,
+    "created_at" TIMESTAMPTZ(6) NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    "expired_at" TIMESTAMPTZ(6),
+    "activated" BOOLEAN NOT NULL DEFAULT false,
+
+    CONSTRAINT "user_features_pkey" PRIMARY KEY ("id")
+);
+
+-- CreateTable
+CREATE TABLE "features" (
+    "id" SERIAL NOT NULL,
+    "feature" VARCHAR NOT NULL,
+    "version" INTEGER NOT NULL DEFAULT 0,
+    "type" INTEGER NOT NULL,
+    "configs" JSON NOT NULL,
+    "created_at" TIMESTAMPTZ(6) NOT NULL DEFAULT CURRENT_TIMESTAMP,
+
+    CONSTRAINT "features_pkey" PRIMARY KEY ("id")
+);
+
+-- CreateIndex
+CREATE UNIQUE INDEX "features_feature_version_key" ON "features"("feature", "version");
+
+-- AddForeignKey
+ALTER TABLE "user_features" ADD CONSTRAINT "user_features_user_id_fkey" FOREIGN KEY ("user_id") REFERENCES "users"("id") ON DELETE CASCADE ON UPDATE CASCADE;
+
+-- AddForeignKey
+ALTER TABLE "user_features" ADD CONSTRAINT "user_features_feature_id_fkey" FOREIGN KEY ("feature_id") REFERENCES "features"("id") ON DELETE CASCADE ON UPDATE CASCADE;

packages/backend/server/schema.prisma

@@ -22,7 +22,7 @@ model User {
 
   accounts             Account[]
   sessions             Session[]
-  features             UserFeatureGates[]
+  features             UserFeatures[]
   customer             UserStripeCustomer?
   subscription         UserSubscription?
   invoices             UserInvoice[]
@@ -113,15 +113,48 @@ model WorkspacePageUserPermission {
   @@map("workspace_page_user_permissions")
 }
 
-model UserFeatureGates {
-  id        String   @id @default(uuid()) @db.VarChar
-  userId    String   @map("user_id") @db.VarChar
-  feature   String   @db.VarChar
-  reason    String   @db.VarChar
-  createdAt DateTime @default(now()) @map("created_at") @db.Timestamptz(6)
-  user      User     @relation(fields: [userId], references: [id], onDelete: Cascade)
+// feature gates is a way to enable/disable features for a user
+// for example:
+// - early access is a feature that allow some users to access the insider version
+// - pro plan is a quota that allow some users access to more resources after they pay
+model UserFeatures {
+  id        Int    @id @default(autoincrement())
+  userId    String @map("user_id") @db.VarChar(36)
+  featureId Int    @map("feature_id") @db.Integer
 
-  @@map("user_feature_gates")
+  // we will record the reason why the feature is enabled/disabled
+  // for example:
+  // - pro_plan_v1: "user buy the pro plan"
+  reason    String    @db.VarChar
+  // record the quota enabled time
+  createdAt DateTime  @default(now()) @map("created_at") @db.Timestamptz(6)
+  // record the quota expired time, pay plan is a subscription, so it will expired
+  expiredAt DateTime? @map("expired_at") @db.Timestamptz(6)
+  // whether the feature is activated
+  // for example:
+  // - if we switch the user to another plan, we will set the old plan to deactivated, but dont delete it
+  activated Boolean   @default(false)
+
+  user    User     @relation(fields: [userId], references: [id], onDelete: Cascade)
+  feature Features @relation(fields: [featureId], references: [id], onDelete: Cascade)
+
+  @@map("user_features")
+}
+
+model Features {
+  id        Int      @id @default(autoincrement())
+  feature   String   @db.VarChar
+  version   Int      @default(0) @db.Integer
+  // 0: feature, 1: quota
+  type      Int      @db.Integer
+  // configs, define by feature conntroller
+  configs   Json     @db.Json
+  createdAt DateTime @default(now()) @map("created_at") @db.Timestamptz(6)
+
+  UserFeatureGates UserFeatures[]
+
+  @@unique([feature, version])
+  @@map("features")
 }
 
 model Account {

packages/backend/server/src/data/commands/run.ts

@@ -14,7 +14,7 @@ interface Migration {
   down: (db: PrismaService) => Promise<void>;
 }
 
-async function collectMigrations(): Promise<Migration[]> {
+export async function collectMigrations(): Promise<Migration[]> {
   const folder = join(fileURLToPath(import.meta.url), '../../migrations');
 
   const migrationFiles = readdirSync(folder)

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

@@ -0,0 +1,93 @@
+import {
+  FeatureKind,
+  Features,
+  FeatureType,
+  upsertFeature,
+} from '../../modules/features';
+import { Quotas } from '../../modules/quota';
+import { PrismaService } from '../../prisma';
+
+export class UserFeaturesInit1698652531198 {
+  // do the migration
+  static async up(db: PrismaService) {
+    // upgrade features from lower version to higher version
+    for (const feature of Features) {
+      await upsertFeature(db, feature);
+    }
+    await migrateNewFeatureTable(db);
+
+    for (const quota of Quotas) {
+      await upsertFeature(db, quota);
+    }
+  }
+
+  // revert the migration
+  static async down(_db: PrismaService) {
+    // TODO: revert the migration
+  }
+}
+
+async function migrateNewFeatureTable(prisma: PrismaService) {
+  const waitingList = await prisma.newFeaturesWaitingList.findMany();
+  for (const oldUser of waitingList) {
+    const user = await prisma.user.findFirst({
+      where: {
+        email: oldUser.email,
+      },
+    });
+    if (user) {
+      const hasEarlyAccess = await prisma.userFeatures.count({
+        where: {
+          userId: user.id,
+          feature: {
+            feature: FeatureType.EarlyAccess,
+          },
+          activated: true,
+        },
+      });
+      if (hasEarlyAccess === 0) {
+        await prisma.$transaction(async tx => {
+          const latestFlag = await tx.userFeatures.findFirst({
+            where: {
+              userId: user.id,
+              feature: {
+                feature: FeatureType.EarlyAccess,
+                type: FeatureKind.Feature,
+              },
+              activated: true,
+            },
+            orderBy: {
+              createdAt: 'desc',
+            },
+          });
+          if (latestFlag) {
+            return latestFlag.id;
+          } else {
+            return tx.userFeatures
+              .create({
+                data: {
+                  reason: 'Early access user',
+                  activated: true,
+                  user: {
+                    connect: {
+                      id: user.id,
+                    },
+                  },
+                  feature: {
+                    connect: {
+                      feature_version: {
+                        feature: FeatureType.EarlyAccess,
+                        version: 1,
+                      },
+                      type: FeatureKind.Feature,
+                    },
+                  },
+                },
+              })
+              .then(r => r.id);
+          }
+        });
+      }
+    }
+  }
+}

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: {