Plug assume organization session

Signed-off-by: Émile Ré <nemile.re@gmail.com>

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

@@ -8,11 +8,13 @@ import {
   IconClock,
   IconLock,
 } from "@probo/ui";
-import { Link } from "react-router";
+import { useNavigate } from "react-router";
 import { graphql } from "relay-runtime";
-import { useFragment } from "react-relay";
+import { useFragment, useMutation } from "react-relay";
 import type { MembershipCardFragment$key } from "./__generated__/MembershipCardFragment.graphql";
 import { parseDate } from "@probo/helpers";
+import { useCallback } from "react";
+import type { MembershipCard_assumeMutation } from "./__generated__/MembershipCard_assumeMutation.graphql";
 
 const fragment = graphql`
   fragment MembershipCardFragment on Membership {
@@ -28,6 +30,25 @@ const fragment = graphql`
   }
 `;
 
+const assumeOrganizationSessionMutation = graphql`
+  mutation MembershipCard_assumeMutation(
+    $input: AssumeOrganizationSessionInput!
+  ) {
+    assumeOrganizationSession(input: $input) {
+      result {
+        __typename
+        ... on PasswordRequired {
+          reason
+        }
+        ... on SAMLAuthenticationRequired {
+          reason
+          redirectUrl
+        }
+      }
+    }
+  }
+`;
+
 interface MembershipCardProps {
   fKey: MembershipCardFragment$key;
 }
@@ -35,6 +56,7 @@ interface MembershipCardProps {
 export function MembershipCard(props: MembershipCardProps) {
   const { fKey } = props;
   const { __ } = useTranslate();
+  const navigate = useNavigate();
 
   const { lastSession, organization } = useFragment<MembershipCardFragment$key>(
     fragment,
@@ -42,13 +64,40 @@ export function MembershipCard(props: MembershipCardProps) {
   );
   const isAuthenticated = !!lastSession;
   const isExpired =
-    lastSession && parseDate(lastSession.expiresAt) >= new Date();
+    lastSession && parseDate(lastSession.expiresAt) < new Date();
 
-  // Determine target URL and button text based on auth status
+  const [assumeOrganizationSession] =
-  // const targetUrl = isAuthenticated
+    useMutation<MembershipCard_assumeMutation>(
-  //   ? `/organizations/${organization.id}`
+      assumeOrganizationSessionMutation,
-  //   : organization.loginUrl;
+    );
-  const targetUrl = `/organizations/${organization.id}`;
+
+  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":
+            navigate("auth/login");
+            break;
+          case "SAMLAuthenticationRequired":
+            window.location.href = result.redirectUrl;
+            break;
+          default:
+            navigate(`/organizations/${organization.id}`);
+        }
+      },
+    });
+  }, [assumeOrganizationSession, navigate, organization.id]);
 
   const getAuthBadge = () => {
     if (isAuthenticated) {
@@ -75,58 +124,23 @@ export function MembershipCard(props: MembershipCardProps) {
     }
   };
 
-  // const getButtonText = () => {
-  //   if (isAuthenticated) return __("Select");
-  //   if (organization.authenticationMethod === "saml")
-  //     return __("Login with SAML");
-  //   return __("Login");
-  // };
-
-  // Check if the URL is a backend SAML endpoint
-  const isSAMLUrl = targetUrl.includes("/connect/saml/");
-
   return (
     <Card padded className="w-full">
       <div className="flex items-center justify-between">
-        {isSAMLUrl ? (
+        <div className="flex items-center gap-4 hover:text-primary flex-1">
-          <a
+          <Avatar
-            href={targetUrl}
+            src={organization.logoUrl}
-            className="flex items-center gap-4 hover:text-primary flex-1"
+            name={organization.name}
-          >
+            size="l"
-            <Avatar
+          />
-              src={organization.logoUrl}
+          <div className="flex flex-col gap-1">
-              name={organization.name}
+            <h2 className="font-semibold text-xl">{organization.name}</h2>
-              size="l"
+            {getAuthBadge()}
-            />
+          </div>
-            <div className="flex flex-col gap-1">
+        </div>
-              <h2 className="font-semibold text-xl">{organization.name}</h2>
-              {getAuthBadge()}
-            </div>
-          </a>
-        ) : (
-          <Link
-            to={targetUrl}
-            className="flex items-center gap-4 hover:text-primary flex-1"
-          >
-            <Avatar
-              src={organization.logoUrl}
-              name={organization.name}
-              size="l"
-            />
-            <div className="flex flex-col gap-1">
-              <h2 className="font-semibold text-xl">{organization.name}</h2>
-              {getAuthBadge()}
-            </div>
-          </Link>
-        )}
         <div className="flex items-center gap-3">
-          <Button asChild>
+          <Button onClick={handleAssumeOrganizationSession}>
-            {/* {isSAMLUrl ? (
+            {__("Login")}
-              <a href={targetUrl}>{getButtonText()}</a>
-            ) : (
-              <Link to={targetUrl}>{getButtonText()}</Link>
-            )} */}
-            <Link to={targetUrl}>LOGIN</Link>
           </Button>
         </div>
       </div>

apps/console/src/pages/iam/organizations/_components/OrganizationDropdownMenuItem.tsx

@@ -7,8 +7,11 @@ import {
 } from "@probo/ui";
 import { graphql } from "relay-runtime";
 import type { OrganizationDropdownMenuItemFragment$key } from "./__generated__/OrganizationDropdownMenuItemFragment.graphql";
-import { useFragment } from "react-relay";
+import { useFragment, useMutation } from "react-relay";
 import { parseDate } from "@probo/helpers";
+import { useCallback } from "react";
+import { useNavigate } from "react-router";
+import type { OrganizationDropdownMenuItem_assumeMutation } from "./__generated__/OrganizationDropdownMenuItem_assumeMutation.graphql";
 
 const fragment = graphql`
   fragment OrganizationDropdownMenuItemFragment on Membership {
@@ -18,25 +21,81 @@ const fragment = graphql`
       expiresAt
     }
     organization @required(action: THROW) {
+      id
       logoUrl
       name
     }
   }
 `;
 
+const assumeOrganizationSessionMutation = graphql`
+  mutation OrganizationDropdownMenuItem_assumeMutation(
+    $input: AssumeOrganizationSessionInput!
+  ) {
+    assumeOrganizationSession(input: $input) {
+      result {
+        __typename
+        ... on PasswordRequired {
+          reason
+        }
+        ... on SAMLAuthenticationRequired {
+          reason
+          redirectUrl
+        }
+      }
+    }
+  }
+`;
+
 export function OrganizationDropdownMenuItem(props: {
   fKey: OrganizationDropdownMenuItemFragment$key;
 }) {
   const { fKey } = props;
+
+  const navigate = useNavigate();
+
   const { id, lastSession, organization } =
     useFragment<OrganizationDropdownMenuItemFragment$key>(fragment, fKey);
 
   const isAuthenticated = !!lastSession;
   const isExpired =
-    lastSession && parseDate(lastSession.expiresAt) >= new Date();
+    lastSession && parseDate(lastSession.expiresAt) < new Date();
+
+  const [assumeOrganizationSession] =
+    useMutation<OrganizationDropdownMenuItem_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":
+            navigate("auth/login");
+            break;
+          case "SAMLAuthenticationRequired":
+            window.location.href = result.redirectUrl;
+            break;
+          default:
+            navigate(`/organizations/${organization.id}`);
+        }
+      },
+    });
+  }, [assumeOrganizationSession, navigate, organization.id]);
 
   return (
-    <DropdownItem key={id}>
+    <DropdownItem key={id} onClick={handleAssumeOrganizationSession}>
       {/* TODO add link or anchor */}
       <Avatar name={organization.name} src={organization.logoUrl} />
       <span className="flex-1">{organization.name}</span>

apps/console/src/pages/iam/organizations/_components/__generated__/MembershipCard_assumeMutation.graphql.ts

@@ -0,0 +1,145 @@
+/**
+ * @generated SignedSource<<2574b3b6a557336992dff4f923fc0870>>
+ * @lightSyntaxTransform
+ * @nogrep
+ */
+
+/* tslint:disable */
+/* eslint-disable */
+// @ts-nocheck
+
+import { ConcreteRequest } from 'relay-runtime';
+export type ReauthenticationReason = "POLICY_REQUIREMENT" | "SENSITIVE_ACTION" | "SESSION_EXPIRED";
+export type AssumeOrganizationSessionInput = {
+  organizationId: string;
+};
+export type MembershipCard_assumeMutation$variables = {
+  input: AssumeOrganizationSessionInput;
+};
+export type MembershipCard_assumeMutation$data = {
+  readonly assumeOrganizationSession: {
+    readonly result: {
+      readonly __typename: "PasswordRequired";
+      readonly reason: ReauthenticationReason;
+    } | {
+      readonly __typename: "SAMLAuthenticationRequired";
+      readonly reason: ReauthenticationReason;
+      readonly redirectUrl: string;
+    } | {
+      // This will never be '%other', but we need some
+      // value in case none of the concrete values match.
+      readonly __typename: "%other";
+    };
+  } | null | undefined;
+};
+export type MembershipCard_assumeMutation = {
+  response: MembershipCard_assumeMutation$data;
+  variables: MembershipCard_assumeMutation$variables;
+};
+
+const node: ConcreteRequest = (function(){
+var v0 = [
+  {
+    "defaultValue": null,
+    "kind": "LocalArgument",
+    "name": "input"
+  }
+],
+v1 = {
+  "alias": null,
+  "args": null,
+  "kind": "ScalarField",
+  "name": "reason",
+  "storageKey": null
+},
+v2 = [
+  {
+    "alias": null,
+    "args": [
+      {
+        "kind": "Variable",
+        "name": "input",
+        "variableName": "input"
+      }
+    ],
+    "concreteType": "AssumeOrganizationSessionPayload",
+    "kind": "LinkedField",
+    "name": "assumeOrganizationSession",
+    "plural": false,
+    "selections": [
+      {
+        "alias": null,
+        "args": null,
+        "concreteType": null,
+        "kind": "LinkedField",
+        "name": "result",
+        "plural": false,
+        "selections": [
+          {
+            "alias": null,
+            "args": null,
+            "kind": "ScalarField",
+            "name": "__typename",
+            "storageKey": null
+          },
+          {
+            "kind": "InlineFragment",
+            "selections": [
+              (v1/*: any*/)
+            ],
+            "type": "PasswordRequired",
+            "abstractKey": null
+          },
+          {
+            "kind": "InlineFragment",
+            "selections": [
+              (v1/*: any*/),
+              {
+                "alias": null,
+                "args": null,
+                "kind": "ScalarField",
+                "name": "redirectUrl",
+                "storageKey": null
+              }
+            ],
+            "type": "SAMLAuthenticationRequired",
+            "abstractKey": null
+          }
+        ],
+        "storageKey": null
+      }
+    ],
+    "storageKey": null
+  }
+];
+return {
+  "fragment": {
+    "argumentDefinitions": (v0/*: any*/),
+    "kind": "Fragment",
+    "metadata": null,
+    "name": "MembershipCard_assumeMutation",
+    "selections": (v2/*: any*/),
+    "type": "Mutation",
+    "abstractKey": null
+  },
+  "kind": "Request",
+  "operation": {
+    "argumentDefinitions": (v0/*: any*/),
+    "kind": "Operation",
+    "name": "MembershipCard_assumeMutation",
+    "selections": (v2/*: any*/)
+  },
+  "params": {
+    "cacheID": "0d93c6ee32107a4263a910e21597e982",
+    "id": null,
+    "metadata": {},
+    "name": "MembershipCard_assumeMutation",
+    "operationKind": "mutation",
+    "text": "mutation MembershipCard_assumeMutation(\n  $input: AssumeOrganizationSessionInput!\n) {\n  assumeOrganizationSession(input: $input) {\n    result {\n      __typename\n      ... on PasswordRequired {\n        reason\n      }\n      ... on SAMLAuthenticationRequired {\n        reason\n        redirectUrl\n      }\n    }\n  }\n}\n"
+  }
+};
+})();
+
+(node as any).hash = "aa7a5daec98d5fa51a56da0a7e0b2178";
+
+export default node;

apps/console/src/pages/iam/organizations/_components/__generated__/OrganizationDropdownMenuItemFragment.graphql.ts

@@ -1,5 +1,5 @@
 /**
- * @generated SignedSource<<ef66c9eab961d5c3c28669629b82cfd6>>
+ * @generated SignedSource<<e917cb1af13350772d46c02bd6b37da5>>
  * @lightSyntaxTransform
  * @nogrep
  */
@@ -17,6 +17,7 @@ export type OrganizationDropdownMenuItemFragment$data = {
     readonly id: string;
   } | null | undefined;
   readonly organization: {
+    readonly id: string;
     readonly logoUrl: string | null | undefined;
     readonly name: string;
   };
@@ -71,6 +72,7 @@ return {
         "name": "organization",
         "plural": false,
         "selections": [
+          (v0/*: any*/),
           {
             "alias": null,
             "args": null,
@@ -96,6 +98,6 @@ return {
 };
 })();
 
-(node as any).hash = "6be7f721fd63d6d758b479e58c68b754";
+(node as any).hash = "22fea7c8463774a5fa9ea5f56d1e876c";
 
 export default node;

apps/console/src/pages/iam/organizations/_components/__generated__/OrganizationDropdownMenuItem_assumeMutation.graphql.ts

@@ -0,0 +1,145 @@
+/**
+ * @generated SignedSource<<469707d0eca0a0378fa589809db964bc>>
+ * @lightSyntaxTransform
+ * @nogrep
+ */
+
+/* tslint:disable */
+/* eslint-disable */
+// @ts-nocheck
+
+import { ConcreteRequest } from 'relay-runtime';
+export type ReauthenticationReason = "POLICY_REQUIREMENT" | "SENSITIVE_ACTION" | "SESSION_EXPIRED";
+export type AssumeOrganizationSessionInput = {
+  organizationId: string;
+};
+export type OrganizationDropdownMenuItem_assumeMutation$variables = {
+  input: AssumeOrganizationSessionInput;
+};
+export type OrganizationDropdownMenuItem_assumeMutation$data = {
+  readonly assumeOrganizationSession: {
+    readonly result: {
+      readonly __typename: "PasswordRequired";
+      readonly reason: ReauthenticationReason;
+    } | {
+      readonly __typename: "SAMLAuthenticationRequired";
+      readonly reason: ReauthenticationReason;
+      readonly redirectUrl: string;
+    } | {
+      // This will never be '%other', but we need some
+      // value in case none of the concrete values match.
+      readonly __typename: "%other";
+    };
+  } | null | undefined;
+};
+export type OrganizationDropdownMenuItem_assumeMutation = {
+  response: OrganizationDropdownMenuItem_assumeMutation$data;
+  variables: OrganizationDropdownMenuItem_assumeMutation$variables;
+};
+
+const node: ConcreteRequest = (function(){
+var v0 = [
+  {
+    "defaultValue": null,
+    "kind": "LocalArgument",
+    "name": "input"
+  }
+],
+v1 = {
+  "alias": null,
+  "args": null,
+  "kind": "ScalarField",
+  "name": "reason",
+  "storageKey": null
+},
+v2 = [
+  {
+    "alias": null,
+    "args": [
+      {
+        "kind": "Variable",
+        "name": "input",
+        "variableName": "input"
+      }
+    ],
+    "concreteType": "AssumeOrganizationSessionPayload",
+    "kind": "LinkedField",
+    "name": "assumeOrganizationSession",
+    "plural": false,
+    "selections": [
+      {
+        "alias": null,
+        "args": null,
+        "concreteType": null,
+        "kind": "LinkedField",
+        "name": "result",
+        "plural": false,
+        "selections": [
+          {
+            "alias": null,
+            "args": null,
+            "kind": "ScalarField",
+            "name": "__typename",
+            "storageKey": null
+          },
+          {
+            "kind": "InlineFragment",
+            "selections": [
+              (v1/*: any*/)
+            ],
+            "type": "PasswordRequired",
+            "abstractKey": null
+          },
+          {
+            "kind": "InlineFragment",
+            "selections": [
+              (v1/*: any*/),
+              {
+                "alias": null,
+                "args": null,
+                "kind": "ScalarField",
+                "name": "redirectUrl",
+                "storageKey": null
+              }
+            ],
+            "type": "SAMLAuthenticationRequired",
+            "abstractKey": null
+          }
+        ],
+        "storageKey": null
+      }
+    ],
+    "storageKey": null
+  }
+];
+return {
+  "fragment": {
+    "argumentDefinitions": (v0/*: any*/),
+    "kind": "Fragment",
+    "metadata": null,
+    "name": "OrganizationDropdownMenuItem_assumeMutation",
+    "selections": (v2/*: any*/),
+    "type": "Mutation",
+    "abstractKey": null
+  },
+  "kind": "Request",
+  "operation": {
+    "argumentDefinitions": (v0/*: any*/),
+    "kind": "Operation",
+    "name": "OrganizationDropdownMenuItem_assumeMutation",
+    "selections": (v2/*: any*/)
+  },
+  "params": {
+    "cacheID": "024d06c3bba88bb18eb1aa7e2888edd0",
+    "id": null,
+    "metadata": {},
+    "name": "OrganizationDropdownMenuItem_assumeMutation",
+    "operationKind": "mutation",
+    "text": "mutation OrganizationDropdownMenuItem_assumeMutation(\n  $input: AssumeOrganizationSessionInput!\n) {\n  assumeOrganizationSession(input: $input) {\n    result {\n      __typename\n      ... on PasswordRequired {\n        reason\n      }\n      ... on SAMLAuthenticationRequired {\n        reason\n        redirectUrl\n      }\n    }\n  }\n}\n"
+  }
+};
+})();
+
+(node as any).hash = "bbc8443e4a756b37f8b46ee98972bdcb";
+
+export default node;

apps/console/src/pages/iam/organizations/_components/__generated__/OrganizationDropdownMenuQuery.graphql.ts

@@ -1,5 +1,5 @@
 /**
- * @generated SignedSource<<2bde2fba4e4897436af08d9dd700b265>>
+ * @generated SignedSource<<034b9a96de4b1efda9529e8165779342>>
  * @lightSyntaxTransform
  * @nogrep
  */
@@ -248,12 +248,12 @@ return {
     ]
   },
   "params": {
-    "cacheID": "8338998501d59b0036d1eaf98dff1ffb",
+    "cacheID": "b880e24e59e2bb500bbb2b9026423c20",
     "id": null,
     "metadata": {},
     "name": "OrganizationDropdownMenuQuery",
     "operationKind": "query",
-    "text": "query OrganizationDropdownMenuQuery {\n  viewer {\n    memberships(first: 1000, orderBy: {direction: DESC, field: CREATED_AT}) {\n      edges {\n        node {\n          id\n          organization {\n            name\n            id\n          }\n          ...OrganizationDropdownMenuItemFragment\n        }\n      }\n    }\n    id\n  }\n}\n\nfragment OrganizationDropdownMenuItemFragment on Membership {\n  id\n  lastSession {\n    id\n    expiresAt\n  }\n  organization {\n    logoUrl\n    name\n    id\n  }\n}\n"
+    "text": "query OrganizationDropdownMenuQuery {\n  viewer {\n    memberships(first: 1000, orderBy: {direction: DESC, field: CREATED_AT}) {\n      edges {\n        node {\n          id\n          organization {\n            name\n            id\n          }\n          ...OrganizationDropdownMenuItemFragment\n        }\n      }\n    }\n    id\n  }\n}\n\nfragment OrganizationDropdownMenuItemFragment on Membership {\n  id\n  lastSession {\n    id\n    expiresAt\n  }\n  organization {\n    id\n    logoUrl\n    name\n  }\n}\n"
   }
 };
 })();

pkg/coredata/membership.go

@@ -62,18 +62,32 @@ func (m Membership) CursorKey(orderBy MembershipOrderField) page.CursorKey {
 
 func (m *Membership) LoadByIdentityInOrganization(ctx context.Context, conn pg.Conn, identityID gid.GID, organizationID gid.GID) error {
 	q := `
+WITH mbr AS (
+    SELECT
+        id,
+        identity_id,
+        organization_id,
+        role,
+        created_at,
+        updated_at
+    FROM
+        iam_memberships
+    WHERE
+        identity_id = @identity_id
+        AND organization_id = @organization_id
+)
 SELECT
-		id,
+    mbr.id,
-		identity_id,
+    mbr.identity_id,
-		organization_id,
+    mbr.organization_id,
-		role,
+    mbr.role,
-		created_at,
+    i.fullname AS full_name,
-		updated_at
+    i.email_address,
+    mbr.created_at,
+    mbr.updated_at
 FROM
-	iam_memberships
+    mbr
-WHERE
+JOIN identities i ON mbr.identity_id = i.id
-	identity_id = @identity_id
-	AND organization_id = @organization_id
 `
 
 	args := pgx.StrictNamedArgs{
@@ -157,18 +171,18 @@ func (m *Membership) LoadByID(
 ) error {
 	query := `
 WITH mbr AS (
-	SELECT
+    SELECT
-		id,
+        id,
-		identity_id,
+        identity_id,
-		organization_id,
+        organization_id,
-		role,
+        role,
-		created_at,
+        created_at,
-		updated_at
+        updated_at
-	FROM
+    FROM
-		iam_memberships
+        iam_memberships
-	WHERE
+    WHERE
-		id = @membership_id
+        id = @membership_id
-		AND %s
+        AND %s
 )
 SELECT
     mbr.id,
@@ -237,19 +251,19 @@ func (m *Membership) LoadRoleByIdentityAndEntityID(
 
 	query := fmt.Sprintf(`
 SELECT
-	m.id,
+    m.id,
-	m.identity_id,
+    m.identity_id,
-	m.organization_id,
+    m.organization_id,
-	m.role,
+    m.role,
-	m.created_at,
+    m.created_at,
-	m.updated_at
+    m.updated_at
 FROM
-	iam_memberships m
+    iam_memberships m
-	INNER JOIN %s e ON e.id = @entity_id
+    INNER JOIN %s e ON e.id = @entity_id
 WHERE
-	%s
+    %s
-	AND m.identity_id = @identity_id
+    AND m.identity_id = @identity_id
-	AND m.organization_id = e.organization_id
+    AND m.organization_id = e.organization_id
 LIMIT 1;
 `, tableName, scopeFragment)
 
@@ -296,19 +310,19 @@ func (m *Membership) LoadByIdentityAndOrg(
 ) error {
 	q := `
 WITH mbr AS (
-	SELECT
+    SELECT
-		am.id,
+        am.id,
-		am.identity_id,
+        am.identity_id,
-		am.organization_id,
+        am.organization_id,
-		am.role,
+        am.role,
-		am.created_at,
+        am.created_at,
-		am.updated_at
+        am.updated_at
-	FROM
+    FROM
-		iam_memberships am
+        iam_memberships am
-	WHERE
+    WHERE
-		am.identity_id = @identity_id
+        am.identity_id = @identity_id
-		AND am.organization_id = @organization_id
+        AND am.organization_id = @organization_id
-		AND %s
+        AND %s
 )
 SELECT
     mbr.id,
@@ -389,7 +403,7 @@ func (m *Membership) Delete(ctx context.Context, conn pg.Conn, scope Scoper, mem
 DELETE FROM
     iam_memberships
 WHERE
-	%s
+    %s
     AND id = @membership_id
 `
 
@@ -421,20 +435,20 @@ func (m *Memberships) LoadByIdentityID(
 ) error {
 	query := `
 WITH mbr AS (
-	SELECT
+    SELECT
-		id,
+        id,
-		identity_id,
+        identity_id,
-		organization_id,
+        organization_id,
-		role,
+        role,
-		created_at,
+        created_at,
-		updated_at
+        updated_at
-	FROM
+    FROM
-		iam_memberships
+        iam_memberships
-	WHERE
+    WHERE
-		identity_id = @identity_id
+        identity_id = @identity_id
-		AND %s
+        AND %s
-	ORDER BY
+    ORDER BY
-		created_at DESC
+        created_at DESC
 )
 SELECT
     mbr.id,
@@ -483,18 +497,18 @@ func (m *Memberships) LoadByOrganizationID(
 ) error {
 	query := `
 WITH mbr AS (
-	SELECT
+    SELECT
-		id,
+        id,
-		identity_id,
+        identity_id,
-		organization_id,
+        organization_id,
-		role,
+        role,
-		created_at,
+        created_at,
-		updated_at
+        updated_at
-	FROM
+    FROM
-		iam_memberships
+        iam_memberships
-	WHERE
+    WHERE
-		organization_id = @organization_id
+        organization_id = @organization_id
-		AND %s
+        AND %s
 )
 SELECT
     id,
@@ -506,19 +520,19 @@ SELECT
     created_at,
     updated_at
 FROM (
-	SELECT
+    SELECT
-		mbr.id,
+        mbr.id,
-		mbr.identity_id,
+        mbr.identity_id,
-		mbr.organization_id,
+        mbr.organization_id,
-		mbr.role,
+        mbr.role,
-		i.fullname as full_name,
+        i.fullname as full_name,
-		i.email_address,
+        i.email_address,
-		mbr.created_at,
+        mbr.created_at,
-		mbr.updated_at
+        mbr.updated_at
-	FROM
+    FROM
-		mbr
+        mbr
-	JOIN
+    JOIN
-		identities i ON mbr.identity_id = i.id
+        identities i ON mbr.identity_id = i.id
 ) AS membership_with_identity
 WHERE %s
 `

pkg/coredata/saml_configuration.go

@@ -18,6 +18,7 @@ import (
 	"context"
 	"crypto/x509"
 	"encoding/pem"
+	"errors"
 	"fmt"
 	"maps"
 	"time"
@@ -125,6 +126,10 @@ LIMIT 1;
 
 	config, err := pgx.CollectExactlyOneRow(rows, pgx.RowToStructByName[SAMLConfiguration])
 	if err != nil {
+		if errors.Is(err, pgx.ErrNoRows) {
+			return ErrResourceNotFound
+		}
+
 		return fmt.Errorf("cannot collect saml_configuration: %w", err)
 	}
 

pkg/coredata/session.go

@@ -195,9 +195,9 @@ UPDATE iam_sessions
 SET
     expired_at = @expired_at,
     updated_at = @updated_at,
-	user_agent = @user_agent,
+    user_agent = @user_agent,
-	ip_address = @ip_address,
+    ip_address = @ip_address,
-	expire_reason = @expire_reason,
+    expire_reason = @expire_reason,
     data = @data
 WHERE
     id = @session_id
@@ -246,7 +246,7 @@ FROM
     iam_sessions
 WHERE
     identity_id = @identity_id
-	AND %s
+    AND %s
 `
 
 	q = fmt.Sprintf(q, cursor.SQLFragment())
@@ -272,12 +272,12 @@ WHERE
 func (s *Sessions) CountByIdentityID(ctx context.Context, conn pg.Conn, identityID gid.GID) (int, error) {
 	q := `
 SELECT
-	COUNT(*)
+    COUNT(*)
 FROM
-	iam_sessions
+    iam_sessions
 WHERE
-	identity_id = @identity_id
+    identity_id = @identity_id
-	`
+`
 
 	args := pgx.StrictNamedArgs{"identity_id": identityID}
 
@@ -297,11 +297,11 @@ UPDATE iam_sessions
 SET
     expired_at = NOW(),
     updated_at = NOW(),
-	expire_reason = 'revoked'
+    expire_reason = 'revoked'
 WHERE
     id != @session_id
-	AND identity_id = @identity_id
+    AND identity_id = @identity_id
-	AND expire_reason IS NULL
+    AND expire_reason IS NULL
 `
 
 	args := pgx.StrictNamedArgs{
@@ -320,25 +320,27 @@ WHERE
 func (s *Session) LoadByRootSessionIDAndMembershipID(ctx context.Context, conn pg.Conn, rootSessionID gid.GID, membershipID gid.GID) error {
 	q := `
 SELECT
-	id,
+    id,
-	identity_id,
+    identity_id,
-	tenant_id,
+    tenant_id,
-	membership_id,
+    membership_id,
-	data,
+    data,
-	parent_session_id,
+    parent_session_id,
-	auth_method,
+    auth_method,
-	authenticated_at,
+    authenticated_at,
-	expire_reason,
+    expire_reason,
-	user_agent,
+    user_agent,
-	ip_address,
+    ip_address,
-	expired_at,
+    expired_at,
-	created_at,
+    created_at,
-	updated_at
+    updated_at
 FROM
-	iam_sessions
+    iam_sessions
 WHERE
-	parent_session_id = @root_session_id
+    parent_session_id = @root_session_id
-	AND membership_id = @membership_id
+    AND membership_id = @membership_id
+ORDER BY created_at DESC
+LIMIT 1
 `
 
 	args := pgx.StrictNamedArgs{