Handle new usage tracking with multi org

server/routers/site/createSite.ts

@@ -6,7 +6,7 @@ import response from "@server/lib/response";
 import HttpCode from "@server/types/HttpCode";
 import createHttpError from "http-errors";
 import logger from "@server/logger";
-import { eq, and } from "drizzle-orm";
+import { eq, and, count } from "drizzle-orm";
 import { getUniqueSiteName } from "../../db/names";
 import { addPeer } from "../gerbil/peers";
 import { fromError } from "zod-validation-error";
@@ -288,7 +288,6 @@ export async function createSite(
         const niceId = await getUniqueSiteName(orgId);
 
         let newSite: Site | undefined;
-        let numSites: Site[] | undefined;
         await db.transaction(async (trx) => {
             if (type == "newt") {
                 [newSite] = await trx
@@ -443,20 +442,9 @@ export async function createSite(
                 });
             }
 
-            numSites = await trx
-                .select()
-                .from(sites)
-                .where(eq(sites.orgId, orgId));
+            await usageService.add(orgId, FeatureId.SITES, 1, trx);
         });
 
-        if (numSites) {
-            await usageService.updateCount(
-                orgId,
-                FeatureId.SITES,
-                numSites.length
-            );
-        }
-
         if (!newSite) {
             return next(
                 createHttpError(

server/routers/site/deleteSite.ts

@@ -64,7 +64,6 @@ export async function deleteSite(
         }
 
         let deletedNewtId: string | null = null;
-        let numSites: Site[] | undefined;
 
         await db.transaction(async (trx) => {
             if (site.type == "wireguard") {
@@ -101,21 +100,9 @@ export async function deleteSite(
                 }
             }
 
-            await trx.delete(sites).where(eq(sites.siteId, siteId));
-
-            numSites = await trx
-                .select()
-                .from(sites)
-                .where(eq(sites.orgId, site.orgId));
+            await usageService.add(site.orgId, FeatureId.SITES, -1, trx);
         });
 
-        if (numSites) {
-            await usageService.updateCount(
-                site.orgId,
-                FeatureId.SITES,
-                numSites.length
-            );
-        }
         // Send termination message outside of transaction to prevent blocking
         if (deletedNewtId) {
             const payload = {