From 10af0ab48dcba853853e4134ecc32c0fd199eb21 Mon Sep 17 00:00:00 2001 From: Brooooooklyn Date: Fri, 22 Mar 2024 08:39:17 +0000 Subject: [PATCH] feat(server): support ai plan (#6216) --- .../migration.sql | 11 ++ packages/backend/server/schema.prisma | 5 +- .../server/src/plugins/payment/resolver.ts | 69 ++++++++++- .../server/src/plugins/payment/service.ts | 113 ++++++++++++------ .../server/src/plugins/payment/types.ts | 1 + packages/backend/server/src/schema.gql | 10 +- 6 files changed, 158 insertions(+), 51 deletions(-) create mode 100644 packages/backend/server/migrations/20240319104623_user_subscriptions/migration.sql diff --git a/packages/backend/server/migrations/20240319104623_user_subscriptions/migration.sql b/packages/backend/server/migrations/20240319104623_user_subscriptions/migration.sql new file mode 100644 index 000000000000..42a46f6b9a8e --- /dev/null +++ b/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"); diff --git a/packages/backend/server/schema.prisma b/packages/backend/server/schema.prisma index 508f4052c52f..21bc89988313 100644 --- a/packages/backend/server/schema.prisma +++ b/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") } diff --git a/packages/backend/server/src/plugins/payment/resolver.ts b/packages/backend/server/src/plugins/payment/resolver.ts index 7fcb95c2fc6d..95882b4d2704 100644 --- a/packages/backend/server/src/plugins/payment/resolver.ts +++ b/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 { + 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, diff --git a/packages/backend/server/src/plugins/payment/service.ts b/packages/backend/server/src/plugins/payment/service.ts index 57728032ac7d..b13ac10936fe 100644 --- a/packages/backend/server/src/plugins/payment/service.ts +++ b/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 { 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'); + } + + 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' + ); } // 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 { 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'); + } + + 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 not been canceled'); + if (!subscriptionInDB.canceledAt) { + throw new BadRequestException('Your subscription has not been canceled'); } - if (user.subscription.end < new Date()) { - throw new Error('Your subscription is expired, please checkout again.'); + if (subscriptionInDB.end < new Date()) { + throw new BadRequestException( + 'Your subscription is expired, please checkout again.' + ); } - if (user.subscription.stripeScheduleId) { + 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 { 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` ); } diff --git a/packages/backend/server/src/plugins/payment/types.ts b/packages/backend/server/src/plugins/payment/types.ts index 7954ad7f4163..d42c87778e3f 100644 --- a/packages/backend/server/src/plugins/payment/types.ts +++ b/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', diff --git a/packages/backend/server/src/schema.gql b/packages/backend/server/src/schema.gql index fd605ad8be90..e91fb650261d 100644 --- a/packages/backend/server/src/schema.gql +++ b/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]") }