From 42cc0a2781248b5dd3dbc21d8fb17e28ef9b90a0 Mon Sep 17 00:00:00 2001
From: jc <46619361+juancwu@users.noreply.github.com>
Date: Tue, 10 Dec 2024 00:14:07 -0500
Subject: [PATCH] do not throw on company not found on login

---
 backend/internal/server/company.go    |  3 +++
 frontend/src/contexts/AuthContext.tsx |  1 -
 frontend/src/pages/Register.tsx       | 12 +++++++++---
 frontend/src/services/company.ts      |  6 +++++-
 4 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/backend/internal/server/company.go b/backend/internal/server/company.go
index 928fe60..611b515 100644
--- a/backend/internal/server/company.go
+++ b/backend/internal/server/company.go
@@ -48,6 +48,9 @@ func (s *Server) handleGetUserCompany(c echo.Context) error {
 
 	company, err := s.queries.GetCompanyByUser(c.Request().Context(), claims.UserID)
 	if err != nil {
+		if isNoRowsError(err) {
+			return echo.NewHTTPError(http.StatusNotFound, "No company found")
+		}
 		log.Error().Err(err).Msg("Failed to get company by user")
 		return echo.NewHTTPError(http.StatusInternalServerError, "Failed to get company")
 	}
diff --git a/frontend/src/contexts/AuthContext.tsx b/frontend/src/contexts/AuthContext.tsx
index 7a53692..0707e6a 100644
--- a/frontend/src/contexts/AuthContext.tsx
+++ b/frontend/src/contexts/AuthContext.tsx
@@ -60,4 +60,3 @@ export function useAuth() {
     }
     return context;
 }
-
diff --git a/frontend/src/pages/Register.tsx b/frontend/src/pages/Register.tsx
index 8ef8528..83b37d2 100644
--- a/frontend/src/pages/Register.tsx
+++ b/frontend/src/pages/Register.tsx
@@ -332,7 +332,9 @@ const Register = () => {
     const location = useLocation();
     const { user, setAuth, clearAuth } = useAuth();
     const [currentStep, setCurrentStep] = useState<RegistrationStep>(() => {
-        const locationState = location.state as { step?: RegistrationStep } | null;
+        const locationState = location.state as {
+            step?: RegistrationStep;
+        } | null;
 
         if (locationState?.step) {
             return locationState.step;
@@ -440,7 +442,11 @@ const Register = () => {
         try {
             const signinResp = await signin(formData.email, formData.password);
             const company = await getCompany(signinResp.access_token);
-            setAuth(signinResp.user, signinResp.access_token, company.ID);
+            setAuth(
+                signinResp.user,
+                signinResp.access_token,
+                company ? company.ID : null
+            );
 
             if (!signinResp.user.isEmailVerified) {
                 setCurrentStep('verify-email');
@@ -451,7 +457,7 @@ const Register = () => {
                 setCurrentStep('form-details');
                 return;
             }
-            
+
             // Redirect based on user role
             if (signinResp.user.role === 'admin') {
                 navigate('/admin/dashboard', { replace: true });
diff --git a/frontend/src/services/company.ts b/frontend/src/services/company.ts
index 6a66732..dd0894c 100644
--- a/frontend/src/services/company.ts
+++ b/frontend/src/services/company.ts
@@ -42,7 +42,7 @@ export async function createCompany(
 
 export async function getCompany(
     accessToken: string
-): Promise<CreateCompanyResponse> {
+): Promise<CreateCompanyResponse | null> {
     const url = getApiUrl('/company');
 
     const res = await fetch(url, {
@@ -54,6 +54,10 @@ export async function getCompany(
 
     const json = await res.json();
 
+    if (res.status === HttpStatusCode.NOT_FOUND) {
+        return null;
+    }
+
     if (res.status !== HttpStatusCode.OK) {
         throw new ApiError('Failed to create company', res.status, json);
     }