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

@@ -1,8 +1,8 @@
import { Request, Response, NextFunction } from "express";
import { z } from "zod";
import { db, UserOrg } from "@server/db";
import { db, orgs, UserOrg } from "@server/db";
import { roles, userInvites, userOrgs, users } from "@server/db";
import { eq } from "drizzle-orm";
import { eq, and, inArray, ne } from "drizzle-orm";
import response from "@server/lib/response";
import HttpCode from "@server/types/HttpCode";
import createHttpError from "http-errors";
@@ -125,8 +125,22 @@ export async function acceptInvite(
}
}
const [org] = await db
.select()
.from(orgs)
.where(eq(orgs.orgId, existingInvite.orgId))
.limit(1);
if (!org) {
return next(
createHttpError(
HttpCode.BAD_REQUEST,
"Organization does not exist. Please contact an admin."
)
);
}
let roleId: number;
let totalUsers: UserOrg[] | undefined;
// get the role to make sure it exists
const existingRole = await db
.select()
@@ -160,25 +174,45 @@ export async function acceptInvite(
await calculateUserClientsForOrgs(existingUser[0].userId, trx);
// Get the total number of users in the org now
totalUsers = await trx
.select()
.from(userOrgs)
.where(eq(userOrgs.orgId, existingInvite.orgId));
// 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, existingInvite.orgId)
)
);
const billingOrgIds = otherBillingOrgs.map((o) => o.orgId);
const orgsInBillingDomainThatTheUserIsStillIn = await trx
.select()
.from(userOrgs)
.where(
and(
eq(userOrgs.userId, existingUser[0].userId),
inArray(userOrgs.orgId, billingOrgIds)
)
);
if (orgsInBillingDomainThatTheUserIsStillIn.length === 0) {
await usageService.add(
existingInvite.orgId,
FeatureId.USERS,
1,
trx
);
}
}
logger.debug(
`User ${existingUser[0].userId} accepted invite to org ${existingInvite.orgId}. Total users in org: ${totalUsers.length}`
`User ${existingUser[0].userId} accepted invite to org ${existingInvite.orgId}`
);
});
if (totalUsers) {
await usageService.updateCount(
existingInvite.orgId,
FeatureId.USERS,
totalUsers.length
);
}
return response<AcceptInviteResponse>(res, {
data: { accepted: true, orgId: existingInvite.orgId },
success: true,

View File

@@ -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")

View File

@@ -1,8 +1,16 @@
import { Request, Response, NextFunction } from "express";
import { z } from "zod";
import { db, resources, sites, UserOrg } from "@server/db";
import {
db,
orgs,
resources,
siteResources,
sites,
UserOrg,
userSiteResources
} from "@server/db";
import { userOrgs, userResources, users, userSites } from "@server/db";
import { and, count, eq, exists } from "drizzle-orm";
import { and, count, eq, exists, inArray } from "drizzle-orm";
import response from "@server/lib/response";
import HttpCode from "@server/types/HttpCode";
import createHttpError from "http-errors";
@@ -50,16 +58,16 @@ export async function removeUserOrg(
const { userId, orgId } = parsedParams.data;
// get the user first
const user = await db
const [user] = await db
.select()
.from(userOrgs)
.where(and(eq(userOrgs.userId, userId), eq(userOrgs.orgId, orgId)));
if (!user || user.length === 0) {
if (!user) {
return next(createHttpError(HttpCode.NOT_FOUND, "User not found"));
}
if (user[0].isOwner) {
if (user.isOwner) {
return next(
createHttpError(
HttpCode.BAD_REQUEST,
@@ -68,7 +76,17 @@ export async function removeUserOrg(
);
}
let userCount: UserOrg[] | 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) => {
await trx
@@ -97,6 +115,26 @@ export async function removeUserOrg(
)
);
await db.delete(userSiteResources).where(
and(
eq(userSiteResources.userId, userId),
exists(
db
.select()
.from(siteResources)
.where(
and(
eq(
siteResources.siteResourceId,
userSiteResources.siteResourceId
),
eq(siteResources.orgId, orgId)
)
)
)
)
);
await db.delete(userSites).where(
and(
eq(userSites.userId, userId),
@@ -114,11 +152,6 @@ export async function removeUserOrg(
)
);
userCount = await trx
.select()
.from(userOrgs)
.where(eq(userOrgs.orgId, orgId));
// if (build === "saas") {
// const [rootUser] = await trx
// .select()
@@ -137,15 +170,31 @@ export async function removeUserOrg(
// }
await calculateUserClientsForOrgs(userId, trx);
});
if (userCount) {
await usageService.updateCount(
orgId,
FeatureId.USERS,
userCount.length
);
}
// 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 billingOrgs = await trx
.select()
.from(orgs)
.where(eq(orgs.billingOrgId, org.billingOrgId));
const billingOrgIds = billingOrgs.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);
}
}
});
return response(res, {
data: null,