Handle new usage tracking with multi org

This commit is contained in:
Owen
2026-02-17 17:09:48 -08:00
parent 79cf7c84dc
commit 4d6240c987
17 changed files with 432 additions and 584 deletions

View File

@@ -85,10 +85,14 @@ export async function getOrgUsage(
orgId,
FeatureId.REMOTE_EXIT_NODES
);
const egressData = await usageService.getUsage(
const organizations = await usageService.getUsage(
orgId,
FeatureId.EGRESS_DATA_MB
FeatureId.ORGINIZATIONS
);
// const egressData = await usageService.getUsage(
// orgId,
// FeatureId.EGRESS_DATA_MB
// );
if (sites) {
usageData.push(sites);
@@ -96,15 +100,18 @@ export async function getOrgUsage(
if (users) {
usageData.push(users);
}
if (egressData) {
usageData.push(egressData);
}
// if (egressData) {
// usageData.push(egressData);
// }
if (domains) {
usageData.push(domains);
}
if (remoteExitNodes) {
usageData.push(remoteExitNodes);
}
if (organizations) {
usageData.push(organizations);
}
const orgLimits = await db
.select()

View File

@@ -12,7 +12,14 @@
*/
import { NextFunction, Request, Response } from "express";
import { db, exitNodes, exitNodeOrgs, ExitNode, ExitNodeOrg } from "@server/db";
import {
db,
exitNodes,
exitNodeOrgs,
ExitNode,
ExitNodeOrg,
orgs
} from "@server/db";
import HttpCode from "@server/types/HttpCode";
import { z } from "zod";
import { remoteExitNodes } from "@server/db";
@@ -25,7 +32,7 @@ import { createRemoteExitNodeSession } from "#private/auth/sessions/remoteExitNo
import { fromError } from "zod-validation-error";
import { hashPassword, verifyPassword } from "@server/auth/password";
import logger from "@server/logger";
import { and, eq } from "drizzle-orm";
import { and, eq, inArray, ne } from "drizzle-orm";
import { getNextAvailableSubnet } from "@server/lib/exitNodes";
import { usageService } from "@server/lib/billing/usageService";
import { FeatureId } from "@server/lib/billing";
@@ -169,7 +176,17 @@ export async function createRemoteExitNode(
);
}
let numExitNodeOrgs: ExitNodeOrg[] | undefined;
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")
);
}
await db.transaction(async (trx) => {
if (!existingExitNode) {
@@ -217,19 +234,43 @@ export async function createRemoteExitNode(
});
}
numExitNodeOrgs = await trx
.select()
.from(exitNodeOrgs)
.where(eq(exitNodeOrgs.orgId, orgId));
});
// calculate if the node 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)
)
);
if (numExitNodeOrgs) {
await usageService.updateCount(
orgId,
FeatureId.REMOTE_EXIT_NODES,
numExitNodeOrgs.length
);
}
const billingOrgIds = otherBillingOrgs.map((o) => o.orgId);
const orgsInBillingDomainThatTheNodeIsStillIn = await trx
.select()
.from(exitNodeOrgs)
.where(
and(
eq(
exitNodeOrgs.exitNodeId,
existingExitNode.exitNodeId
),
inArray(exitNodeOrgs.orgId, billingOrgIds)
)
);
if (orgsInBillingDomainThatTheNodeIsStillIn.length === 0) {
await usageService.add(
orgId,
FeatureId.REMOTE_EXIT_NODES,
1,
trx
);
}
}
});
const token = generateSessionToken();
await createRemoteExitNodeSession(token, remoteExitNodeId);

View File

@@ -13,9 +13,9 @@
import { NextFunction, Request, Response } from "express";
import { z } from "zod";
import { db, ExitNodeOrg, exitNodeOrgs, exitNodes } from "@server/db";
import { db, ExitNodeOrg, exitNodeOrgs, exitNodes, orgs } from "@server/db";
import { remoteExitNodes } from "@server/db";
import { and, count, eq } from "drizzle-orm";
import { and, count, eq, inArray } from "drizzle-orm";
import response from "@server/lib/response";
import HttpCode from "@server/types/HttpCode";
import createHttpError from "http-errors";
@@ -50,7 +50,8 @@ export async function deleteRemoteExitNode(
const [remoteExitNode] = await db
.select()
.from(remoteExitNodes)
.where(eq(remoteExitNodes.remoteExitNodeId, remoteExitNodeId));
.where(eq(remoteExitNodes.remoteExitNodeId, remoteExitNodeId))
.limit(1);
if (!remoteExitNode) {
return next(
@@ -70,7 +71,17 @@ export async function deleteRemoteExitNode(
);
}
let numExitNodeOrgs: ExitNodeOrg[] | undefined;
const [org] = await db.select().from(orgs).where(eq(orgs.orgId, orgId));
if (!org) {
return next(
createHttpError(
HttpCode.NOT_FOUND,
`Org with ID ${orgId} not found`
)
);
}
await db.transaction(async (trx) => {
await trx
.delete(exitNodeOrgs)
@@ -81,38 +92,39 @@ export async function deleteRemoteExitNode(
)
);
const [remainingExitNodeOrgs] = await trx
.select({ count: count() })
.from(exitNodeOrgs)
.where(eq(exitNodeOrgs.exitNodeId, remoteExitNode.exitNodeId!));
// calculate if the user is in any other of the orgs before we count it as an remove to the billing org
if (org.billingOrgId) {
const otherBillingOrgs = await trx
.select()
.from(orgs)
.where(eq(orgs.billingOrgId, org.billingOrgId));
if (remainingExitNodeOrgs.count === 0) {
await trx
.delete(remoteExitNodes)
const billingOrgIds = otherBillingOrgs.map((o) => o.orgId);
const orgsInBillingDomainThatTheNodeIsStillIn = await trx
.select()
.from(exitNodeOrgs)
.where(
eq(remoteExitNodes.remoteExitNodeId, remoteExitNodeId)
and(
eq(
exitNodeOrgs.exitNodeId,
remoteExitNode.exitNodeId!
),
inArray(exitNodeOrgs.orgId, billingOrgIds)
)
);
await trx
.delete(exitNodes)
.where(
eq(exitNodes.exitNodeId, remoteExitNode.exitNodeId!)
if (orgsInBillingDomainThatTheNodeIsStillIn.length === 0) {
await usageService.add(
orgId,
FeatureId.REMOTE_EXIT_NODES,
-1,
trx
);
}
}
numExitNodeOrgs = await trx
.select()
.from(exitNodeOrgs)
.where(eq(exitNodeOrgs.orgId, orgId));
});
if (numExitNodeOrgs) {
await usageService.updateCount(
orgId,
FeatureId.REMOTE_EXIT_NODES,
numExitNodeOrgs.length
);
}
return response(res, {
data: null,
success: true,