From a9b7cce49bf9b48dbf7b3560924c6de66dc72909 Mon Sep 17 00:00:00 2001 From: Owen Date: Tue, 23 Jun 2026 15:41:38 -0400 Subject: [PATCH] Improve efficiency of calculateUserClientsForOrgs --- server/lib/calculateUserClientsForOrgs.ts | 75 ++++++++++++++----- .../routers/orgIdp/unassociateOrgIdp.ts | 2 +- server/routers/auth/deleteMyAccount.ts | 2 +- server/routers/idp/validateOidcCallback.ts | 2 +- server/routers/olm/createUserOlm.ts | 2 +- server/routers/org/createOrg.ts | 14 +++- server/routers/user/acceptInvite.ts | 12 ++- server/routers/user/adminRemoveUser.ts | 2 +- server/routers/user/createOrgUser.ts | 17 ++--- server/routers/user/removeUserOrg.ts | 2 +- 10 files changed, 84 insertions(+), 46 deletions(-) diff --git a/server/lib/calculateUserClientsForOrgs.ts b/server/lib/calculateUserClientsForOrgs.ts index 090bf4d8c..fe09539a9 100644 --- a/server/lib/calculateUserClientsForOrgs.ts +++ b/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(); + 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 { + const trx = primaryDb; + const queuedAssociationRebuilds: ClientRow[] = []; + const execute = async (transaction: Transaction | typeof db) => { const orgCache = new Map(); 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 { // 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( diff --git a/server/private/routers/orgIdp/unassociateOrgIdp.ts b/server/private/routers/orgIdp/unassociateOrgIdp.ts index 41b2e6c89..0b5c1ed51 100644 --- a/server/private/routers/orgIdp/unassociateOrgIdp.ts +++ b/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}` ); diff --git a/server/routers/auth/deleteMyAccount.ts b/server/routers/auth/deleteMyAccount.ts index d03af5631..d45486946 100644 --- a/server/routers/auth/deleteMyAccount.ts +++ b/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}` ); diff --git a/server/routers/idp/validateOidcCallback.ts b/server/routers/idp/validateOidcCallback.ts index 6a82f2c12..f6cb656b3 100644 --- a/server/routers/idp/validateOidcCallback.ts +++ b/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 } diff --git a/server/routers/olm/createUserOlm.ts b/server/routers/olm/createUserOlm.ts index 714fb4b35..306317a0c 100644 --- a/server/routers/olm/createUserOlm.ts +++ b/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 diff --git a/server/routers/org/createOrg.ts b/server/routers/org/createOrg.ts index 7b2b1f87a..35466ebc0 100644 --- a/server/routers/org/createOrg.ts +++ b/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( diff --git a/server/routers/user/acceptInvite.ts b/server/routers/user/acceptInvite.ts index e3366a0c5..ef7ddcdbd 100644 --- a/server/routers/user/acceptInvite.ts +++ b/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(res, { data: { accepted: true, orgId: existingInvite.orgId }, diff --git a/server/routers/user/adminRemoveUser.ts b/server/routers/user/adminRemoveUser.ts index 38713ce26..066848b07 100644 --- a/server/routers/user/adminRemoveUser.ts +++ b/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}` ); diff --git a/server/routers/user/createOrgUser.ts b/server/routers/user/createOrgUser.ts index f03dd763b..c6f25e085 100644 --- a/server/routers/user/createOrgUser.ts +++ b/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( diff --git a/server/routers/user/removeUserOrg.ts b/server/routers/user/removeUserOrg.ts index 58fc85b69..902aeed84 100644 --- a/server/routers/user/removeUserOrg.ts +++ b/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}` );