refactor(core): make subscription hook (#4669)

packages/frontend/core/src/components/affine/setting-modal/general-setting/billing/index.tsx

@@ -12,7 +12,6 @@ import {
   pricesQuery,
   resumeSubscriptionMutation,
   SubscriptionPlan,
-  subscriptionQuery,
   SubscriptionRecurring,
   SubscriptionStatus,
 } from '@affine/graphql';
@@ -20,10 +19,14 @@ import { useMutation, useQuery } from '@affine/workspace/affine/gql';
 import { ArrowRightSmallIcon } from '@blocksuite/icons';
 import { Button, IconButton } from '@toeverything/components/button';
 import { useSetAtom } from 'jotai';
-import { Suspense, useCallback, useEffect } from 'react';
+import { Suspense, useCallback } from 'react';
 
 import { openSettingModalAtom } from '../../../../../atoms';
 import { useCurrentLoginStatus } from '../../../../../hooks/affine/use-current-login-status';
+import {
+  type SubscriptionMutator,
+  useUserSubscription,
+} from '../../../../../hooks/use-subscription';
 import * as styles from './style.css';
 
 export const BillingSettings = () => {
@@ -56,14 +59,11 @@ export const BillingSettings = () => {
 };
 
 const SubscriptionSettings = () => {
-  const { data: subscriptionQueryResult } = useQuery({
+  const [subscription, mutateSubscription] = useUserSubscription();
-    query: subscriptionQuery,
-  });
   const { data: pricesQueryResult } = useQuery({
     query: pricesQuery,
   });
 
-  const subscription = subscriptionQueryResult.currentUser?.subscription;
   const plan = subscription?.plan ?? SubscriptionPlan.Free;
   const recurring = subscription?.recurring ?? SubscriptionRecurring.Monthly;
 
@@ -123,7 +123,7 @@ const SubscriptionSettings = () => {
                 subscription.end
               ).toLocaleDateString()}`}
             >
-              <ResumeSubscription />
+              <ResumeSubscription onSubscriptionUpdate={mutateSubscription} />
             </SettingRow>
           ) : (
             <SettingRow
@@ -133,7 +133,7 @@ const SubscriptionSettings = () => {
                 subscription.end
               ).toLocaleDateString()}`}
             >
-              <CancelSubscription />
+              <CancelSubscription onSubscriptionUpdate={mutateSubscription} />
             </SettingRow>
           )}
         </>
@@ -166,20 +166,18 @@ const PlanAction = ({ plan }: { plan: string }) => {
 
 const PaymentMethodUpdater = () => {
   // TODO: open stripe customer portal
-  const { isMutating, trigger, data } = useMutation({
+  const { isMutating, trigger } = useMutation({
     mutation: createCustomerPortalMutation,
   });
 
   const update = useCallback(() => {
-    trigger();
+    trigger(null, {
+      onSuccess: data => {
+        window.open(data.createCustomerPortal, '_blank', 'noopener noreferrer');
+      },
+    });
   }, [trigger]);
 
-  useEffect(() => {
-    if (data?.createCustomerPortal) {
-      window.open(data.createCustomerPortal, '_blank', 'noopener noreferrer');
-    }
-  }, [data]);
-
   return (
     <Button onClick={update} loading={isMutating} disabled={isMutating}>
       Update
@@ -187,14 +185,22 @@ const PaymentMethodUpdater = () => {
   );
 };
 
-const ResumeSubscription = () => {
+const ResumeSubscription = ({
+  onSubscriptionUpdate,
+}: {
+  onSubscriptionUpdate: SubscriptionMutator;
+}) => {
   const { isMutating, trigger } = useMutation({
     mutation: resumeSubscriptionMutation,
   });
 
   const resume = useCallback(() => {
-    trigger();
+    trigger(null, {
-  }, [trigger]);
+      onSuccess: data => {
+        onSubscriptionUpdate(data.resumeSubscription);
+      },
+    });
+  }, [trigger, onSubscriptionUpdate]);
 
   return (
     <Button onClick={resume} loading={isMutating} disabled={isMutating}>
@@ -203,14 +209,22 @@ const ResumeSubscription = () => {
   );
 };
 
-const CancelSubscription = () => {
+const CancelSubscription = ({
+  onSubscriptionUpdate,
+}: {
+  onSubscriptionUpdate: SubscriptionMutator;
+}) => {
   const { isMutating, trigger } = useMutation({
     mutation: cancelSubscriptionMutation,
   });
 
   const cancel = useCallback(() => {
-    trigger();
+    trigger(null, {
-  }, [trigger]);
+      onSuccess: data => {
+        onSubscriptionUpdate(data.cancelSubscription);
+      },
+    });
+  }, [trigger, onSubscriptionUpdate]);
 
   return (
     <IconButton

packages/frontend/core/src/components/affine/setting-modal/general-setting/plans/index.tsx

@@ -5,7 +5,6 @@ import {
   checkoutMutation,
   pricesQuery,
   SubscriptionPlan,
-  subscriptionQuery,
   SubscriptionRecurring,
   updateSubscriptionMutation,
 } from '@affine/graphql';
@@ -20,6 +19,11 @@ import {
   useState,
 } from 'react';
 
+import { useCurrentLoginStatus } from '../../../../../hooks/affine/use-current-login-status';
+import {
+  type SubscriptionMutator,
+  useUserSubscription,
+} from '../../../../../hooks/use-subscription';
 import * as styles from './style.css';
 
 interface FixedPrice {
@@ -102,9 +106,8 @@ const planDetail = new Map<SubscriptionPlan, FixedPrice | DynamicPrice>([
 ]);
 
 const Settings = () => {
-  const { data, mutate } = useQuery({
+  const [subscription, mutateSubscription] = useUserSubscription();
-    query: subscriptionQuery,
+  const loggedIn = useCurrentLoginStatus() === 'authenticated';
-  });
 
   const {
     data: { prices },
@@ -125,9 +128,6 @@ const Settings = () => {
     }
   });
 
-  const loggedIn = !!data.currentUser;
-  const subscription = data.currentUser?.subscription;
-
   const [recurring, setRecurring] = useState<string>(
     subscription?.recurring ?? SubscriptionRecurring.Yearly
   );
@@ -135,10 +135,6 @@ const Settings = () => {
   const currentPlan = subscription?.plan ?? SubscriptionPlan.Free;
   const currentRecurring = subscription?.recurring;
 
-  const refresh = useCallback(() => {
-    mutate();
-  }, [mutate]);
-
   const yearlyDiscount = (
     planDetail.get(SubscriptionPlan.Pro) as FixedPrice | undefined
   )?.discount;
@@ -228,19 +224,19 @@ const Settings = () => {
                           detail.plan === SubscriptionPlan.Free)) ? (
                         <CurrentPlan />
                       ) : detail.plan === SubscriptionPlan.Free ? (
-                        <Downgrade onActionDone={refresh} />
+                        <Downgrade onSubscriptionUpdate={mutateSubscription} />
                       ) : currentRecurring !== recurring &&
                         currentPlan === detail.plan ? (
                         <ChangeRecurring
                           // @ts-expect-error must exist
                           from={currentRecurring}
                           to={recurring as SubscriptionRecurring}
-                          onActionDone={refresh}
+                          onSubscriptionUpdate={mutateSubscription}
                         />
                       ) : (
                         <Upgrade
                           recurring={recurring as SubscriptionRecurring}
-                          onActionDone={refresh}
+                          onSubscriptionUpdate={mutateSubscription}
                         />
                       )
                     ) : (
@@ -280,14 +276,22 @@ const Settings = () => {
   );
 };
 
-const Downgrade = ({ onActionDone }: { onActionDone: () => void }) => {
+const Downgrade = ({
+  onSubscriptionUpdate,
+}: {
+  onSubscriptionUpdate: SubscriptionMutator;
+}) => {
   const { isMutating, trigger } = useMutation({
     mutation: cancelSubscriptionMutation,
   });
 
   const downgrade = useCallback(() => {
-    trigger(null, { onSuccess: onActionDone });
+    trigger(null, {
-  }, [trigger, onActionDone]);
+      onSuccess: data => {
+        onSubscriptionUpdate(data.cancelSubscription);
+      },
+    });
+  }, [trigger, onSubscriptionUpdate]);
 
   return (
     <Button
@@ -304,48 +308,57 @@ const Downgrade = ({ onActionDone }: { onActionDone: () => void }) => {
 
 const Upgrade = ({
   recurring,
-  onActionDone,
+  onSubscriptionUpdate,
 }: {
   recurring: SubscriptionRecurring;
-  onActionDone: () => void;
+  onSubscriptionUpdate: SubscriptionMutator;
 }) => {
-  const { isMutating, trigger, data } = useMutation({
+  const { isMutating, trigger } = useMutation({
     mutation: checkoutMutation,
   });
 
-  const upgrade = useCallback(() => {
-    trigger({ recurring });
-  }, [trigger, recurring]);
-
   const newTabRef = useRef<Window | null>(null);
 
-  useEffect(() => {
+  const onClose = useCallback(() => {
-    if (data?.checkout) {
+    newTabRef.current = null;
-      if (newTabRef.current) {
+    onSubscriptionUpdate();
-        newTabRef.current.focus();
+  }, [onSubscriptionUpdate]);
-      } else {
-        // FIXME: safari prevents from opening new tab by window api
-        // TODO(@xp): what if electron?
-        const newTab = window.open(
-          data.checkout,
-          '_blank',
-          'noopener noreferrer'
-        );
 
-        if (newTab) {
+  const upgrade = useCallback(() => {
-          newTabRef.current = newTab;
+    if (newTabRef.current) {
-          const update = () => {
+      newTabRef.current.focus();
-            onActionDone();
+    } else {
-          };
+      trigger(
-          newTab.addEventListener('close', update);
+        { recurring },
+        {
+          onSuccess: data => {
+            // FIXME: safari prevents from opening new tab by window api
+            // TODO(@xp): what if electron?
+            const newTab = window.open(
+              data.checkout,
+              '_blank',
+              'noopener noreferrer'
+            );
 
-          return () => newTab.removeEventListener('close', update);
+            if (newTab) {
+              newTabRef.current = newTab;
+
+              newTab.addEventListener('close', onClose);
+            }
+          },
         }
-      }
+      );
     }
+  }, [trigger, recurring, onClose]);
 
-    return;
+  useEffect(() => {
-  }, [data?.checkout, onActionDone]);
+    return () => {
+      if (newTabRef.current) {
+        newTabRef.current.removeEventListener('close', onClose);
+        newTabRef.current = null;
+      }
+    };
+  }, [onClose]);
 
   return (
     <Button
@@ -363,19 +376,26 @@ const Upgrade = ({
 const ChangeRecurring = ({
   from: _from /* TODO: from can be useful when showing confirmation modal  */,
   to,
-  onActionDone,
+  onSubscriptionUpdate,
 }: {
   from: SubscriptionRecurring;
   to: SubscriptionRecurring;
-  onActionDone: () => void;
+  onSubscriptionUpdate: SubscriptionMutator;
 }) => {
   const { isMutating, trigger } = useMutation({
     mutation: updateSubscriptionMutation,
   });
 
   const change = useCallback(() => {
-    trigger({ recurring: to }, { onSuccess: onActionDone });
+    trigger(
-  }, [trigger, onActionDone, to]);
+      { recurring: to },
+      {
+        onSuccess: data => {
+          onSubscriptionUpdate(data.updateSubscriptionRecurring);
+        },
+      }
+    );
+  }, [trigger, onSubscriptionUpdate, to]);
 
   return (
     <Button

packages/frontend/core/src/hooks/use-subscription.ts

@@ -0,0 +1,40 @@
+import { type SubscriptionQuery, subscriptionQuery } from '@affine/graphql';
+import { useQuery } from '@affine/workspace/affine/gql';
+import { useCallback } from 'react';
+
+export type Subscription = NonNullable<
+  NonNullable<SubscriptionQuery['currentUser']>['subscription']
+>;
+
+export type SubscriptionMutator = (update?: Partial<Subscription>) => void;
+
+const selector = (data: SubscriptionQuery) =>
+  data.currentUser?.subscription ?? null;
+
+export const useUserSubscription = () => {
+  const { data, mutate } = useQuery({
+    query: subscriptionQuery,
+  });
+
+  const set: SubscriptionMutator = useCallback(
+    (update?: Partial<Subscription>) => {
+      mutate(prev => {
+        if (!update || !prev?.currentUser?.subscription) {
+          return;
+        }
+
+        return {
+          currentUser: {
+            subscription: {
+              ...prev.currentUser?.subscription,
+              ...update,
+            },
+          },
+        };
+      });
+    },
+    [mutate]
+  );
+
+  return [selector(data), set] as const;
+};