feat(server): support ai plan (#6216)

packages/backend/server/migrations/20240319104623_user_subscriptions/migration.sql

@@ -0,0 +1,11 @@
+/*
+  Warnings:
+
+  - A unique constraint covering the columns `[user_id,plan]` on the table `user_subscriptions` will be added. If there are existing duplicate values, this will fail.
+
+*/
+-- DropIndex
+DROP INDEX "user_subscriptions_user_id_key";
+
+-- CreateIndex
+CREATE UNIQUE INDEX "user_subscriptions_user_id_plan_key" ON "user_subscriptions"("user_id", "plan");

packages/backend/server/schema.prisma

@@ -24,7 +24,7 @@ model User {
 
   features             UserFeatures[]
   customer             UserStripeCustomer?
-  subscription         UserSubscription?
+  subscriptions        UserSubscription[]
   invoices             UserInvoice[]
   workspacePermissions WorkspaceUserPermission[]
   pagePermissions      WorkspacePageUserPermission[]
@@ -369,7 +369,7 @@ model UserStripeCustomer {
 
 model UserSubscription {
   id                   Int       @id @default(autoincrement()) @db.Integer
-  userId               String    @unique @map("user_id") @db.VarChar(36)
+  userId               String    @map("user_id") @db.VarChar(36)
   plan                 String    @db.VarChar(20)
   // yearly/monthly
   recurring            String    @db.VarChar(20)
@@ -395,6 +395,7 @@ model UserSubscription {
   updatedAt DateTime @updatedAt @map("updated_at") @db.Timestamptz(6)
   user      User     @relation(fields: [userId], references: [id], onDelete: Cascade)
 
+  @@unique([userId, plan])
   @@map("user_subscriptions")
 }
 

packages/backend/server/src/plugins/payment/resolver.ts

@@ -190,7 +190,7 @@ export class SubscriptionResolver {
     }
 
     // extend it when new plans are added
-    const fixedPlans = [SubscriptionPlan.Pro];
+    const fixedPlans = [SubscriptionPlan.Pro, SubscriptionPlan.AI];
 
     return fixedPlans.reduce((prices, plan) => {
       const price = findPrice(plan);
@@ -242,17 +242,35 @@ export class SubscriptionResolver {
   @Mutation(() => UserSubscriptionType)
   async cancelSubscription(
     @CurrentUser() user: CurrentUser,
+    @Args({
+      name: 'plan',
+      type: () => SubscriptionPlan,
+      nullable: true,
+      defaultValue: SubscriptionPlan.Pro,
+    })
+    plan: SubscriptionPlan,
     @Args('idempotencyKey') idempotencyKey: string
   ) {
-    return this.service.cancelSubscription(idempotencyKey, user.id);
+    return this.service.cancelSubscription(idempotencyKey, user.id, plan);
   }
 
   @Mutation(() => UserSubscriptionType)
   async resumeSubscription(
     @CurrentUser() user: CurrentUser,
+    @Args({
+      name: 'plan',
+      type: () => SubscriptionPlan,
+      nullable: true,
+      defaultValue: SubscriptionPlan.Pro,
+    })
+    plan: SubscriptionPlan,
     @Args('idempotencyKey') idempotencyKey: string
   ) {
-    return this.service.resumeCanceledSubscription(idempotencyKey, user.id);
+    return this.service.resumeCanceledSubscription(
+      idempotencyKey,
+      user.id,
+      plan
+    );
   }
 
   @Mutation(() => UserSubscriptionType)
@@ -260,11 +278,19 @@ export class SubscriptionResolver {
     @CurrentUser() user: CurrentUser,
     @Args({ name: 'recurring', type: () => SubscriptionRecurring })
     recurring: SubscriptionRecurring,
+    @Args({
+      name: 'plan',
+      type: () => SubscriptionPlan,
+      nullable: true,
+      defaultValue: SubscriptionPlan.Pro,
+    })
+    plan: SubscriptionPlan,
     @Args('idempotencyKey') idempotencyKey: string
   ) {
     return this.service.updateSubscriptionRecurring(
       idempotencyKey,
       user.id,
+      plan,
       recurring
     );
   }
@@ -277,11 +303,21 @@ export class UserSubscriptionResolver {
     private readonly db: PrismaClient
   ) {}
 
-  @ResolveField(() => UserSubscriptionType, { nullable: true })
+  @ResolveField(() => UserSubscriptionType, {
+    nullable: true,
+    deprecationReason: 'use `UserType.subscriptions`',
+  })
   async subscription(
     @Context() ctx: { isAdminQuery: boolean },
     @CurrentUser() me: User,
-    @Parent() user: User
+    @Parent() user: User,
+    @Args({
+      name: 'plan',
+      type: () => SubscriptionPlan,
+      nullable: true,
+      defaultValue: SubscriptionPlan.Pro,
+    })
+    plan: SubscriptionPlan
   ) {
     // allow admin to query other user's subscription
     if (!ctx.isAdminQuery && me.id !== user.id) {
@@ -311,12 +347,33 @@ export class UserSubscriptionResolver {
 
     return this.db.userSubscription.findUnique({
       where: {
-        userId: user.id,
+        userId_plan: {
+          userId: user.id,
+          plan,
+        },
         status: SubscriptionStatus.Active,
       },
     });
   }
 
+  @ResolveField(() => [UserSubscriptionType])
+  async subscriptions(
+    @CurrentUser() me: User,
+    @Parent() user: User
+  ): Promise<UserSubscription[]> {
+    if (me.id !== user.id) {
+      throw new ForbiddenException(
+        'You are not allowed to access this subscription.'
+      );
+    }
+
+    return this.db.userSubscription.findMany({
+      where: {
+        userId: user.id,
+      },
+    });
+  }
+
   @ResolveField(() => [UserInvoiceType])
   async invoices(
     @CurrentUser() me: User,

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

@@ -1,4 +1,4 @@
-import { Injectable, Logger } from '@nestjs/common';
+import { BadRequestException, Injectable, Logger } from '@nestjs/common';
 import { OnEvent as RawOnEvent } from '@nestjs/event-emitter';
 import type {
   Prisma,
@@ -88,12 +88,15 @@ export class SubscriptionService {
     const currentSubscription = await this.db.userSubscription.findFirst({
       where: {
         userId: user.id,
+        plan,
         status: SubscriptionStatus.Active,
       },
     });
 
     if (currentSubscription) {
-      throw new Error('You already have a subscription');
+      throw new BadRequestException(
+        `You've already subscripted to the ${plan} plan`
+      );
     }
 
     const price = await this.getPrice(plan, recurring);
@@ -154,35 +157,47 @@ export class SubscriptionService {
 
   async cancelSubscription(
     idempotencyKey: string,
-    userId: string
+    userId: string,
+    plan: SubscriptionPlan
   ): Promise<UserSubscription> {
     const user = await this.db.user.findUnique({
       where: {
         id: userId,
       },
       include: {
-        subscription: true,
+        subscriptions: {
+          where: {
+            plan,
+          },
+        },
       },
     });
 
-    if (!user?.subscription) {
-      throw new Error('You do not have any subscription');
+    if (!user) {
+      throw new BadRequestException('Unknown user');
     }
 
-    if (user.subscription.canceledAt) {
-      throw new Error('Your subscription has already been canceled');
+    const subscriptionInDB = user?.subscriptions.find(s => s.plan === plan);
+    if (!subscriptionInDB) {
+      throw new BadRequestException(`You didn't subscript to the ${plan} plan`);
+    }
+
+    if (subscriptionInDB.canceledAt) {
+      throw new BadRequestException(
+        'Your subscription has already been canceled'
+      );
     }
 
     // should release the schedule first
-    if (user.subscription.stripeScheduleId) {
+    if (subscriptionInDB.stripeScheduleId) {
       const manager = await this.scheduleManager.fromSchedule(
-        user.subscription.stripeScheduleId
+        subscriptionInDB.stripeScheduleId
       );
       await manager.cancel(idempotencyKey);
       return this.saveSubscription(
         user,
         await this.stripe.subscriptions.retrieve(
-          user.subscription.stripeSubscriptionId
+          subscriptionInDB.stripeSubscriptionId
         ),
         false
       );
@@ -190,7 +205,7 @@ export class SubscriptionService {
       // let customer contact support if they want to cancel immediately
       // see https://stripe.com/docs/billing/subscriptions/cancel
       const subscription = await this.stripe.subscriptions.update(
-        user.subscription.stripeSubscriptionId,
+        subscriptionInDB.stripeSubscriptionId,
         { cancel_at_period_end: true },
         { idempotencyKey }
       );
@@ -200,44 +215,52 @@ export class SubscriptionService {
 
   async resumeCanceledSubscription(
     idempotencyKey: string,
-    userId: string
+    userId: string,
+    plan: SubscriptionPlan
   ): Promise<UserSubscription> {
     const user = await this.db.user.findUnique({
       where: {
         id: userId,
       },
       include: {
-        subscription: true,
+        subscriptions: true,
       },
     });
 
-    if (!user?.subscription) {
-      throw new Error('You do not have any subscription');
+    if (!user) {
+      throw new BadRequestException('Unknown user');
     }
 
-    if (!user.subscription.canceledAt) {
-      throw new Error('Your subscription has not been canceled');
+    const subscriptionInDB = user?.subscriptions.find(s => s.plan === plan);
+    if (!subscriptionInDB) {
+      throw new BadRequestException(`You didn't subscript to the ${plan} plan`);
     }
 
-    if (user.subscription.end < new Date()) {
-      throw new Error('Your subscription is expired, please checkout again.');
+    if (!subscriptionInDB.canceledAt) {
+      throw new BadRequestException('Your subscription has not been canceled');
     }
 
-    if (user.subscription.stripeScheduleId) {
+    if (subscriptionInDB.end < new Date()) {
+      throw new BadRequestException(
+        'Your subscription is expired, please checkout again.'
+      );
+    }
+
+    if (subscriptionInDB.stripeScheduleId) {
       const manager = await this.scheduleManager.fromSchedule(
-        user.subscription.stripeScheduleId
+        subscriptionInDB.stripeScheduleId
       );
       await manager.resume(idempotencyKey);
       return this.saveSubscription(
         user,
         await this.stripe.subscriptions.retrieve(
-          user.subscription.stripeSubscriptionId
+          subscriptionInDB.stripeSubscriptionId
         ),
         false
       );
     } else {
       const subscription = await this.stripe.subscriptions.update(
-        user.subscription.stripeSubscriptionId,
+        subscriptionInDB.stripeSubscriptionId,
         { cancel_at_period_end: false },
         { idempotencyKey }
       );
@@ -249,6 +272,7 @@ export class SubscriptionService {
   async updateSubscriptionRecurring(
     idempotencyKey: string,
     userId: string,
+    plan: SubscriptionPlan,
     recurring: SubscriptionRecurring
   ): Promise<UserSubscription> {
     const user = await this.db.user.findUnique({
@@ -256,30 +280,38 @@ export class SubscriptionService {
         id: userId,
       },
       include: {
-        subscription: true,
+        subscriptions: true,
       },
     });
 
-    if (!user?.subscription) {
-      throw new Error('You do not have any subscription');
+    if (!user) {
+      throw new BadRequestException('Unknown user');
+    }
+    const subscriptionInDB = user?.subscriptions.find(s => s.plan === plan);
+    if (!subscriptionInDB) {
+      throw new BadRequestException(`You didn't subscript to the ${plan} plan`);
     }
 
-    if (user.subscription.canceledAt) {
-      throw new Error('Your subscription has already been canceled ');
+    if (subscriptionInDB.canceledAt) {
+      throw new BadRequestException(
+        'Your subscription has already been canceled '
+      );
     }
 
-    if (user.subscription.recurring === recurring) {
-      throw new Error('You have already subscribed to this plan');
+    if (subscriptionInDB.recurring === recurring) {
+      throw new BadRequestException(
+        `You are already in ${recurring} recurring`
+      );
     }
 
     const price = await this.getPrice(
-      user.subscription.plan as SubscriptionPlan,
+      subscriptionInDB.plan as SubscriptionPlan,
       recurring
     );
 
     const manager = await this.scheduleManager.fromSubscription(
       `${idempotencyKey}-fromSubscription`,
-      user.subscription.stripeSubscriptionId
+      subscriptionInDB.stripeSubscriptionId
     );
 
     await manager.update(
@@ -295,7 +327,7 @@ export class SubscriptionService {
 
     return await this.db.userSubscription.update({
       where: {
-        id: user.subscription.id,
+        id: subscriptionInDB.id,
       },
       data: {
         stripeScheduleId: manager.schedule?.id ?? null, // update schedule id or set to null(undefined means untouched)
@@ -312,7 +344,7 @@ export class SubscriptionService {
     });
 
     if (!user) {
-      throw new Error('Unknown user');
+      throw new BadRequestException('Unknown user');
     }
 
     try {
@@ -323,7 +355,7 @@ export class SubscriptionService {
       return portal.url;
     } catch (e) {
       this.logger.error('Failed to create customer portal.', e);
-      throw new Error('Failed to create customer portal');
+      throw new BadRequestException('Failed to create customer portal');
     }
   }
 
@@ -520,7 +552,10 @@ export class SubscriptionService {
 
     const currentSubscription = await this.db.userSubscription.findUnique({
       where: {
-        userId: user.id,
+        userId_plan: {
+          userId: user.id,
+          plan,
+        },
       },
     });
 
@@ -643,8 +678,8 @@ export class SubscriptionService {
     });
 
     if (!prices.data.length) {
-      throw new Error(
-        `Unknown subscription plan ${plan} with recurring ${recurring}`
+      throw new BadRequestException(
+        `Unknown subscription plan ${plan} with ${recurring} recurring`
       );
     }
 

packages/backend/server/src/plugins/payment/types.ts

@@ -20,6 +20,7 @@ export enum SubscriptionRecurring {
 export enum SubscriptionPlan {
   Free = 'free',
   Pro = 'pro',
+  AI = 'ai',
   Team = 'team',
   Enterprise = 'enterprise',
   SelfHosted = 'selfhosted',

packages/backend/server/src/schema.gql

@@ -110,7 +110,7 @@ type Mutation {
   acceptInviteById(inviteId: String!, sendAcceptMail: Boolean, workspaceId: String!): Boolean!
   addToEarlyAccess(email: String!): Int!
   addWorkspaceFeature(feature: FeatureType!, workspaceId: String!): Int!
-  cancelSubscription(idempotencyKey: String!): UserSubscription!
+  cancelSubscription(idempotencyKey: String!, plan: SubscriptionPlan = Pro): UserSubscription!
   changeEmail(email: String!, token: String!): UserType!
   changePassword(newPassword: String!, token: String!): UserType!
 
@@ -134,7 +134,7 @@ type Mutation {
   removeAvatar: RemoveAvatar!
   removeEarlyAccess(email: String!): Int!
   removeWorkspaceFeature(feature: FeatureType!, workspaceId: String!): Int!
-  resumeSubscription(idempotencyKey: String!): UserSubscription!
+  resumeSubscription(idempotencyKey: String!, plan: SubscriptionPlan = Pro): UserSubscription!
   revoke(userId: String!, workspaceId: String!): Boolean!
   revokePage(pageId: String!, workspaceId: String!): Boolean! @deprecated(reason: "use revokePublicPage")
   revokePublicPage(pageId: String!, workspaceId: String!): WorkspacePage!
@@ -149,7 +149,7 @@ type Mutation {
   signIn(email: String!, password: String!): UserType!
   signUp(email: String!, name: String!, password: String!): UserType!
   updateProfile(input: UpdateUserInput!): UserType!
-  updateSubscriptionRecurring(idempotencyKey: String!, recurring: SubscriptionRecurring!): UserSubscription!
+  updateSubscriptionRecurring(idempotencyKey: String!, plan: SubscriptionPlan = Pro, recurring: SubscriptionRecurring!): UserSubscription!
 
   """Update workspace"""
   updateWorkspace(input: UpdateWorkspaceInput!): WorkspaceType!
@@ -264,6 +264,7 @@ enum ServerFeature {
 }
 
 enum SubscriptionPlan {
+  AI
   Enterprise
   Free
   Pro
@@ -385,7 +386,8 @@ type UserType {
   """User name"""
   name: String!
   quota: UserQuota
-  subscription: UserSubscription
+  subscription(plan: SubscriptionPlan = Pro): UserSubscription @deprecated(reason: "use `UserType.subscriptions`")
+  subscriptions: [UserSubscription!]!
   token: tokenType! @deprecated(reason: "use [/api/auth/authorize]")
 }