From af2158cb0c3f3475dbe4479e0d636024733a05ee Mon Sep 17 00:00:00 2001
From: CatsJuice <wp1298554944@hotmail.com>
Date: Mon, 1 Apr 2024 07:28:33 +0000
Subject: [PATCH] feat(core): compatible with multiple subscriptions (#6277)

---
 .../affine/auth/subscription-redirect.tsx     |  8 ++++--
 .../core/src/hooks/use-subscription.ts        | 25 ++++++++++---------
 .../frontend/graphql/src/graphql/index.ts     |  2 +-
 .../graphql/src/graphql/subscription.gql      |  2 +-
 packages/frontend/graphql/src/schema.ts       |  4 +--
 5 files changed, 23 insertions(+), 18 deletions(-)

diff --git a/packages/frontend/core/src/components/affine/auth/subscription-redirect.tsx b/packages/frontend/core/src/components/affine/auth/subscription-redirect.tsx
index 65a2987fdc..d66f17bfbd 100644
--- a/packages/frontend/core/src/components/affine/auth/subscription-redirect.tsx
+++ b/packages/frontend/core/src/components/affine/auth/subscription-redirect.tsx
@@ -4,7 +4,7 @@ import { Loading } from '@affine/component/ui/loading';
 import { AffineShapeIcon } from '@affine/core/components/page-list';
 import { useCredentialsRequirement } from '@affine/core/hooks/affine/use-server-config';
 import { useAsyncCallback } from '@affine/core/hooks/affine-async-hooks';
-import type { SubscriptionPlan, SubscriptionRecurring } from '@affine/graphql';
+import { SubscriptionPlan, type SubscriptionRecurring } from '@affine/graphql';
 import {
   changePasswordMutation,
   createCheckoutSessionMutation,
@@ -138,7 +138,11 @@ const SubscriptionRedirectWithData = () => {
     );
   }
 
-  if (subscriptionData.currentUser?.subscription) {
+  if (
+    subscriptionData.currentUser?.subscriptions?.some(
+      sub => sub.plan === SubscriptionPlan.Pro
+    )
+  ) {
     return <SubscriptionExisting />;
   }
 
diff --git a/packages/frontend/core/src/hooks/use-subscription.ts b/packages/frontend/core/src/hooks/use-subscription.ts
index c3e81816f7..8191b2a9fa 100644
--- a/packages/frontend/core/src/hooks/use-subscription.ts
+++ b/packages/frontend/core/src/hooks/use-subscription.ts
@@ -1,20 +1,22 @@
 import { useAsyncCallback } from '@affine/core/hooks/affine-async-hooks';
 import type { SubscriptionQuery } from '@affine/graphql';
-import { subscriptionQuery } from '@affine/graphql';
+import { SubscriptionPlan, subscriptionQuery } from '@affine/graphql';
 
 import { useServerFeatures } from './affine/use-server-config';
 import { useQuery } from './use-query';
 
 export type Subscription = NonNullable<
-  NonNullable<SubscriptionQuery['currentUser']>['subscription']
+  NonNullable<SubscriptionQuery['currentUser']>['subscriptions'][number]
 >;
 
 export type SubscriptionMutator = (update?: Partial<Subscription>) => void;
 
-const selector = (data: SubscriptionQuery) =>
-  data.currentUser?.subscription ?? null;
+const selector = (data: SubscriptionQuery, plan: SubscriptionPlan) =>
+  (data.currentUser?.subscriptions ?? []).find(p => p.plan === plan);
 
-export const useUserSubscription = () => {
+export const useUserSubscription = (
+  plan: SubscriptionPlan = SubscriptionPlan.Pro
+) => {
   const { payment: hasPaymentFeature } = useServerFeatures();
   const { data, mutate } = useQuery(
     hasPaymentFeature ? { query: subscriptionQuery } : undefined
@@ -23,26 +25,25 @@ export const useUserSubscription = () => {
   const set: SubscriptionMutator = useAsyncCallback(
     async (update?: Partial<Subscription>) => {
       await mutate(prev => {
-        if (!update || !prev?.currentUser?.subscription) {
+        if (!update || !prev?.currentUser?.subscriptions?.length) {
           return;
         }
 
         return {
           currentUser: {
-            subscription: {
-              ...prev.currentUser?.subscription,
-              ...update,
-            },
+            subscriptions: (prev.currentUser?.subscriptions ?? []).map(sub =>
+              sub.plan !== plan ? sub : { ...sub, ...update }
+            ),
           },
         };
       });
     },
-    [mutate]
+    [mutate, plan]
   );
 
   if (!hasPaymentFeature) {
     return [null, () => {}] as const;
   }
 
-  return [selector(data), set] as const;
+  return [selector(data, plan), set] as const;
 };
diff --git a/packages/frontend/graphql/src/graphql/index.ts b/packages/frontend/graphql/src/graphql/index.ts
index 188af142f8..014afe243f 100644
--- a/packages/frontend/graphql/src/graphql/index.ts
+++ b/packages/frontend/graphql/src/graphql/index.ts
@@ -748,7 +748,7 @@ export const subscriptionQuery = {
   query: `
 query subscription {
   currentUser {
-    subscription {
+    subscriptions {
       id
       status
       plan
diff --git a/packages/frontend/graphql/src/graphql/subscription.gql b/packages/frontend/graphql/src/graphql/subscription.gql
index 0072350f1f..48811aa690 100644
--- a/packages/frontend/graphql/src/graphql/subscription.gql
+++ b/packages/frontend/graphql/src/graphql/subscription.gql
@@ -1,6 +1,6 @@
 query subscription {
   currentUser {
-    subscription {
+    subscriptions {
       id
       status
       plan
diff --git a/packages/frontend/graphql/src/schema.ts b/packages/frontend/graphql/src/schema.ts
index 258ab904f4..a0f408e969 100644
--- a/packages/frontend/graphql/src/schema.ts
+++ b/packages/frontend/graphql/src/schema.ts
@@ -743,7 +743,7 @@ export type SubscriptionQuery = {
   __typename?: 'Query';
   currentUser: {
     __typename?: 'UserType';
-    subscription: {
+    subscriptions: Array<{
       __typename?: 'UserSubscription';
       id: string;
       status: SubscriptionStatus;
@@ -753,7 +753,7 @@ export type SubscriptionQuery = {
       end: string;
       nextBillAt: string | null;
       canceledAt: string | null;
-    } | null;
+    }>;
   } | null;
 };