Improve efficiency of calculateUserClientsForOrgs

server/lib/calculateUserClientsForOrgs.ts

@@ -6,6 +6,7 @@ import {
     db,
     olms,
     orgs,
+    primaryDb,
     roleClients,
     roles,
     Transaction,
@@ -23,10 +24,44 @@ import { rebuildClientAssociationsFromClient } from "./rebuildClientAssociations
 import { OlmErrorCodes } from "@server/routers/olm/error";
 import { tierMatrix } from "./billing/tierMatrix";
 
-export async function calculateUserClientsForOrgs(
+type ClientRow = typeof clients.$inferSelect;
+
+function runQueuedClientAssociationRebuilds(
     userId: string,
-    trx: Transaction | typeof db = db
+    queuedClients: ClientRow[]
+): void {
+    if (queuedClients.length === 0) {
+        return;
+    }
+
+    const uniqueClientsById = new Map<number, ClientRow>();
+    for (const client of queuedClients) {
+        uniqueClientsById.set(client.clientId, client);
+    }
+
+    void (async () => {
+        for (const client of uniqueClientsById.values()) {
+            try {
+                await rebuildClientAssociationsFromClient(client, db);
+            } catch (error) {
+                logger.error(
+                    `Failed rebuilding associations for client ${client.clientId} (user ${userId}): ${String(error)}`
+                );
+            }
+        }
+
+        logger.debug(
+            `Queued association rebuild completed for ${uniqueClientsById.size} client(s) (user ${userId})`
+        );
+    })();
+}
+
+export async function calculateUserClientsForOrgs(
+    userId: string
 ): Promise<void> {
+    const trx = primaryDb;
+    const queuedAssociationRebuilds: ClientRow[] = [];
+
     const execute = async (transaction: Transaction | typeof db) => {
         const orgCache = new Map<string, typeof orgs.$inferSelect | null>();
         const adminRoleCache = new Map<
@@ -189,7 +224,12 @@ export async function calculateUserClientsForOrgs(
 
         if (userOlms.length === 0) {
             // No OLMs for this user, but we should still clean up any orphaned clients
-            await cleanupOrphanedClients(userId, transaction);
+            await cleanupOrphanedClients(
+                userId,
+                transaction,
+                [],
+                queuedAssociationRebuilds
+            );
             return;
         }
 
@@ -382,10 +422,7 @@ export async function calculateUserClientsForOrgs(
                         .returning();
                 }
 
-                await rebuildClientAssociationsFromClient(
-                    newClient,
-                    transaction
-                );
+                queuedAssociationRebuilds.push(newClient);
 
                 // Grant admin role access to the client
                 await transaction.insert(roleClients).values({
@@ -414,24 +451,22 @@ export async function calculateUserClientsForOrgs(
         }
 
         // Clean up clients in orgs the user is no longer in
-        await cleanupOrphanedClients(userId, transaction, userOrgIds);
+        await cleanupOrphanedClients(
+            userId,
+            transaction,
+            userOrgIds,
+            queuedAssociationRebuilds
+        );
     };
 
-    if (trx) {
-        // Use provided transaction
-        await execute(trx);
-    } else {
-        // Create new transaction
-        await db.transaction(async (transaction) => {
-            await execute(transaction);
-        });
-    }
+    runQueuedClientAssociationRebuilds(userId, queuedAssociationRebuilds);
 }
 
 async function cleanupOrphanedClients(
     userId: string,
     trx: Transaction | typeof db,
-    userOrgIds: string[] = []
+    userOrgIds: string[] = [],
+    queuedAssociationRebuilds: ClientRow[] = []
 ): Promise<void> {
     // Find all OLM clients for this user that should be deleted
     // If userOrgIds is empty, delete all OLM clients (user has no orgs)
@@ -461,9 +496,9 @@ async function cleanupOrphanedClients(
             )
             .returning();
 
-        // Rebuild associations for each deleted client to clean up related data
+        // Queue deleted clients for post-transaction association cleanup.
         for (const deletedClient of deletedClients) {
-            await rebuildClientAssociationsFromClient(deletedClient, trx);
+            queuedAssociationRebuilds.push(deletedClient);
 
             if (deletedClient.olmId) {
                 await sendTerminateClient(

server/private/routers/orgIdp/unassociateOrgIdp.ts

@@ -121,7 +121,7 @@ export async function unassociateOrgIdp(
         });
 
         for (const userId of userIdsToRemove) {
-            calculateUserClientsForOrgs(userId, primaryDb).catch((e) => {
+            calculateUserClientsForOrgs(userId).catch((e) => {
                 logger.error(
                     `Failed to calculate user clients after removing user ${userId} from org ${orgId} during IdP unassociation: ${e}`
                 );

server/routers/auth/deleteMyAccount.ts

@@ -224,7 +224,7 @@ export async function deleteMyAccount(
             }
         });
 
-        calculateUserClientsForOrgs(userId, primaryDb).catch((e) => {
+        calculateUserClientsForOrgs(userId).catch((e) => {
             logger.error(
                 `Failed to calculate user clients after deleting account for user ${userId}: ${e}`
             );

server/routers/idp/validateOidcCallback.ts

@@ -635,7 +635,7 @@ export async function validateOidcCallback(
                 }
             });
 
-            calculateUserClientsForOrgs(userId!, primaryDb).catch((err) => {
+            calculateUserClientsForOrgs(userId!).catch((err) => {
                 logger.error(
                     "Error calculating user clients after syncing orgs and roles for OIDC user",
                     { error: err }

server/routers/olm/createUserOlm.ts

@@ -104,7 +104,7 @@ export async function createUserOlm(
             dateCreated: moment().toISOString()
         });
 
-        calculateUserClientsForOrgs(userId, primaryDb).catch((e) => {
+        calculateUserClientsForOrgs(userId).catch((e) => {
             console.error(
                 "Error calculating user clients after creating olm:",
                 e

server/routers/org/createOrg.ts

@@ -1,6 +1,6 @@
 import { Request, Response, NextFunction } from "express";
 import { z } from "zod";
-import { db } from "@server/db";
+import { db, primaryDb } from "@server/db";
 import { and, count, eq } from "drizzle-orm";
 import {
     domains,
@@ -233,6 +233,7 @@ export async function createOrg(
         let error = "";
         let org: Org | null = null;
         let numOrgs: number | null = null;
+        let ownerUserId: string | null = null;
 
         await db.transaction(async (trx) => {
             const allDomains = await trx
@@ -326,7 +327,6 @@ export async function createOrg(
                 );
             }
 
-            let ownerUserId: string | null = null;
             if (req.user) {
                 await trx.insert(userOrgs).values({
                     userId: req.user!.userId,
@@ -382,8 +382,6 @@ export async function createOrg(
                 }))
             );
 
-            await calculateUserClientsForOrgs(ownerUserId, trx);
-
             if (billingOrgIdForNewOrg) {
                 const [numOrgsResult] = await trx
                     .select({ count: count() })
@@ -396,6 +394,14 @@ export async function createOrg(
             }
         });
 
+        if (ownerUserId) {
+            calculateUserClientsForOrgs(ownerUserId).catch((e) => {
+                logger.error(
+                    `Failed to calculate user clients after creating org ${orgId} for user ${ownerUserId}: ${e}`
+                );
+            });
+        }
+
         if (!org) {
             return next(
                 createHttpError(

server/routers/user/acceptInvite.ts

@@ -202,13 +202,11 @@ export async function acceptInvite(
             );
         });
 
-        calculateUserClientsForOrgs(existingUser[0].userId, primaryDb).catch(
-            (e) => {
-                logger.error(
-                    `Failed to calculate user clients after accepting invite for user ${existingUser[0].userId}: ${e}`
-                );
-            }
-        );
+        calculateUserClientsForOrgs(existingUser[0].userId).catch((e) => {
+            logger.error(
+                `Failed to calculate user clients after accepting invite for user ${existingUser[0].userId}: ${e}`
+            );
+        });
 
         return response<AcceptInviteResponse>(res, {
             data: { accepted: true, orgId: existingInvite.orgId },

server/routers/user/adminRemoveUser.ts

@@ -55,7 +55,7 @@ export async function adminRemoveUser(
             await trx.delete(users).where(eq(users.userId, userId));
         });
 
-        calculateUserClientsForOrgs(userId, primaryDb).catch((e) => {
+        calculateUserClientsForOrgs(userId).catch((e) => {
             logger.error(
                 `Failed to calculate user clients after removing user ${userId}: ${e}`
             );

server/routers/user/createOrgUser.ts

@@ -56,7 +56,6 @@ const bodySchema = z
 export type CreateOrgUserResponse = {};
 const CreateOrgUserResponseDataSchema = z.object({});
 
-
 registry.registerPath({
     method: "put",
     path: "/org/{orgId}/user",
@@ -77,7 +76,9 @@ registry.registerPath({
             description: "Successful response",
             content: {
                 "application/json": {
-                    schema: createApiResponseSchema(CreateOrgUserResponseDataSchema)
+                    schema: createApiResponseSchema(
+                        CreateOrgUserResponseDataSchema
+                    )
                 }
             }
         }
@@ -326,13 +327,11 @@ export async function createOrgUser(
             });
 
             if (userIdForClients) {
-                calculateUserClientsForOrgs(userIdForClients, primaryDb).catch(
-                    (e) => {
-                        logger.error(
-                            `Failed to calculate user clients after creating org user: ${e}`
-                        );
-                    }
-                );
+                calculateUserClientsForOrgs(userIdForClients).catch((e) => {
+                    logger.error(
+                        `Failed to calculate user clients after creating org user: ${e}`
+                    );
+                });
             }
         } else {
             return next(

server/routers/user/removeUserOrg.ts

@@ -109,7 +109,7 @@ export async function removeUserOrg(
             await removeUserFromOrg(org, userId, trx);
         });
 
-        calculateUserClientsForOrgs(userId, primaryDb).catch((e) => {
+        calculateUserClientsForOrgs(userId).catch((e) => {
             logger.error(
                 `Failed to calculate user clients after removing user ${userId} from org ${orgId}: ${e}`
             );