Assume organization when loading org layout

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

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

@@ -1,5 +1,5 @@
 /**
- * @generated SignedSource<<0f048726f7e7371223919d924fa15d54>>
+ * @generated SignedSource<<09963c0483a8416abf9e411f4a0b716b>>
  * @lightSyntaxTransform
  * @nogrep
  */
@@ -13,10 +13,10 @@ export type ReauthenticationReason = "POLICY_REQUIREMENT" | "SENSITIVE_ACTION" |
 export type AssumeOrganizationSessionInput = {
   organizationId: string;
 };
-export type MembershipCard_assumeMutation$variables = {
+export type ViewerMembershipLayoutLoader_assumeMutation$variables = {
   input: AssumeOrganizationSessionInput;
 };
-export type MembershipCard_assumeMutation$data = {
+export type ViewerMembershipLayoutLoader_assumeMutation$data = {
   readonly assumeOrganizationSession: {
     readonly result: {
       readonly __typename: "OrganizationSessionCreated";
@@ -41,9 +41,9 @@ export type MembershipCard_assumeMutation$data = {
     };
   } | null | undefined;
 };
-export type MembershipCard_assumeMutation = {
+export type ViewerMembershipLayoutLoader_assumeMutation = {
-  response: MembershipCard_assumeMutation$data;
+  response: ViewerMembershipLayoutLoader_assumeMutation$data;
-  variables: MembershipCard_assumeMutation$variables;
+  variables: ViewerMembershipLayoutLoader_assumeMutation$variables;
 };
 
 const node: ConcreteRequest = (function(){
@@ -171,7 +171,7 @@ return {
     "argumentDefinitions": (v0/*: any*/),
     "kind": "Fragment",
     "metadata": null,
-    "name": "MembershipCard_assumeMutation",
+    "name": "ViewerMembershipLayoutLoader_assumeMutation",
     "selections": (v3/*: any*/),
     "type": "Mutation",
     "abstractKey": null
@@ -180,20 +180,20 @@ return {
   "operation": {
     "argumentDefinitions": (v0/*: any*/),
     "kind": "Operation",
-    "name": "MembershipCard_assumeMutation",
+    "name": "ViewerMembershipLayoutLoader_assumeMutation",
     "selections": (v3/*: any*/)
   },
   "params": {
-    "cacheID": "d106c9d8a5a2c3256da2f1f83fb8fde7",
+    "cacheID": "3c73e9c20476f178c97c687dfd90b7e6",
     "id": null,
     "metadata": {},
-    "name": "MembershipCard_assumeMutation",
+    "name": "ViewerMembershipLayoutLoader_assumeMutation",
     "operationKind": "mutation",
-    "text": "mutation MembershipCard_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 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"
   }
 };
 })();
 
-(node as any).hash = "7acaa03357731962d7d29eb2c046660a";
+(node as any).hash = "d885b41930ea8c7e5c14cadb7ae40eb3";
 
 export default node;

apps/console/src/pages/iam/memberships/_components/MembershipCard.tsx

@@ -9,12 +9,10 @@ import {
   IconClock,
   IconLock,
 } from "@probo/ui";
-import { useCallback } from "react";
+import { useFragment } from "react-relay";
-import { useFragment, useMutation } from "react-relay";
+import { Link } from "react-router";
-import { Link, useNavigate } from "react-router";
 import { graphql } from "relay-runtime";
 
-import type { MembershipCard_assumeMutation } from "#/__generated__/iam/MembershipCard_assumeMutation.graphql";
 import type { MembershipCardFragment$key } from "#/__generated__/iam/MembershipCardFragment.graphql";
 
 const fragment = graphql`
@@ -31,34 +29,6 @@ const fragment = graphql`
   }
 `;
 
-const assumeOrganizationSessionMutation = graphql`
-  mutation MembershipCard_assumeMutation(
-    $input: AssumeOrganizationSessionInput!
-  ) {
-    assumeOrganizationSession(input: $input) {
-      result {
-        __typename
-        ... on OrganizationSessionCreated {
-          membership {
-            id
-            lastSession {
-              id
-              expiresAt
-            }
-          }
-        }
-        ... on PasswordRequired {
-          reason
-        }
-        ... on SAMLAuthenticationRequired {
-          reason
-          redirectUrl
-        }
-      }
-    }
-  }
-`;
-
 interface MembershipCardProps {
   fKey: MembershipCardFragment$key;
 }
@@ -66,51 +36,17 @@ interface MembershipCardProps {
 export function MembershipCard(props: MembershipCardProps) {
   const { fKey } = props;
   const { __ } = useTranslate();
-  const navigate = useNavigate();
 
   const { lastSession, organization } = useFragment<MembershipCardFragment$key>(
     fragment,
     fKey,
   );
-  const isAuthenticated = !!lastSession;
+  const isAssuming = !!lastSession;
   const isExpired
     = lastSession && parseDate(lastSession.expiresAt) < new Date();
 
-  const [assumeOrganizationSession]
-    = useMutation<MembershipCard_assumeMutation>(
-      assumeOrganizationSessionMutation,
-    );
-
-  const handleAssumeOrganizationSession = useCallback(() => {
-    assumeOrganizationSession({
-      variables: {
-        input: {
-          organizationId: organization.id,
-        },
-      },
-      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:
-            void navigate(`/organizations/${organization.id}`);
-        }
-      },
-    });
-  }, [assumeOrganizationSession, navigate, organization.id]);
-
   const getAuthBadge = () => {
-    if (isAuthenticated) {
+    if (isAssuming) {
       return (
         <Badge variant="success" className="flex items-center gap-1">
           <IconCheckmark1 size={14} />
@@ -149,17 +85,11 @@ export function MembershipCard(props: MembershipCardProps) {
           </div>
         </div>
         <div className="flex items-center gap-3">
-          {isAuthenticated
+          <Link to={`/organizations/${organization.id}`}>
-            ? (
+            {isAssuming
-                <Link to={`/organizations/${organization.id}`}>
+              ? <Button variant="secondary">{__("Start")}</Button>
-                  <Button variant="secondary">{__("Start")}</Button>
+              : <Button>{__("Login")}</Button>}
-                </Link>
+          </Link>
-              )
-            : (
-                <Button onClick={handleAssumeOrganizationSession}>
-                  {__("Login")}
-                </Button>
-              )}
         </div>
       </div>
     </Card>

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

@@ -1,7 +1,10 @@
+import { UnAuthenticatedError } from "@probo/relay";
 import { Skeleton } from "@probo/ui";
 import { Suspense, useEffect } from "react";
-import { useQueryLoader } from "react-relay";
+import { graphql, useMutation, 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 { useOrganizationId } from "#/hooks/useOrganizationId";
 import { IAMRelayProvider } from "#/providers/IAMRelayProvider";
@@ -11,18 +14,80 @@ 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(() => {
-    loadQuery({
+    assumeOrganizationSession({
-      organizationId,
+      variables: {
-      hideSidebar: false,
+        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:
+            loadQuery({
+              organizationId,
+              hideSidebar: false,
+            });
+        }
+      },
     });
-  }, [loadQuery, organizationId]);
+  }, [navigate, assumeOrganizationSession, loadQuery, organizationId]);
 
   if (!queryRef) {
     return <Skeleton className="w-full h-screen" />;

pkg/iam/session_service.go

@@ -417,8 +417,7 @@ func (s SessionService) AssumeOrganizationSession(
 	err := s.pg.WithTx(
 		ctx,
 		func(tx pg.Conn) error {
-			err := rootSession.LoadByID(ctx, tx, sessionID)
+			if err := rootSession.LoadByID(ctx, tx, sessionID); err != nil {
-			if err != nil {
 				if err == coredata.ErrResourceNotFound {
 					return NewSessionNotFoundError(sessionID)
 				}
@@ -433,13 +432,11 @@ func (s SessionService) AssumeOrganizationSession(
 				return NewSessionExpiredError(sessionID)
 			}
 
-			err = identity.LoadByID(ctx, tx, rootSession.IdentityID)
+			if err := identity.LoadByID(ctx, tx, rootSession.IdentityID); err != nil {
-			if err != nil {
 				return fmt.Errorf("cannot load identity: %w", err)
 			}
 
-			err = membership.LoadByIdentityInOrganization(ctx, tx, rootSession.IdentityID, organizationID)
+			if err := membership.LoadByIdentityInOrganization(ctx, tx, rootSession.IdentityID, organizationID); err != nil {
-			if err != nil {
 				if err == coredata.ErrResourceNotFound {
 					return NewMembershipNotFoundError(organizationID)
 				}
@@ -450,8 +447,21 @@ func (s SessionService) AssumeOrganizationSession(
 				return NewMembershipInactiveError(membership.ID)
 			}
 
+			// If child session already exists use it
+			if err := childSession.LoadByRootSessionIDAndMembershipID(ctx, tx, rootSession.IdentityID, membership.ID); err == nil {
+				if childSession.ExpireReason != nil || now.After(childSession.ExpiredAt) {
+					return NewSessionExpiredError(childSession.ID)
+				}
+
+				return nil
+			} else {
+				if err != coredata.ErrResourceNotFound {
+					return fmt.Errorf("cannot load child session: %w", err)
+				}
+			}
+
 			samlConfig := &coredata.SAMLConfiguration{}
-			err = samlConfig.LoadByOrganizationIDAndEmailDomain(
+			err := samlConfig.LoadByOrganizationIDAndEmailDomain(
 				ctx,
 				tx,
 				scope,