Count everything when deleting the org

server/lib/deleteOrg.ts

@@ -4,14 +4,18 @@ import {
     clientSitesAssociationsCache,
     db,
     domains,
+    exitNodeOrgs,
+    exitNodes,
     olms,
     orgDomains,
     orgs,
+    remoteExitNodes,
     resources,
-    sites
+    sites,
+    userOrgs
 } from "@server/db";
 import { newts, newtSessions } from "@server/db";
-import { eq, and, inArray, sql } from "drizzle-orm";
+import { eq, and, inArray, sql, count, countDistinct } from "drizzle-orm";
 import HttpCode from "@server/types/HttpCode";
 import createHttpError from "http-errors";
 import logger from "@server/logger";
@@ -62,6 +66,11 @@ export async function deleteOrgById(
     const deletedNewtIds: string[] = [];
     const olmsToTerminate: string[] = [];
 
+    let domainCount: number | null = null;
+    let siteCount: number | null = null;
+    let userCount: number | null = null;
+    let remoteExitNodeCount: number | null = null;
+
     await db.transaction(async (trx) => {
         for (const site of orgSites) {
             if (site.pubKey) {
@@ -141,8 +150,70 @@ export async function deleteOrgById(
         await usageService.add(orgId, FeatureId.ORGINIZATIONS, -1, trx); // here we are decreasing the org count BEFORE deleting the org because we need to still be able to get the org to get the billing org inside of here
 
         await trx.delete(orgs).where(eq(orgs.orgId, orgId));
+
+        if (org.billingOrgId) {
+            const billingOrgs = await trx
+                .select()
+                .from(orgs)
+                .where(eq(orgs.billingOrgId, org.billingOrgId));
+
+            if (billingOrgs.length > 0) {
+                const billingOrgIds = billingOrgs.map((org) => org.orgId);
+
+                const [domainCountRes] = await trx
+                    .select({ count: count() })
+                    .from(orgDomains)
+                    .where(inArray(orgDomains.orgId, billingOrgIds));
+
+                domainCount = domainCountRes.count;
+
+                const [siteCountRes] = await trx
+                    .select({ count: count() })
+                    .from(sites)
+                    .where(inArray(sites.orgId, billingOrgIds));
+
+                siteCount = siteCountRes.count;
+
+                const [userCountRes] = await trx
+                    .select({ count: countDistinct(userOrgs.userId) })
+                    .from(userOrgs)
+                    .where(inArray(userOrgs.orgId, billingOrgIds));
+
+                userCount = userCountRes.count;
+
+                const [remoteExitNodeCountRes] = await trx
+                    .select({ count: countDistinct(exitNodeOrgs.exitNodeId) })
+                    .from(exitNodeOrgs)
+                    .where(inArray(exitNodeOrgs.orgId, billingOrgIds));
+
+                remoteExitNodeCount = remoteExitNodeCountRes.count;
+            }
+        }
     });
 
+    if (org.billingOrgId) {
+        usageService.updateCount(
+            org.billingOrgId,
+            FeatureId.DOMAINS,
+            domainCount ?? 0
+        );
+        usageService.updateCount(
+            org.billingOrgId,
+            FeatureId.SITES,
+            siteCount ?? 0
+        );
+        usageService.updateCount(
+            org.billingOrgId,
+            FeatureId.USERS,
+            userCount ?? 0
+        );
+        usageService.updateCount(
+            org.billingOrgId,
+            FeatureId.REMOTE_EXIT_NODES,
+            remoteExitNodeCount ?? 0
+        );
+    }
+
     return { deletedNewtIds, olmsToTerminate };
 }