cleanup

Author: zdql

packages/app/control/src/app/(app)/app/[id]/(overview)/_components/referral-handler.tsx

@@ -13,54 +13,25 @@ export const ReferralHandler: React.FC<Props> = ({ appId }) => {
   const referralCode = searchParams.get('referral_code');
   const [processed, setProcessed] = useState(false);
 
-  const utils = api.useUtils();
-  const { mutateAsync: createMembership } =
-    api.apps.app.memberships.create.useMutation();
-  const { mutateAsync: updateReferrer } =
-    api.apps.app.memberships.update.referrer.useMutation();
+  const { mutateAsync: registerReferral } =
+    api.apps.app.registerReferral.useMutation();
 
   useEffect(() => {
     if (!referralCode || processed) return;
 
     const processReferralCode = async () => {
-      const referralCodeData = await utils.apps.app.referralCode.get.byCode
-        .fetch(referralCode)
-        .catch(() => null);
-
-      if (!referralCodeData) {
-        setProcessed(true);
-        return;
-      }
-
-      const membership = await utils.apps.app.memberships.get
-        .fetch({ appId })
-        .catch(() => null);
-
-      // Only update if user has no referrer yet
-      if (membership?.referrerId === null) {
-        await updateReferrer({
-          appId,
-          referrerId: referralCodeData.id,
-        }).catch(() => {
-          // Silently fail - user might already have a referrer from a race condition
-        });
-      }
-
-      // Create membership if it doesn't exist
-      if (!membership) {
-        await createMembership({
-          appId,
-          referrerId: referralCodeData.id,
-        }).catch(() => {
-          // Silently fail - membership might have been created in the meantime
-        });
-      }
+      await registerReferral({
+        appId,
+        code: referralCode,
+      }).catch(() => {
+        // Silently fail - referral code may be invalid, expired, or user may already have a referrer
+      });
 
       setProcessed(true);
     };
 
     void processReferralCode();
-  }, [referralCode, appId, createMembership, updateReferrer, processed, utils]);
+  }, [referralCode, appId, registerReferral, processed]);
 
   return null;
 };

packages/app/control/src/trpc/routers/apps/index.ts

@@ -23,6 +23,7 @@ import {
   createAppMembershipSchema,
   updateAppMembershipReferrer,
   updateAppMembershipReferrerSchema,
+  setAppMembershipReferrer,
 } from '@/services/db/apps/membership';
 import {
   listAppsSchema,
@@ -262,6 +263,26 @@ export const appsRouter = createTRPCRouter({
         }),
     },
 
+    registerReferral: protectedProcedure
+      .input(z.object({ appId: appIdSchema, code: z.string() }))
+      .mutation(async ({ input, ctx }) => {
+        const success = await setAppMembershipReferrer(
+          ctx.session.user.id,
+          input.appId,
+          input.code
+        );
+
+        if (!success) {
+          throw new TRPCError({
+            code: 'BAD_REQUEST',
+            message:
+              'Referral code could not be applied. It may be invalid, expired, or you may already have a referrer for this app.',
+          });
+        }
+
+        return { success: true };
+      }),
+
     transactions: {
       list: paginatedProcedure
         .concat(protectedProcedure)