Update billing subscription endpoints/webhook logic

apps/api/src/app/controllers/billing.controller.ts

@@ -1,7 +1,6 @@
 import { ENV } from '@jetstream/api-config';
 import { Request, Response } from 'express';
 import { z } from 'zod';
-import * as subscriptionDbService from '../db/subscription.db';
 import * as userDbService from '../db/user.db';
 import * as stripeService from '../services/stripe.service';
 import { redirect, sendJson } from '../utils/response.handlers';
@@ -42,6 +41,12 @@ export const routeDefinition = {
       hasSourceOrg: false,
     },
   },
+  createBillingPortalSession: {
+    controllerFn: () => createBillingPortalSession,
+    validators: {
+      hasSourceOrg: false,
+    },
+  },
 };
 
 const stripeWebhookHandler = async (req: Request, res: Response) => {
@@ -60,12 +65,12 @@ const createCheckoutSessionHandler = createRoute(
   async ({ user: sessionUser, body }, req, res) => {
     const { priceId } = body;
 
-    const user = await userDbService.findById(sessionUser.id);
+    const user = await userDbService.findByIdWithSubscriptions(sessionUser.id);
 
     const sessions = await stripeService.createCheckoutSession({
       mode: 'subscription',
       priceId,
-      customerId: user.subscriptions?.[0]?.customerId,
+      customerId: user.billingAccount?.customerId,
       user,
     });
 
@@ -82,22 +87,42 @@ const processCheckoutSuccessHandler = createRoute(
       return redirect(res, `${ENV.JETSTREAM_CLIENT_URL}/settings/billing`);
     }
 
-    const user = await userDbService.findById(sessionUser.id);
+    const user = await userDbService.findByIdWithSubscriptions(sessionUser.id);
     if (!user.subscriptions?.[0]?.customerId) {
       return redirect(res, `${ENV.JETSTREAM_CLIENT_URL}/settings/billing`);
     }
 
-    await stripeService.saveSubscriptionFromCompletedSession({
-      sessionId,
-      userId: user.id,
-    });
+    await stripeService.saveSubscriptionFromCompletedSession({ sessionId });
 
     redirect(res, `${ENV.JETSTREAM_CLIENT_URL}/settings/billing`);
   }
 );
 
 const getSubscriptionsHandler = createRoute(routeDefinition.getSubscriptions.validators, async ({ user }, req, res) => {
-  const subscriptions = await subscriptionDbService.findByUserId(user.id);
+  const userProfile = await userDbService.findById(user.id);
+  // No billing account, so there are no subscriptions
+  if (!userProfile.billingAccount?.customerId) {
+    sendJson(res, { customer: null });
+    return;
+  }
 
-  sendJson(res, subscriptions);
+  const customer = await stripeService.getUserFacingStripeCustomer({ customerId: userProfile.billingAccount.customerId });
+  sendJson(res, { customer });
 });
+
+const createBillingPortalSession = createRoute(
+  routeDefinition.createBillingPortalSession.validators,
+  async ({ user: sessionUser }, req, res) => {
+    const user = await userDbService.findByIdWithSubscriptions(sessionUser.id);
+
+    if (!user.billingAccount) {
+      throw new Error('User does not have a billing account');
+    }
+
+    const sessions = await stripeService.createBillingPortalSession({
+      customerId: user.billingAccount?.customerId,
+    });
+
+    redirect(res, sessions.url);
+  }
+);

apps/api/src/app/db/subscription.db.ts

@@ -1,13 +1,15 @@
 /* eslint-disable @typescript-eslint/no-non-null-assertion */
 import { prisma } from '@jetstream/api-config';
+import { EntitlementsAccess, EntitlementsAccessSchema } from '@jetstream/types';
 import { Prisma } from '@prisma/client';
+import Stripe from 'stripe';
 
 const SELECT = Prisma.validator<Prisma.SubscriptionSelect>()({
   id: true,
   customerId: true,
-  providerId: true,
+  subscriptionId: true,
   status: true,
-  planId: true,
+  priceId: true,
   userId: true,
   createdAt: true,
   updatedAt: true,
@@ -21,63 +23,75 @@ export const findById = async ({ id, userId }: { id: string; userId: string }) =
   return await prisma.subscription.findUniqueOrThrow({ where: { id, userId }, select: SELECT });
 };
 
-export const upsertSubscription = async ({
-  userId,
+export const findSubscriptionsByCustomerId = async ({
   customerId,
-  providerId,
-  planId,
+  status,
 }: {
-  userId: string;
   customerId: string;
-  providerId: string;
-  planId: string;
+  status?: 'ACTIVE' | 'CANCELLED' | 'PAST_DUE' | 'PAUSED';
 }) => {
-  return await prisma.subscription.upsert({
-    create: { userId, status: 'ACTIVE', customerId, providerId, planId },
-    update: { userId, status: 'ACTIVE', customerId, providerId, planId },
-    where: { uniqueSubscription: { userId, providerId, planId } },
-    select: SELECT,
+  return await prisma.subscription.findMany({ where: { customerId, status }, select: SELECT });
+};
+
+export const updateUserEntitlements = async (customerId: string, entitlementAccessUntrusted: EntitlementsAccess) => {
+  const entitlementAccess = EntitlementsAccessSchema.parse(entitlementAccessUntrusted);
+  const user = await prisma.user.findFirstOrThrow({
+    where: { billingAccount: { customerId } },
+    select: { id: true },
+  });
+
+  await prisma.entitlement.upsert({
+    create: {
+      userId: user.id,
+      ...entitlementAccess,
+    },
+    update: entitlementAccess,
+    where: { userId: user.id },
   });
 };
 
-// export const create = async (
-//   userId: string,
-//   payload: {
-//     name: string;
-//     description?: Maybe<string>;
-//   }
-// ) => {
-//   return await prisma.subscription.create({
-//     select: SELECT,
-//     data: {
-//       userId,
-//       name: payload.name.trim(),
-//       description: payload.description?.trim(),
-//     },
-//   });
-// };
+/**
+ * Given a customer's current subscriptions, cancel all other subscriptions, create any needed subscriptions, and update the subscription state
+ * In addition, entitlements are also updated to reflect the user's current subscription state
+ */
+export const updateSubscriptionStateForCustomer = async ({
+  userId,
+  customerId,
+  subscriptions,
+}: {
+  userId: string;
+  customerId: string;
+  subscriptions: Stripe.Subscription[];
+}) => {
+  const priceIds = subscriptions.flatMap((subscription) => subscription.items.data.map((item) => item.price.id));
 
-// export const update = async (
-//   userId,
-//   id,
-//   payload: {
-//     name: string;
-//     description?: Maybe<string>;
-//   }
-// ) => {
-//   return await prisma.subscription.update({
-//     select: SELECT,
-//     where: { userId, id },
-//     data: {
-//       name: payload.name.trim(),
-//       description: payload.description?.trim() ?? null,
-//     },
-//   });
-// };
+  await prisma.$transaction([
+    // Delete all subscriptions that are no longer active in Stripe
+    prisma.subscription.deleteMany({
+      where: { userId, customerId, priceId: { notIn: priceIds } },
+    }),
+    // Create/Update all current subscriptions from Stripe
+    ...subscriptions.flatMap((subscription) =>
+      subscription.items.data.map((item) =>
+        prisma.subscription.upsert({
+          create: {
+            userId,
+            subscriptionId: subscription.id,
+            status: subscription.status.toUpperCase(),
+            customerId,
+            priceId: item.price.id,
+          },
+          update: { status: subscription.status.toUpperCase() },
+          where: { uniqueSubscription: { userId, subscriptionId: subscription.id, priceId: item.price.id } },
+        })
+      )
+    ),
+  ]);
+};
 
-// export const deleteOrganization = async (userId, id) => {
-//   return await prisma.subscription.delete({
-//     select: SELECT,
-//     where: { userId, id },
-//   });
-// };
+export const cancelAllSubscriptionsForUser = async ({ customerId }: { customerId: string }) => {
+  await prisma.subscription.updateMany({
+    where: { customerId },
+    data: { status: 'CANCELED' },
+  });
+};

apps/api/src/app/db/user.db.ts

@@ -21,15 +21,9 @@ const userSelect: Prisma.UserSelect = {
       chromeExtension: true,
     },
   },
-  subscriptions: {
-    where: { status: 'ACTIVE' },
-    take: 1,
+  billingAccount: {
     select: {
-      id: true,
-      providerId: true,
       customerId: true,
-      planId: true,
-      status: true,
     },
   },
   updatedAt: true,
@@ -71,15 +65,9 @@ const FullUserFacingProfileSelect = Prisma.validator<Prisma.UserSelect & { hasPa
       updatedAt: true,
     },
   },
-  subscriptions: {
-    where: { status: 'ACTIVE' },
-    take: 1,
+  billingAccount: {
     select: {
-      id: true,
-      providerId: true,
       customerId: true,
-      planId: true,
-      status: true,
     },
   },
   createdAt: true,
@@ -94,14 +82,23 @@ const UserFacingProfileSelect = Prisma.validator<Prisma.UserSelect>()({
   emailVerified: true,
   picture: true,
   preferences: true,
+  billingAccount: {
+    select: {
+      customerId: true,
+    },
+  },
+  entitlements: {
+    select: {
+      chromeExtension: true,
+      recordSync: true,
+    },
+  },
   subscriptions: {
-    where: { status: 'ACTIVE' },
-    take: 1,
     select: {
       id: true,
-      providerId: true,
-      customerId: true,
-      planId: true,
+      productId: true,
+      subscriptionId: true,
+      priceId: true,
       status: true,
     },
   },
@@ -118,8 +115,28 @@ export const findById = (id: string) => {
   return prisma.user.findFirstOrThrow({ where: { id }, select: UserFacingProfileSelect });
 };
 
+export const findByIdWithSubscriptions = (id: string) => {
+  return prisma.user.findFirstOrThrow({
+    where: { id },
+    select: {
+      ...userSelect,
+      subscriptions: {
+        where: { status: 'ACTIVE' },
+        select: {
+          id: true,
+          customerId: true,
+          productId: true,
+          subscriptionId: true,
+          priceId: true,
+          status: true,
+        },
+      },
+    },
+  });
+};
+
 export const findIdByUserIdUserFacing = ({ userId }: { userId: string }) => {
-  return prisma.user.findFirstOrThrow({ where: { id: userId }, select: UserFacingProfileSelect }).then(({ id }) => id);
+  return prisma.user.findFirstOrThrow({ where: { id: userId }, select: UserFacingProfileSelect });
 };
 
 export const checkUserEntitlement = ({
@@ -129,7 +146,7 @@ export const checkUserEntitlement = ({
   userId: string;
   entitlement: keyof Omit<Entitlement, 'id' | 'userId' | 'createdAt' | 'updatedAt'>;
 }): Promise<boolean> => {
-  return prisma.entitlement.count({ where: { userId, [entitlement]: true } }).then((result) => result > 0);
+  return prisma.entitlement.count({ where: { id: userId, [entitlement]: true } }).then((result) => result > 0);
 };
 
 export async function updateUser(
@@ -178,3 +195,18 @@ export async function deleteUserAndAllRelatedData(userId: string): Promise<void>
     },
   });
 }
+
+export async function findByBillingAccountByCustomerId({ customerId }: { customerId: string }) {
+  const billingAccount = await prisma.billingAccount.findFirst({ where: { customerId } });
+  return billingAccount;
+}
+
+export async function createBillingAccountIfNotExists({ userId, customerId }: { userId: string; customerId: string }) {
+  const existingCustomer = await prisma.billingAccount.findUnique({ where: { uniqueCustomer: { customerId, userId } } });
+  if (existingCustomer) {
+    return existingCustomer;
+  }
+  return await prisma.billingAccount.create({
+    data: { customerId, userId },
+  });
+}

apps/api/src/app/routes/api.routes.ts

@@ -64,6 +64,7 @@ routes.delete('/me/profile/2fa/:type', userController.deleteAuthFactor.controlle
 routes.post('/billing/checkout-session', billingController.createCheckoutSession.controllerFn());
 routes.get('/billing/checkout-session/complete', billingController.processCheckoutSuccess.controllerFn());
 routes.get('/billing/subscriptions', billingController.getSubscriptions.controllerFn());
+routes.post('/billing/portal', billingController.createBillingPortalSession.controllerFn());
 
 /**
  * ************************************