From 2632f2282ee256a365600c83050dd814145c53ea Mon Sep 17 00:00:00 2001
From: Din <dinsterizer@gmail.com>
Date: Tue, 23 Jan 2024 10:46:59 +0700
Subject: [PATCH] optimize

---
 @api/features/auth/email.login.ts             |  2 ++
 @api/features/auth/infos.ts                   | 26 ++++++++------
 .../auth/notification.push.post-test.ts       |  2 ++
 @api/features/auth/oauth.login.ts             |  2 ++
 @api/features/auth/organization-switch.ts     | 28 ++++++++-------
 @web/pages/profile-general.tsx                |  2 +-
 @web/providers/push-notification.tsx          | 35 +++++++++++--------
 @web/stores/auth.ts                           |  2 ++
 8 files changed, 59 insertions(+), 40 deletions(-)

diff --git a/@api/features/auth/email.login.ts b/@api/features/auth/email.login.ts
index cdbc77a..8650d96 100644
--- a/@api/features/auth/email.login.ts
+++ b/@api/features/auth/email.login.ts
@@ -91,6 +91,7 @@ export const authEmailLoginRoute = procedure
       return {
         auth: {
           jwt,
+          session,
           user: existingUser,
           organization: organizationMember.organization,
           organizationMember,
@@ -130,6 +131,7 @@ export const authEmailLoginRoute = procedure
     return {
       auth: {
         jwt,
+        session,
         user: {
           ...user,
           subscriptions: [],
diff --git a/@api/features/auth/infos.ts b/@api/features/auth/infos.ts
index c4f9c80..9964c67 100644
--- a/@api/features/auth/infos.ts
+++ b/@api/features/auth/infos.ts
@@ -5,20 +5,24 @@ import { signAuthJwt } from '@api/lib/auth'
 
 export const authInfosRoute = authProcedure.query(async ({ ctx }) => {
   const session = await findSessionForAuth({ ctx, sessionSecretKey: ctx.auth.sessionSecretKey })
+  const jwt = await signAuthJwt({
+    env: ctx.env,
+    payload: {
+      sessionSecretKey: session.secretKey,
+      userId: session.user.id,
+      organizationId: session.organizationMember.organizationId,
+      organizationRole: session.organizationMember.role,
+      activeSubscriptionVariantIds: getActiveSubscriptionVariantIds(session.user.subscriptions),
+    },
+  })
 
   return {
     auth: {
-      ...session,
-      jwt: await signAuthJwt({
-        env: ctx.env,
-        payload: {
-          sessionSecretKey: session.secretKey,
-          userId: session.user.id,
-          organizationId: session.organizationMember.organizationId,
-          organizationRole: session.organizationMember.role,
-          activeSubscriptionVariantIds: getActiveSubscriptionVariantIds(session.user.subscriptions),
-        },
-      }),
+      jwt,
+      session,
+      user: session.user,
+      organization: session.organization,
+      organizationMember: session.organizationMember,
     },
   }
 })
diff --git a/@api/features/auth/notification.push.post-test.ts b/@api/features/auth/notification.push.post-test.ts
index 557c9a9..54fbedf 100644
--- a/@api/features/auth/notification.push.post-test.ts
+++ b/@api/features/auth/notification.push.post-test.ts
@@ -23,6 +23,8 @@ export const authNotificationPushPostTestRoute = authProcedure.mutation(async ({
     })
   }
 
+  await new Promise((resolve) => setTimeout(resolve, 10_000))
+
   const payload = await buildPushPayload(
     {
       data: "You've got mail!",
diff --git a/@api/features/auth/oauth.login.ts b/@api/features/auth/oauth.login.ts
index 40cfabd..888d453 100644
--- a/@api/features/auth/oauth.login.ts
+++ b/@api/features/auth/oauth.login.ts
@@ -84,6 +84,7 @@ export const authOauthLoginRoute = procedure
       return {
         auth: {
           jwt,
+          session,
           user: oauthAccount.user,
           organization: organizationMember.organization,
           organizationMember,
@@ -140,6 +141,7 @@ export const authOauthLoginRoute = procedure
     return {
       auth: {
         jwt,
+        session,
         user: {
           ...user,
           subscriptions: [],
diff --git a/@api/features/auth/organization-switch.ts b/@api/features/auth/organization-switch.ts
index 15336e8..11d7daf 100644
--- a/@api/features/auth/organization-switch.ts
+++ b/@api/features/auth/organization-switch.ts
@@ -20,22 +20,24 @@ export const authOrganizationSwitchRoute = authProcedure
       .where(eq(Sessions.secretKey, ctx.auth.sessionSecretKey))
 
     const session = await findSessionForAuth({ ctx, sessionSecretKey: ctx.auth.sessionSecretKey })
+    const jwt = await signAuthJwt({
+      env: ctx.env,
+      payload: {
+        sessionSecretKey: session.secretKey,
+        userId: session.user.id,
+        organizationId: session.organizationMember.organizationId,
+        organizationRole: session.organizationMember.role,
+        activeSubscriptionVariantIds: getActiveSubscriptionVariantIds(session.user.subscriptions),
+      },
+    })
 
     return {
       auth: {
-        ...session,
-        jwt: await signAuthJwt({
-          env: ctx.env,
-          payload: {
-            sessionSecretKey: session.secretKey,
-            userId: session.user.id,
-            organizationId: session.organizationMember.organizationId,
-            organizationRole: session.organizationMember.role,
-            activeSubscriptionVariantIds: getActiveSubscriptionVariantIds(
-              session.user.subscriptions,
-            ),
-          },
-        }),
+        jwt,
+        session,
+        user: session.user,
+        organization: session.organization,
+        organizationMember: session.organizationMember,
       },
     }
   })
diff --git a/@web/pages/profile-general.tsx b/@web/pages/profile-general.tsx
index 9d5d241..556c503 100644
--- a/@web/pages/profile-general.tsx
+++ b/@web/pages/profile-general.tsx
@@ -32,7 +32,7 @@ export function Component() {
                 <div className="pt-6 sm:flex items-center">
                   <dt className="font-medium text-foreground sm:w-64 sm:flex-none sm:pr-6">ID</dt>
                   <dd className="mt-1 flex justify-between gap-x-6 sm:mt-0 sm:flex-auto">
-                    <div className="text-muted-foreground">{query.data.auth.userId}</div>
+                    <div className="text-muted-foreground">{query.data.auth.user.id}</div>
                   </dd>
                 </div>
                 <div className="pt-6 sm:flex items-center">
diff --git a/@web/providers/push-notification.tsx b/@web/providers/push-notification.tsx
index 8276124..51ec83b 100644
--- a/@web/providers/push-notification.tsx
+++ b/@web/providers/push-notification.tsx
@@ -6,32 +6,37 @@ import { useEffect } from 'react'
 import { P, match } from 'ts-pattern'
 
 export function PushNotificationProvider(props: { children: React.ReactNode }) {
-  const useId = useAuthStore().state?.user.id
+  const userId = useAuthStore().state?.session.userId
+  const registeredPushNotification = !!useAuthStore().state?.session.pushSubscription
   const { permission } = useNotificationPermission()
+  const { mutate } = api.auth.notification.push.register.useMutation()
 
-  const mutation = api.auth.notification.push.register.useMutation()
-
-  // TODO: optimize it
   useEffect(() => {
-    match(permission)
+    if (!userId) return
+
+    return match(permission)
       .with(P.union('default', 'denied'), () => {
-        mutation.mutate({
+        if (!registeredPushNotification) return
+
+        mutate({
           subscription: null,
         })
       })
       .with('granted', () => {
+        if (registeredPushNotification) return
+
+        const unsub = onSWMessage('handlePushSubscription', async (data) => {
+          mutate({
+            subscription: data.subscription,
+          })
+        })
+
         postMessageToSW('subscribePushNotification', {})
-      })
-      .exhaustive()
-  }, [mutation, permission, useId])
 
-  useEffect(() => {
-    return onSWMessage('handlePushSubscription', async (data) => {
-      mutation.mutate({
-        subscription: data.subscription,
+        return unsub
       })
-    })
-  }, [mutation])
+      .exhaustive()
+  }, [mutate, permission, userId, registeredPushNotification])
 
   return props.children
 }
diff --git a/@web/stores/auth.ts b/@web/stores/auth.ts
index 527d019..bddf86f 100644
--- a/@web/stores/auth.ts
+++ b/@web/stores/auth.ts
@@ -1,6 +1,7 @@
 import {
   organizationMemberSchema,
   organizationSchema,
+  sessionSchema,
   subscriptionSchema,
   userSchema,
 } from '@api/database/schema'
@@ -13,6 +14,7 @@ import { persist } from 'zustand/middleware'
 export const authStateSchema = z
   .object({
     jwt: z.string(),
+    session: sessionSchema,
     user: userSchema.and(
       z.object({
         subscriptions: z.array(subscriptionSchema),