Move shared assume logic into hook to use on employee pages

Signed-off-by: Émile Ré <emile@getprobo.com>

apps/console/relay.config.json

@@ -36,6 +36,7 @@
   },
   "sources": {
     "apps/console/src/pages/iam": "iam",
+    "apps/console/src/hooks/iam": "iam",
     "apps/console/src": "core"
   }
 }

apps/console/src/__generated__/iam/PasswordSignInPageMutation.graphql.ts

@@ -1,5 +1,5 @@
 /**
- * @generated SignedSource<<f242c4003408e8b86541c8f98b2c6d18>>
+ * @generated SignedSource<<3bef621b9a5eb43afa619cd76b2bb6b6>>
  * @lightSyntaxTransform
  * @nogrep
  */
@@ -11,6 +11,7 @@
 import { ConcreteRequest } from 'relay-runtime';
 export type SignInInput = {
   email: string;
+  organizationId?: string | null | undefined;
   password: string;
 };
 export type PasswordSignInPageMutation$variables = {

apps/console/src/__generated__/iam/useAssumeMutation.graphql.ts

@@ -1,5 +1,5 @@
 /**
- * @generated SignedSource<<09963c0483a8416abf9e411f4a0b716b>>
+ * @generated SignedSource<<b5a89cb2acee801ab7d41cf2afb471b6>>
  * @lightSyntaxTransform
  * @nogrep
  */
@@ -13,10 +13,10 @@ export type ReauthenticationReason = "POLICY_REQUIREMENT" | "SENSITIVE_ACTION" |
 export type AssumeOrganizationSessionInput = {
   organizationId: string;
 };
-export type ViewerMembershipLayoutLoader_assumeMutation$variables = {
+export type useAssumeMutation$variables = {
   input: AssumeOrganizationSessionInput;
 };
-export type ViewerMembershipLayoutLoader_assumeMutation$data = {
+export type useAssumeMutation$data = {
   readonly assumeOrganizationSession: {
     readonly result: {
       readonly __typename: "OrganizationSessionCreated";
@@ -41,9 +41,9 @@ export type ViewerMembershipLayoutLoader_assumeMutation$data = {
     };
   } | null | undefined;
 };
-export type ViewerMembershipLayoutLoader_assumeMutation = {
+export type useAssumeMutation = {
-  response: ViewerMembershipLayoutLoader_assumeMutation$data;
+  response: useAssumeMutation$data;
-  variables: ViewerMembershipLayoutLoader_assumeMutation$variables;
+  variables: useAssumeMutation$variables;
 };
 
 const node: ConcreteRequest = (function(){
@@ -171,7 +171,7 @@ return {
     "argumentDefinitions": (v0/*: any*/),
     "kind": "Fragment",
     "metadata": null,
-    "name": "ViewerMembershipLayoutLoader_assumeMutation",
+    "name": "useAssumeMutation",
     "selections": (v3/*: any*/),
     "type": "Mutation",
     "abstractKey": null
@@ -180,20 +180,20 @@ return {
   "operation": {
     "argumentDefinitions": (v0/*: any*/),
     "kind": "Operation",
-    "name": "ViewerMembershipLayoutLoader_assumeMutation",
+    "name": "useAssumeMutation",
     "selections": (v3/*: any*/)
   },
   "params": {
-    "cacheID": "3c73e9c20476f178c97c687dfd90b7e6",
+    "cacheID": "6c65712730108428976b1d0007ca2993",
     "id": null,
     "metadata": {},
-    "name": "ViewerMembershipLayoutLoader_assumeMutation",
+    "name": "useAssumeMutation",
     "operationKind": "mutation",
-    "text": "mutation ViewerMembershipLayoutLoader_assumeMutation(\n  $input: AssumeOrganizationSessionInput!\n) {\n  assumeOrganizationSession(input: $input) {\n    result {\n      __typename\n      ... on OrganizationSessionCreated {\n        membership {\n          id\n          lastSession {\n            id\n            expiresAt\n          }\n        }\n      }\n      ... on PasswordRequired {\n        reason\n      }\n      ... on SAMLAuthenticationRequired {\n        reason\n        redirectUrl\n      }\n    }\n  }\n}\n"
+    "text": "mutation useAssumeMutation(\n  $input: AssumeOrganizationSessionInput!\n) {\n  assumeOrganizationSession(input: $input) {\n    result {\n      __typename\n      ... on OrganizationSessionCreated {\n        membership {\n          id\n          lastSession {\n            id\n            expiresAt\n          }\n        }\n      }\n      ... on PasswordRequired {\n        reason\n      }\n      ... on SAMLAuthenticationRequired {\n        reason\n        redirectUrl\n      }\n    }\n  }\n}\n"
   }
 };
 })();
 
-(node as any).hash = "d885b41930ea8c7e5c14cadb7ae40eb3";
+(node as any).hash = "1ccd539abfc276084939857e65421be9";
 
 export default node;

apps/console/src/hooks/iam/useAssume.ts

@@ -0,0 +1,84 @@
+import { UnAuthenticatedError } from "@probo/relay";
+import { useEffect } from "react";
+import { useMutation } from "react-relay";
+import { useNavigate } from "react-router";
+import { graphql } from "relay-runtime";
+
+import type { useAssumeMutation } from "#/__generated__/iam/useAssumeMutation.graphql";
+
+import { useOrganizationId } from "../useOrganizationId";
+
+interface UseAssumeParameters {
+  onSuccess: () => void;
+}
+
+const assumeMutation = graphql`
+  mutation useAssumeMutation(
+    $input: AssumeOrganizationSessionInput!
+  ) {
+    assumeOrganizationSession(input: $input) {
+      result {
+        __typename
+        ... on OrganizationSessionCreated {
+          membership {
+            id
+            lastSession {
+              id
+              expiresAt
+            }
+          }
+        }
+        ... on PasswordRequired {
+          reason
+        }
+        ... on SAMLAuthenticationRequired {
+          reason
+          redirectUrl
+        }
+      }
+    }
+  }
+`;
+
+export function useAssume(params: UseAssumeParameters) {
+  const { onSuccess } = params;
+
+  const organizationId = useOrganizationId();
+  const navigate = useNavigate();
+
+  const [assumeOrganizationSession] = useMutation<useAssumeMutation>(assumeMutation);
+
+  useEffect(() => {
+    assumeOrganizationSession({
+      variables: {
+        input: { organizationId },
+      },
+      onError: (error) => {
+        if (error instanceof UnAuthenticatedError) {
+          void navigate("/auth/login");
+          return;
+        }
+      },
+      onCompleted: ({ assumeOrganizationSession }) => {
+        if (!assumeOrganizationSession) {
+          throw new Error("complete mutation result is empty");
+        }
+
+        const { result } = assumeOrganizationSession;
+
+        switch (result.__typename) {
+          case "PasswordRequired":
+            void navigate("/auth/login");
+            break;
+          case "SAMLAuthenticationRequired":
+            window.location.href = result.redirectUrl;
+            break;
+          default:
+            onSuccess();
+        }
+      },
+    });
+  }, [onSuccess, navigate, assumeOrganizationSession, organizationId]);
+
+  return;
+}

apps/console/src/pages/iam/organizations/ViewerMembershipLayoutLoader.tsx

@@ -1,11 +1,9 @@
-import { UnAuthenticatedError } from "@probo/relay";
 import { Skeleton } from "@probo/ui";
-import { Suspense, useEffect } from "react";
+import { Suspense, useCallback } from "react";
-import { graphql, useMutation, useQueryLoader } from "react-relay";
+import { useQueryLoader } from "react-relay";
-import { useNavigate } from "react-router";
 
-import type { ViewerMembershipLayoutLoader_assumeMutation } from "#/__generated__/iam/ViewerMembershipLayoutLoader_assumeMutation.graphql";
 import type { ViewerMembershipLayoutQuery } from "#/__generated__/iam/ViewerMembershipLayoutQuery.graphql";
+import { useAssume } from "#/hooks/iam/useAssume";
 import { useOrganizationId } from "#/hooks/useOrganizationId";
 import { IAMRelayProvider } from "#/providers/IAMRelayProvider";
 
@@ -14,80 +12,23 @@ import {
   viewerMembershipLayoutQuery,
 } from "./ViewerMembershipLayout";
 
-const ensureAssumingMutation = graphql`
-  mutation ViewerMembershipLayoutLoader_assumeMutation(
-    $input: AssumeOrganizationSessionInput!
-  ) {
-    assumeOrganizationSession(input: $input) {
-      result {
-        __typename
-        ... on OrganizationSessionCreated {
-          membership {
-            id
-            lastSession {
-              id
-              expiresAt
-            }
-          }
-        }
-        ... on PasswordRequired {
-          reason
-        }
-        ... on SAMLAuthenticationRequired {
-          reason
-          redirectUrl
-        }
-      }
-    }
-  }
-`;
-
 function ViewerMembershipLayoutQueryLoader() {
   const organizationId = useOrganizationId();
-  const navigate = useNavigate();
 
-  const [assumeOrganizationSession]
-    = useMutation<ViewerMembershipLayoutLoader_assumeMutation>(
-      ensureAssumingMutation,
-    );
   const [queryRef, loadQuery] = useQueryLoader<ViewerMembershipLayoutQuery>(
     viewerMembershipLayoutQuery,
   );
 
-  useEffect(() => {
+  const onAssumeSuccess = useCallback(
-    assumeOrganizationSession({
+    () =>
-      variables: {
+      loadQuery({
-        input: { organizationId },
+        organizationId,
-      },
+        hideSidebar: false,
-      onError: (error) => {
+      }),
-        if (error instanceof UnAuthenticatedError) {
+    [loadQuery, organizationId],
-          void navigate("/auth/login");
+  );
-          return;
-        }
-      },
-      onCompleted: ({ assumeOrganizationSession }) => {
-        if (!assumeOrganizationSession) {
-          throw new Error("complete mutation result is empty");
-        }
 
-        const { result } = assumeOrganizationSession;
+  useAssume({ onSuccess: onAssumeSuccess });
-
-        switch (result.__typename) {
-          case "PasswordRequired":
-            void navigate("/auth/login");
-            break;
-          case "SAMLAuthenticationRequired":
-            window.location.href = result.redirectUrl;
-            break;
-          default:
-            loadQuery({
-              organizationId,
-              hideSidebar: false,
-            });
-        }
-      },
-    });
-  }, [navigate, assumeOrganizationSession, loadQuery, organizationId]);
 
   if (!queryRef) {
     return <Skeleton className="w-full h-screen" />;

apps/console/src/pages/organizations/employee/EmployeeLayoutLoader.tsx

@@ -1,8 +1,9 @@
 import { Skeleton } from "@probo/ui";
-import { Suspense, useEffect } from "react";
+import { Suspense, useCallback } from "react";
 import { useQueryLoader } from "react-relay";
 
 import type { ViewerMembershipLayoutQuery } from "#/__generated__/iam/ViewerMembershipLayoutQuery.graphql";
+import { useAssume } from "#/hooks/iam/useAssume";
 import { useOrganizationId } from "#/hooks/useOrganizationId";
 import { IAMRelayProvider } from "#/providers/IAMRelayProvider";
 
@@ -17,12 +18,16 @@ function EmployeeLayoutQueryLoader() {
     viewerMembershipLayoutQuery,
   );
 
-  useEffect(() => {
+  const onAssumeSuccess = useCallback(
-    loadQuery({
+    () =>
-      organizationId,
+      loadQuery({
-      hideSidebar: true,
+        organizationId,
-    });
+        hideSidebar: false,
-  }, [loadQuery, organizationId]);
+      }),
+    [loadQuery, organizationId],
+  );
+
+  useAssume({ onSuccess: onAssumeSuccess });
 
   if (!queryRef) {
     return <Skeleton className="w-full h-screen" />;

apps/console/vite.config.ts

@@ -8,7 +8,11 @@ import { defineConfig } from "vite";
 export default defineConfig({
   plugins: [
     react({
-      exclude: ["src/pages/iam/**/*", "src/components/connectors/**/*"],
+      exclude: [
+        "src/hooks/iam/**/*",
+        "src/pages/iam/**/*",
+        "src/components/connectors/**/*",
+      ],
       babel: {
         plugins: [
           [
@@ -22,7 +26,11 @@ export default defineConfig({
       },
     }),
     react({
-      include: ["src/pages/iam/**/*", "src/components/connectors/**/*"],
+      include: [
+        "src/hooks/iam/**/*",
+        "src/pages/iam/**/*",
+        "src/components/connectors/**/*",
+      ],
       babel: {
         plugins: [
           [