pangolin_mirror/pangolin
1
0
Fork 0
mirror of https://github.com/fosrl/pangolin.git synced 2026-03-03 09:16:40 +00:00
Code Issues Packages Projects Releases Wiki Activity

Handle new usage tracking with multi org

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

14
server/routers/domain/createOrgDomain.ts
View File

@@ -148,7 +148,6 @@ export async function createOrgDomain(
}
}
let numOrgDomains: OrgDomains[] | undefined;
let aRecords: CreateDomainResponse["aRecords"];
let cnameRecords: CreateDomainResponse["cnameRecords"];
let txtRecords: CreateDomainResponse["txtRecords"];
@@ -347,20 +346,9 @@ export async function createOrgDomain(
await trx.insert(dnsRecords).values(recordsToInsert);
}
numOrgDomains = await trx
.select()
.from(orgDomains)
.where(eq(orgDomains.orgId, orgId));
await usageService.add(orgId, FeatureId.DOMAINS, 1, trx);
});
if (numOrgDomains) {
await usageService.updateCount(
orgId,
FeatureId.DOMAINS,
numOrgDomains.length
);
}
if (!returned) {
return next(
createHttpError(

15
server/routers/domain/deleteOrgDomain.ts
View File

@@ -36,8 +36,6 @@ export async function deleteAccountDomain(
}
const { domainId, orgId } = parsed.data;
let numOrgDomains: OrgDomains[] | undefined;
await db.transaction(async (trx) => {
const [existing] = await trx
.select()
@@ -79,20 +77,9 @@ export async function deleteAccountDomain(
await trx.delete(domains).where(eq(domains.domainId, domainId));
numOrgDomains = await trx
.select()
.from(orgDomains)
.where(eq(orgDomains.orgId, orgId));
await usageService.add(orgId, FeatureId.DOMAINS, -1, trx);
});
if (numOrgDomains) {
await usageService.updateCount(
orgId,
FeatureId.DOMAINS,
numOrgDomains.length
);
}
return response<DeleteAccountDomainResponse>(res, {
data: { success: true },
success: true,

80
server/routers/org/createOrg.ts
View File

@@ -1,7 +1,7 @@
import { Request, Response, NextFunction } from "express";
import { z } from "zod";
import { db } from "@server/db";
import { and, eq } from "drizzle-orm";
import { and, count, eq } from "drizzle-orm";
import {
domains,
Org,
@@ -24,11 +24,7 @@ import { OpenAPITags, registry } from "@server/openApi";
import { isValidCIDR } from "@server/lib/validators";
import { createCustomer } from "#dynamic/lib/billing";
import { usageService } from "@server/lib/billing/usageService";
import {
FeatureId,
limitsService,
sandboxLimitSet
} from "@server/lib/billing";
import { FeatureId, limitsService, freeLimitSet } from "@server/lib/billing";
import { build } from "@server/build";
import { calculateUserClientsForOrgs } from "@server/lib/calculateUserClientsForOrgs";
import { doCidrsOverlap } from "@server/lib/ip";
@@ -114,6 +110,7 @@ export async function createOrg(
// )
// );
// }
//
// make sure the orgId is unique
const orgExists = await db
@@ -174,8 +171,46 @@ export async function createOrg(
}
}
if (build == "saas") {
if (!billingOrgIdForNewOrg) {
return next(
createHttpError(
HttpCode.INTERNAL_SERVER_ERROR,
"Billing org not found for user. Cannot create new organization."
)
);
}
const usage = await usageService.getUsage(billingOrgIdForNewOrg, FeatureId.ORGINIZATIONS);
if (!usage) {
return next(
createHttpError(
HttpCode.NOT_FOUND,
"No usage data found for this organization"
)
);
}
const rejectOrgs = await usageService.checkLimitSet(
billingOrgIdForNewOrg,
FeatureId.ORGINIZATIONS,
{
...usage,
instantaneousValue: (usage.instantaneousValue || 0) + 1
} // We need to add one to know if we are violating the limit
);
if (rejectOrgs) {
return next(
createHttpError(
HttpCode.FORBIDDEN,
"Organization limit exceeded. Please upgrade your plan."
)
);
}
}
let error = "";
let org: Org | null = null;
let numOrgs: number | null = null;
await db.transaction(async (trx) => {
const allDomains = await trx
@@ -184,14 +219,14 @@ export async function createOrg(
.where(eq(domains.configManaged, true));
// Generate SSH CA keys for the org
const ca = generateCA(`${orgId}-ca`);
const encryptionKey = config.getRawConfig().server.secret!;
const encryptedCaPrivateKey = encrypt(ca.privateKeyPem, encryptionKey);
// const ca = generateCA(`${orgId}-ca`);
// const encryptionKey = config.getRawConfig().server.secret!;
// const encryptedCaPrivateKey = encrypt(ca.privateKeyPem, encryptionKey);
const saasBillingFields =
build === "saas" && req.user && isFirstOrg !== null
? isFirstOrg
? { isBillingOrg: true as const, billingOrgId: null }
? { isBillingOrg: true as const, billingOrgId: orgId } // if this is the first org, it becomes the billing org for itself
: {
isBillingOrg: false as const,
billingOrgId: billingOrgIdForNewOrg
@@ -206,8 +241,8 @@ export async function createOrg(
subnet,
utilitySubnet,
createdAt: new Date().toISOString(),
sshCaPrivateKey: encryptedCaPrivateKey,
sshCaPublicKey: ca.publicKeyOpenSSH,
// sshCaPrivateKey: encryptedCaPrivateKey,
// sshCaPublicKey: ca.publicKeyOpenSSH,
...saasBillingFields
})
.returning();
@@ -310,6 +345,17 @@ export async function createOrg(
);
await calculateUserClientsForOrgs(ownerUserId, trx);
if (billingOrgIdForNewOrg) {
const [numOrgsResult] = await trx
.select({ count: count() })
.from(orgs)
.where(eq(orgs.billingOrgId, billingOrgIdForNewOrg)); // all the billable orgs including the primary org that is the billing org itself
numOrgs = numOrgsResult.count;
} else {
numOrgs = 1; // we only have one org if there is no billing org found out
}
});
if (!org) {
@@ -326,7 +372,7 @@ export async function createOrg(
}
if (build === "saas" && isFirstOrg === true) {
await limitsService.applyLimitSetToOrg(orgId, sandboxLimitSet);
await limitsService.applyLimitSetToOrg(orgId, freeLimitSet);
const customerId = await createCustomer(orgId, req.user?.email);
if (customerId) {
await usageService.updateCount(
@@ -338,6 +384,14 @@ export async function createOrg(
}
}
if (numOrgs) {
usageService.updateCount(
billingOrgIdForNewOrg || orgId,
FeatureId.ORGINIZATIONS,
numOrgs
);
}
return response(res, {
data: org,
success: true,

16
server/routers/site/createSite.ts
View File

@@ -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(

15
server/routers/site/deleteSite.ts
View File

@@ -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 = {

68
server/routers/user/acceptInvite.ts
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,

64
server/routers/user/createOrgUser.ts
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")

87
server/routers/user/removeUserOrg.ts
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,