Handle new usage tracking with multi org

server/routers/user/createOrgUser.ts

@@ -6,8 +6,8 @@ import createHttpError from "http-errors";
 import logger from "@server/logger";
 import { fromError } from "zod-validation-error";
 import { OpenAPITags, registry } from "@server/openApi";
-import { db, UserOrg } from "@server/db";
-import { and, eq } from "drizzle-orm";
+import { db, orgs, UserOrg } from "@server/db";
+import { and, eq, inArray, ne } from "drizzle-orm";
 import { idp, idpOidcConfig, roles, userOrgs, users } from "@server/db";
 import { generateId } from "@server/auth/sessions/app";
 import { usageService } from "@server/lib/billing/usageService";
@@ -151,6 +151,21 @@ export async function createOrgUser(
                 );
             }
 
+            const [org] = await db
+                .select()
+                .from(orgs)
+                .where(eq(orgs.orgId, orgId))
+                .limit(1);
+
+            if (!org) {
+                return next(
+                    createHttpError(
+                        HttpCode.NOT_FOUND,
+                        "Organization not found"
+                    )
+                );
+            }
+
             const [idpRes] = await db
                 .select()
                 .from(idp)
@@ -172,8 +187,6 @@ export async function createOrgUser(
                 );
             }
 
-            let orgUsers: UserOrg[] | undefined;
-
             await db.transaction(async (trx) => {
                 const [existingUser] = await trx
                     .select()
@@ -244,22 +257,37 @@ export async function createOrgUser(
                         .returning();
                 }
 
-                // List all of the users in the org
-                orgUsers = await trx
-                    .select()
-                    .from(userOrgs)
-                    .where(eq(userOrgs.orgId, orgId));
-
                 await calculateUserClientsForOrgs(userId, trx);
-            });
 
-            if (orgUsers) {
-                await usageService.updateCount(
-                    orgId,
-                    FeatureId.USERS,
-                    orgUsers.length
-                );
-            }
+                // calculate if the user is in any other of the orgs before we count it as an add to the billing org
+                if (org.billingOrgId) {
+                    const otherBillingOrgs = await trx
+                        .select()
+                        .from(orgs)
+                        .where(
+                            and(
+                                eq(orgs.billingOrgId, org.billingOrgId),
+                                ne(orgs.orgId, orgId)
+                            )
+                        );
+
+                    const billingOrgIds = otherBillingOrgs.map((o) => o.orgId);
+
+                    const orgsInBillingDomainThatTheUserIsStillIn = await trx
+                        .select()
+                        .from(userOrgs)
+                        .where(
+                            and(
+                                eq(userOrgs.userId, userId),
+                                inArray(userOrgs.orgId, billingOrgIds)
+                            )
+                        );
+
+                    if (orgsInBillingDomainThatTheUserIsStillIn.length === 0) {
+                        await usageService.add(orgId, FeatureId.USERS, 1, trx);
+                    }
+                }
+            });
         } else {
             return next(
                 createHttpError(HttpCode.BAD_REQUEST, "User type is required")