fix: self hosted register

toeverything · Apr 16, 2024 · de433cc · de433cc
1 parent 0ca8a23
commit de433cc
Show file tree

Hide file tree

Showing 4 changed files with 59 additions and 7 deletions.
diff --git a/packages/backend/server/src/core/quota/schema.ts b/packages/backend/server/src/core/quota/schema.ts
@@ -117,6 +117,23 @@ export const Quotas: Quota[] = [
       copilotActionLimit: 10,
     },
   },
+  {
+    feature: QuotaType.UnlimitedPlanV1,
+    type: FeatureKind.Quota,
+    version: 1,
+    configs: {
+      // quota name
+      name: 'Unlimited',
+      // single blob limit 10MB
+      blobLimit: 100 * OneMB,
+      // total blob limit 10GB
+      storageQuota: 1000 * OneGB,
+      // history period of validity 7 days
+      historyPeriod: 7 * OneDay,
+      // member limit 3
+      memberLimit: 10,
+    },
+  },
 ];
 
 export const Quota_FreePlanV1_1 = {

diff --git a/packages/backend/server/src/core/quota/types.ts b/packages/backend/server/src/core/quota/types.ts
@@ -17,6 +17,7 @@ import { ByteUnit, OneDay, OneKB } from './constant';
 export enum QuotaType {
   FreePlanV1 = 'free_plan_v1',
   ProPlanV1 = 'pro_plan_v1',
+  UnlimitedPlanV1 = 'unlimited_plan_v1',
   // only for test, smaller quota
   RestrictedPlanV1 = 'restricted_plan_v1',
 }
@@ -26,6 +27,7 @@ const quotaPlan = z.object({
     QuotaType.FreePlanV1,
     QuotaType.ProPlanV1,
     QuotaType.RestrictedPlanV1,
+    QuotaType.UnlimitedPlanV1,
   ]),
   configs: z.object({
     name: z.string(),

diff --git a/packages/backend/server/src/data/migrations/1713258139569-refresh-new-plan.ts b/packages/backend/server/src/data/migrations/1713258139569-refresh-new-plan.ts
@@ -0,0 +1,27 @@
+import { ModuleRef } from '@nestjs/core';
+import { PrismaClient } from '@prisma/client';
+
+import { QuotaService, QuotaType } from '../../core/quota';
+import { UserService } from '../../core/user';
+import { Config } from '../../fundamentals';
+import { upgradeLatestQuotaVersion } from './utils/user-quotas';
+
+export class RefreshNewPlan1713258139569 {
+  // do the migration
+  static async up(db: PrismaClient, ref: ModuleRef) {
+    const config = ref.get(Config, { strict: false });
+    const { AFFINE_ADMIN_EMAIL } = process.env;
+    if (config.isSelfhosted && AFFINE_ADMIN_EMAIL) {
+      const user = ref.get(UserService, { strict: false });
+      const quota = ref.get(QuotaService, { strict: false });
+      const { id } = (await user.findUserByEmail(AFFINE_ADMIN_EMAIL)) || {};
+      if (id) {
+        await upgradeLatestQuotaVersion(db, QuotaType.UnlimitedPlanV1, '');
+        await quota.switchUserQuota(id, QuotaType.UnlimitedPlanV1);
+      }
+    }
+  }
+
+  // revert the migration
+  static async down(_db: PrismaClient) {}
+}
diff --git a/packages/backend/server/src/data/migrations/99999999-self-host-admin.ts b/packages/backend/server/src/data/migrations/99999999-self-host-admin.ts
@@ -1,6 +1,7 @@
 import { ModuleRef } from '@nestjs/core';
 import { PrismaClient } from '@prisma/client';
 
+import { QuotaService, QuotaType } from '../../core/quota';
 import { UserService } from '../../core/user';
 import { Config, CryptoHelper } from '../../fundamentals';
 
@@ -10,6 +11,7 @@ export class SelfHostAdmin99999999 {
     const config = ref.get(Config, { strict: false });
     const crypto = ref.get(CryptoHelper, { strict: false });
     const user = ref.get(UserService, { strict: false });
+    const quota = ref.get(QuotaService, { strict: false });
     if (config.isSelfhosted) {
       if (
         !process.env.AFFINE_ADMIN_EMAIL ||
@@ -19,13 +21,17 @@ export class SelfHostAdmin99999999 {
           'You have to set AFFINE_ADMIN_EMAIL and AFFINE_ADMIN_PASSWORD environment variables to generate the initial user for self-hosted AFFiNE Server.'
         );
       }
-      await user.findOrCreateUser(process.env.AFFINE_ADMIN_EMAIL, {
-        name: 'AFFINE First User',
-        emailVerifiedAt: new Date(),
-        password: await crypto.encryptPassword(
-          process.env.AFFINE_ADMIN_PASSWORD
-        ),
-      });
+      const { id } = await user.findOrCreateUser(
+        process.env.AFFINE_ADMIN_EMAIL,
+        {
+          name: 'AFFINE First User',
+          emailVerifiedAt: new Date(),
+          password: await crypto.encryptPassword(
+            process.env.AFFINE_ADMIN_PASSWORD
+          ),
+        }
+      );
+      await quota.switchUserQuota(id, QuotaType.UnlimitedPlanV1);
     }
   }