introduce caching in calculate func

server/lib/calculateUserClientsForOrgs.ts

@@ -28,6 +28,159 @@ export async function calculateUserClientsForOrgs(
     trx?: Transaction
 ): Promise<void> {
     const execute = async (transaction: Transaction) => {
+        const orgCache = new Map<string, typeof orgs.$inferSelect | null>();
+        const adminRoleCache = new Map<
+            string,
+            typeof roles.$inferSelect | null
+        >();
+        const exitNodesCache = new Map<
+            string,
+            Awaited<ReturnType<typeof listExitNodes>>
+        >();
+        const isOrgLicensedCache = new Map<string, boolean>();
+        const existingClientCache = new Map<
+            string,
+            typeof clients.$inferSelect | null
+        >();
+        const roleClientAccessCache = new Map<string, boolean>();
+        const userClientAccessCache = new Map<string, boolean>();
+
+        const getOrgOlmKey = (orgId: string, olmId: string) =>
+            `${orgId}:${olmId}`;
+        const getRoleClientKey = (roleId: number, clientId: number) =>
+            `${roleId}:${clientId}`;
+        const getUserClientKey = (cachedUserId: string, clientId: number) =>
+            `${cachedUserId}:${clientId}`;
+
+        const getOrg = async (orgId: string) => {
+            if (orgCache.has(orgId)) {
+                return orgCache.get(orgId) ?? null;
+            }
+
+            const [org] = await transaction
+                .select()
+                .from(orgs)
+                .where(eq(orgs.orgId, orgId));
+            orgCache.set(orgId, org ?? null);
+
+            return org ?? null;
+        };
+
+        const getAdminRole = async (orgId: string) => {
+            if (adminRoleCache.has(orgId)) {
+                return adminRoleCache.get(orgId) ?? null;
+            }
+
+            const [adminRole] = await transaction
+                .select()
+                .from(roles)
+                .where(and(eq(roles.isAdmin, true), eq(roles.orgId, orgId)))
+                .limit(1);
+            adminRoleCache.set(orgId, adminRole ?? null);
+
+            return adminRole ?? null;
+        };
+
+        const getExitNodes = async (orgId: string) => {
+            if (exitNodesCache.has(orgId)) {
+                return exitNodesCache.get(orgId)!;
+            }
+
+            const exitNodes = await listExitNodes(orgId);
+            exitNodesCache.set(orgId, exitNodes);
+
+            return exitNodes;
+        };
+
+        const getIsOrgLicensed = async (orgId: string) => {
+            if (isOrgLicensedCache.has(orgId)) {
+                return isOrgLicensedCache.get(orgId)!;
+            }
+
+            const isOrgLicensed = await isLicensedOrSubscribed(
+                orgId,
+                tierMatrix.deviceApprovals
+            );
+            isOrgLicensedCache.set(orgId, isOrgLicensed);
+
+            return isOrgLicensed;
+        };
+
+        const getExistingClient = async (orgId: string, olmId: string) => {
+            const key = getOrgOlmKey(orgId, olmId);
+            if (existingClientCache.has(key)) {
+                return existingClientCache.get(key) ?? null;
+            }
+
+            const [existingClient] = await transaction
+                .select()
+                .from(clients)
+                .where(
+                    and(
+                        eq(clients.userId, userId),
+                        eq(clients.orgId, orgId),
+                        eq(clients.olmId, olmId)
+                    )
+                )
+                .limit(1);
+
+            existingClientCache.set(key, existingClient ?? null);
+
+            return existingClient ?? null;
+        };
+
+        const hasRoleClientAccess = async (
+            roleId: number,
+            clientId: number
+        ) => {
+            const key = getRoleClientKey(roleId, clientId);
+            if (roleClientAccessCache.has(key)) {
+                return roleClientAccessCache.get(key)!;
+            }
+
+            const [existingRoleClient] = await transaction
+                .select()
+                .from(roleClients)
+                .where(
+                    and(
+                        eq(roleClients.roleId, roleId),
+                        eq(roleClients.clientId, clientId)
+                    )
+                )
+                .limit(1);
+
+            const hasAccess = Boolean(existingRoleClient);
+            roleClientAccessCache.set(key, hasAccess);
+
+            return hasAccess;
+        };
+
+        const hasUserClientAccess = async (
+            cachedUserId: string,
+            clientId: number
+        ) => {
+            const key = getUserClientKey(cachedUserId, clientId);
+            if (userClientAccessCache.has(key)) {
+                return userClientAccessCache.get(key)!;
+            }
+
+            const [existingUserClient] = await transaction
+                .select()
+                .from(userClients)
+                .where(
+                    and(
+                        eq(userClients.userId, cachedUserId),
+                        eq(userClients.clientId, clientId)
+                    )
+                )
+                .limit(1);
+
+            const hasAccess = Boolean(existingUserClient);
+            userClientAccessCache.set(key, hasAccess);
+
+            return hasAccess;
+        };
+
         // Get all OLMs for this user
         const userOlms = await transaction
             .select()
@@ -54,7 +207,9 @@ export async function calculateUserClientsForOrgs(
             .innerJoin(roles, eq(userOrgRoles.roleId, roles.roleId))
             .where(eq(userOrgs.userId, userId));
 
-        const userOrgIds = [...new Set(userOrgRoleRows.map((r) => r.userOrgs.orgId))];
+        const userOrgIds = [
+            ...new Set(userOrgRoleRows.map((r) => r.userOrgs.orgId))
+        ];
         const orgIdToRoleRows = new Map<
             string,
             (typeof userOrgRoleRows)[0][]
@@ -64,6 +219,13 @@ export async function calculateUserClientsForOrgs(
             list.push(r);
             orgIdToRoleRows.set(r.userOrgs.orgId, list);
         }
+        const orgRequiresDeviceApprovalRole = new Map<string, boolean>();
+        for (const [orgId, roleRowsForOrg] of orgIdToRoleRows.entries()) {
+            orgRequiresDeviceApprovalRole.set(
+                orgId,
+                roleRowsForOrg.some((r) => r.roles.requireDeviceApproval)
+            );
+        }
 
         // For each OLM, ensure there's a client in each org the user is in
         for (const olm of userOlms) {
@@ -71,10 +233,7 @@ export async function calculateUserClientsForOrgs(
                 const roleRowsForOrg = orgIdToRoleRows.get(orgId)!;
                 const userOrg = roleRowsForOrg[0].userOrgs;
 
-                const [org] = await transaction
+                const org = await getOrg(orgId);
-                    .select()
-                    .from(orgs)
-                    .where(eq(orgs.orgId, orgId));
 
                 if (!org) {
                     logger.warn(
@@ -91,11 +250,7 @@ export async function calculateUserClientsForOrgs(
                 }
 
                 // Get admin role for this org (needed for access grants)
-                const [adminRole] = await transaction
+                const adminRole = await getAdminRole(orgId);
-                    .select()
-                    .from(roles)
-                    .where(and(eq(roles.isAdmin, true), eq(roles.orgId, orgId)))
-                    .limit(1);
 
                 if (!adminRole) {
                     logger.warn(
@@ -105,64 +260,50 @@ export async function calculateUserClientsForOrgs(
                 }
 
                 // Check if a client already exists for this OLM+user+org combination
-                const [existingClient] = await transaction
+                const existingClient = await getExistingClient(
-                    .select()
+                    orgId,
-                    .from(clients)
+                    olm.olmId
-                    .where(
+                );
-                        and(
-                            eq(clients.userId, userId),
-                            eq(clients.orgId, orgId),
-                            eq(clients.olmId, olm.olmId)
-                        )
-                    )
-                    .limit(1);
 
                 if (existingClient) {
                     // Ensure admin role has access to the client
-                    const [existingRoleClient] = await transaction
+                    const hasRoleAccess = await hasRoleClientAccess(
-                        .select()
+                        adminRole.roleId,
-                        .from(roleClients)
+                        existingClient.clientId
-                        .where(
+                    );
-                            and(
-                                eq(roleClients.roleId, adminRole.roleId),
-                                eq(
-                                    roleClients.clientId,
-                                    existingClient.clientId
-                                )
-                            )
-                        )
-                        .limit(1);
 
-                    if (!existingRoleClient) {
+                    if (!hasRoleAccess) {
                         await transaction.insert(roleClients).values({
                             roleId: adminRole.roleId,
                             clientId: existingClient.clientId
                         });
+                        roleClientAccessCache.set(
+                            getRoleClientKey(
+                                adminRole.roleId,
+                                existingClient.clientId
+                            ),
+                            true
+                        );
                         logger.debug(
                             `Granted admin role access to existing client ${existingClient.clientId} for OLM ${olm.olmId} in org ${orgId} (user ${userId})`
                         );
                     }
 
                     // Ensure user has access to the client
-                    const [existingUserClient] = await transaction
+                    const hasUserAccess = await hasUserClientAccess(
-                        .select()
+                        userId,
-                        .from(userClients)
+                        existingClient.clientId
-                        .where(
+                    );
-                            and(
-                                eq(userClients.userId, userId),
-                                eq(
-                                    userClients.clientId,
-                                    existingClient.clientId
-                                )
-                            )
-                        )
-                        .limit(1);
 
-                    if (!existingUserClient) {
+                    if (!hasUserAccess) {
                         await transaction.insert(userClients).values({
                             userId,
                             clientId: existingClient.clientId
                         });
+                        userClientAccessCache.set(
+                            getUserClientKey(userId, existingClient.clientId),
+                            true
+                        );
                         logger.debug(
                             `Granted user access to existing client ${existingClient.clientId} for OLM ${olm.olmId} in org ${orgId} (user ${userId})`
                         );
@@ -175,7 +316,7 @@ export async function calculateUserClientsForOrgs(
                 }
 
                 // Get exit nodes for this org
-                const exitNodesList = await listExitNodes(orgId);
+                const exitNodesList = await getExitNodes(orgId);
 
                 if (exitNodesList.length === 0) {
                     logger.warn(
@@ -206,14 +347,11 @@ export async function calculateUserClientsForOrgs(
 
                 const niceId = await getUniqueClientName(orgId);
 
-                const isOrgLicensed = await isLicensedOrSubscribed(
+                const isOrgLicensed = await getIsOrgLicensed(userOrg.orgId);
-                    userOrg.orgId,
-                    tierMatrix.deviceApprovals
-                );
                 const requireApproval =
                     build !== "oss" &&
                     isOrgLicensed &&
-                    roleRowsForOrg.some((r) => r.roles.requireDeviceApproval);
+                    orgRequiresDeviceApprovalRole.get(orgId) === true;
 
                 const newClientData: InferInsertModel<typeof clients> = {
                     userId,
@@ -232,6 +370,10 @@ export async function calculateUserClientsForOrgs(
                     .insert(clients)
                     .values(newClientData)
                     .returning();
+                existingClientCache.set(
+                    getOrgOlmKey(orgId, olm.olmId),
+                    newClient
+                );
 
                 // create approval request
                 if (requireApproval) {
@@ -257,12 +399,20 @@ export async function calculateUserClientsForOrgs(
                     roleId: adminRole.roleId,
                     clientId: newClient.clientId
                 });
+                roleClientAccessCache.set(
+                    getRoleClientKey(adminRole.roleId, newClient.clientId),
+                    true
+                );
 
                 // Grant user access to the client
                 await transaction.insert(userClients).values({
                     userId,
                     clientId: newClient.clientId
                 });
+                userClientAccessCache.set(
+                    getUserClientKey(userId, newClient.clientId),
+                    true
+                );
 
                 logger.debug(
                     `Created client for OLM ${olm.olmId} in org ${orgId} (user ${userId}) with access granted to admin role and user`