Assume organization when loading org layout

Signed-off-by: Émile Ré <emile@getprobo.com>
2026-02-09 15:29:04 +04:00
parent f174a508a6
commit 1e1ec72ee7
4 changed files with 108 additions and 103 deletions
--- a/apps/console/src/generated/iam/ViewerMembershipLayoutLoader_assumeMutation.graphql.ts
+++ b/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 = {
-  response: MembershipCard_assumeMutation$data;
-  variables: MembershipCard_assumeMutation$variables;
+export type ViewerMembershipLayoutLoader_assumeMutation = {
+  response: ViewerMembershipLayoutLoader_assumeMutation$data;
+  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;
--- a/apps/console/src/pages/iam/memberships/_components/MembershipCard.tsx
+++ b/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, useMutation } from "react-relay";
-import { Link, useNavigate } from "react-router";
+import { useFragment } from "react-relay";
+import { Link } 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}`}>
-                  <Button variant="secondary">{__("Start")}</Button>
-                </Link>
-              )
-            : (
-                <Button onClick={handleAssumeOrganizationSession}>
-                  {__("Login")}
-                </Button>
-              )}
+          <Link to={`/organizations/${organization.id}`}>
+            {isAssuming
+              ? <Button variant="secondary">{__("Start")}</Button>
+              : <Button>{__("Login")}</Button>}
+          </Link>
        </div>
      </div>
    </Card>
--- a/apps/console/src/pages/iam/organizations/ViewerMembershipLayoutLoader.tsx
+++ b/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({
-      organizationId,
-      hideSidebar: false,
+    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:
+            loadQuery({
+              organizationId,
+              hideSidebar: false,
+            });
+        }
+      },
    });
-  }, [loadQuery, organizationId]);
+  }, [navigate, assumeOrganizationSession, loadQuery, organizationId]);

  if (!queryRef) {
    return <Skeleton className="w-full h-screen" />;
--- a/pkg/iam/session_service.go
+++ b/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 != nil {
+			if err := rootSession.LoadByID(ctx, tx, sessionID); 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 != nil {
+			if err := identity.LoadByID(ctx, tx, rootSession.IdentityID); err != nil {
 				return fmt.Errorf("cannot load identity: %w", err)
 			}

-			err = membership.LoadByIdentityInOrganization(ctx, tx, rootSession.IdentityID, organizationID)
-			if err != nil {
+			if err := membership.LoadByIdentityInOrganization(ctx, tx, rootSession.IdentityID, organizationID); 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,