Add usage checker

This commit is contained in:
Bill Yang 2025-03-13 22:15:10 -07:00
parent db06705649
commit 3b05aa6757
8 changed files with 302 additions and 89 deletions

16
server/src/cron/index.ts Normal file
View file

@ -0,0 +1,16 @@
import * as cron from "node-cron";
import { updateUsersMonthlyUsage } from "./monthly-usage-checker.js";
export function initializeCronJobs() {
console.log("Initializing cron jobs...");
// Schedule the monthly usage checker to run every 5 minutes
cron.schedule("*/5 * * * *", async () => {
console.log("Running monthly usage checker cron job");
await updateUsersMonthlyUsage();
});
updateUsersMonthlyUsage();
console.log("Cron jobs initialized successfully");
}

View file

@ -0,0 +1,188 @@
import { user, member, sites, subscription } from "../db/postgres/schema.js";
import { clickhouse } from "../db/clickhouse/clickhouse.js";
import { STRIPE_PLANS } from "../lib/const.js";
import { eq, inArray, and } from "drizzle-orm";
import { db } from "../db/postgres/postgres.js";
import { processResults } from "../api/utils.js";
// Default event limit for users without an active subscription
const DEFAULT_EVENT_LIMIT = 20_000;
/**
* Gets the first day of the current month in YYYY-MM-DD format
*/
function getStartOfMonth(): string {
const date = new Date();
return `${date.getFullYear()}-${String(date.getMonth() + 1).padStart(
2,
"0"
)}-01`;
}
/**
* Gets all site IDs for organizations owned by a user
*/
async function getSiteIdsForUser(userId: string): Promise<number[]> {
try {
// Find the organizations this user is an owner of
const userOrgs = await db
.select({ organizationId: member.organizationId })
.from(member)
.where(and(eq(member.userId, userId), eq(member.role, "owner")));
if (!userOrgs.length) {
return [];
}
const orgIds = userOrgs.map((org) => org.organizationId);
// Get all sites for these organizations
const siteRecords = await db
.select({ siteId: sites.siteId })
.from(sites)
.where(inArray(sites.organizationId, orgIds));
return siteRecords.map((record) => record.siteId);
} catch (error) {
console.error(`Error getting sites for user ${userId}:`, error);
return [];
}
}
/**
* Gets event limit for a user based on their subscription plan
* @returns [eventLimit, periodStartDate]
*/
async function getUserSubscriptionInfo(
userId: string
): Promise<[number, string | null]> {
try {
// Find active subscription
const userSubscription = await db
.select()
.from(subscription)
.where(
and(
eq(subscription.referenceId, userId),
inArray(subscription.status, ["active", "trialing"])
)
)
.limit(1);
if (!userSubscription.length) {
return [DEFAULT_EVENT_LIMIT, null];
}
// Find the plan in STRIPE_PLANS
const plan = STRIPE_PLANS.find((p) => p.name === userSubscription[0].plan);
const eventLimit = plan ? plan.limits.events : DEFAULT_EVENT_LIMIT;
// Get period start date - if not available, use first day of month
const periodStart = userSubscription[0].periodStart
? new Date(userSubscription[0].periodStart).toISOString().split("T")[0]
: getStartOfMonth();
return [eventLimit, periodStart];
} catch (error) {
console.error(`Error getting subscription info for user ${userId}:`, error);
return [DEFAULT_EVENT_LIMIT, null];
}
}
/**
* Gets monthly pageview count from ClickHouse for the given site IDs
*/
async function getMonthlyPageviews(
siteIds: number[],
startDate: string | null
): Promise<number> {
if (!siteIds.length) {
return 0;
}
// If no startDate is provided (no subscription), default to start of month
const periodStart = startDate || getStartOfMonth();
try {
const result = await clickhouse.query({
query: `
SELECT COUNT(*) as count
FROM pageviews
WHERE site_id IN (${siteIds.join(",")})
AND timestamp >= toDate('${periodStart}')
`,
format: "JSONEachRow",
});
const rows = await processResults<{ count: string }>(result);
return parseInt(rows[0].count, 10);
} catch (error) {
console.error(
`Error querying ClickHouse for pageviews for sites ${siteIds}:`,
error
);
return 0;
}
}
/**
* Updates monthly event usage for all users
*/
export async function updateUsersMonthlyUsage() {
console.log(
"[Monthly Usage Checker] Starting check of monthly event usage..."
);
try {
// Get all users
const users = await db.select().from(user);
for (const userData of users) {
try {
// Get site IDs for organizations owned by this user
const siteIds = await getSiteIdsForUser(userData.id);
// If user has no sites, continue to next user
if (!siteIds.length) {
continue;
}
// Get user's subscription information (limit and period start)
const [eventLimit, periodStart] = await getUserSubscriptionInfo(
userData.id
);
// Get monthly pageview count from ClickHouse using the subscription period
const pageviewCount = await getMonthlyPageviews(siteIds, periodStart);
// Update user's monthlyEventCount and overMonthlyLimit fields
await db
.update(user)
.set({
monthlyEventCount: pageviewCount,
overMonthlyLimit: pageviewCount > eventLimit,
})
.where(eq(user.id, userData.id));
console.log(
`[Monthly Usage Checker] Updated user ${
userData.email
}: ${pageviewCount.toLocaleString()} events, limit ${eventLimit.toLocaleString()}, period started ${
periodStart || "this month"
}`
);
} catch (error) {
console.error(
`[Monthly Usage Checker] Error processing user ${userData.id}:`,
error
);
}
}
console.log("[Monthly Usage Checker] Completed monthly event usage check");
} catch (error) {
console.error(
"[Monthly Usage Checker] Error updating monthly usage:",
error
);
}
}

View file

@ -26,7 +26,7 @@ export const initializeClickhouse = async () => {
operating_system LowCardinality(String),
operating_system_version LowCardinality(String),
language LowCardinality(String),
country LowCardinality(String),
country LowCardinality(FixedString(2)),
iso_3166_2 LowCardinality(String),
screen_width UInt16,
screen_height UInt16,
@ -34,7 +34,7 @@ export const initializeClickhouse = async () => {
)
ENGINE = MergeTree()
PARTITION BY toYYYYMM(timestamp)
ORDER BY (timestamp, session_id)
ORDER BY (site_id, timestamp)
`,
});
@ -60,7 +60,7 @@ export const initializeClickhouse = async () => {
)
ENGINE = MergeTree()
PARTITION BY toYYYYMM(start_time)
ORDER BY (session_id, start_time)
ORDER BY (site_id, start_time)
`,
});
};

View file

@ -28,6 +28,8 @@ export const user = pgTable(
banReason: text(),
banExpires: timestamp({ mode: "string" }),
stripeCustomerId: text(),
overMonthlyLimit: boolean().default(false),
monthlyEventCount: integer().default(0),
},
(table) => [
unique("user_username_unique").on(table.username),

View file

@ -43,7 +43,7 @@ export async function cleanupOldSessions() {
RETURNING *
`;
console.log(`Cleaned up ${deletedSessions.length} sessions`);
// console.log(`Cleaned up ${deletedSessions.length} sessions`);
if (deletedSessions.length > 0) {
await insertSessions(

View file

@ -28,6 +28,7 @@ import { mapHeaders } from "./lib/auth-utils.js";
import { trackPageView } from "./tracker/trackPageView.js";
import { listOrganizationMembers } from "./api/listOrganizationMembers.js";
import { getUserOrganizations } from "./api/getUserOrganizations.js";
import { initializeCronJobs } from "./cron/index.js";
const __filename = fileURLToPath(import.meta.url);
const __dirname = dirname(__filename);
@ -150,10 +151,14 @@ const start = async () => {
await loadAllowedDomains();
// Start the server
await server.listen({ port: 3001, host: "0.0.0.0" });
// Start session cleanup cron job
cron.schedule("*/60 * * * * *", () => {
console.log("Cleaning up old sessions");
cleanupOldSessions();
});
// Initialize all cron jobs including monthly usage checker
initializeCronJobs();
} catch (err) {
server.log.error(err);
process.exit(1);

View file

@ -4,7 +4,7 @@ import dotenv from "dotenv";
import pg from "pg";
import { drizzleAdapter } from "better-auth/adapters/drizzle";
import { db } from "../db/postgres/postgres.js";
import { IS_CLOUD } from "./const.js";
import { IS_CLOUD, STRIPE_PLANS } from "./const.js";
import * as schema from "../db/postgres/schema.js";
import { eq } from "drizzle-orm";
import { stripe } from "@better-auth/stripe";
@ -31,89 +31,7 @@ const pluginList = IS_CLOUD
createCustomerOnSignUp: true,
subscription: {
enabled: true,
plans: [
{
priceId: "price_1R1fIVDFVprnAny2yJtRRPBm",
name: "basic100k",
interval: "month",
limits: {
events: 100_000,
},
},
{
priceId: "price_1R1fKJDFVprnAny2mfiBjkAQ",
name: "basic250k",
interval: "month",
limits: {
events: 250_000,
},
},
{
name: "basic500k",
priceId: "price_1R1fQlDFVprnAny2WwNdiRgT",
interval: "month",
limits: {
events: 500_000,
},
},
{
name: "basic1m",
priceId: "price_1R1fR2DFVprnAny28tPEQAwh",
interval: "month",
limits: {
events: 1_000_000,
},
},
{
name: "basic2m",
priceId: "price_1R1fRMDFVprnAny24AMo0Vuu",
interval: "month",
limits: {
events: 2_000_000,
},
},
{
name: "pro100k",
priceId: "price_1R1fRmDFVprnAny27gL7XFCY",
interval: "month",
limits: {
events: 100_000,
},
},
{
name: "pro250k",
priceId: "price_1R1fSADFVprnAny2d7d4tXTs",
interval: "month",
limits: {
events: 250_000,
},
},
{
name: "pro500k",
priceId: "price_1R1fSkDFVprnAny2MzBvhPKs",
interval: "month",
limits: {
events: 500_000,
},
},
{
name: "pro1m",
priceId: "price_1R1fTMDFVprnAny2IdeB1bLV",
interval: "month",
limits: {
events: 1_000_000,
},
},
{
name: "pro2m",
priceId: "price_1R1fTXDFVprnAny2JBLVtkIU",
interval: "month",
limits: {
events: 2_000_000,
},
},
],
plans: STRIPE_PLANS,
},
}),
]

View file

@ -3,3 +3,87 @@ import dotenv from "dotenv";
dotenv.config();
export const IS_CLOUD = process.env.CLOUD === "true";
export const STRIPE_PLANS = [
{
priceId: "price_1R1fIVDFVprnAny2yJtRRPBm",
name: "basic100k",
interval: "month",
limits: {
events: 100_000,
},
},
{
priceId: "price_1R1fKJDFVprnAny2mfiBjkAQ",
name: "basic250k",
interval: "month",
limits: {
events: 250_000,
},
},
{
name: "basic500k",
priceId: "price_1R1fQlDFVprnAny2WwNdiRgT",
interval: "month",
limits: {
events: 500_000,
},
},
{
name: "basic1m",
priceId: "price_1R1fR2DFVprnAny28tPEQAwh",
interval: "month",
limits: {
events: 1_000_000,
},
},
{
name: "basic2m",
priceId: "price_1R1fRMDFVprnAny24AMo0Vuu",
interval: "month",
limits: {
events: 2_000_000,
},
},
{
name: "pro100k",
priceId: "price_1R1fRmDFVprnAny27gL7XFCY",
interval: "month",
limits: {
events: 100_000,
},
},
{
name: "pro250k",
priceId: "price_1R1fSADFVprnAny2d7d4tXTs",
interval: "month",
limits: {
events: 250_000,
},
},
{
name: "pro500k",
priceId: "price_1R1fSkDFVprnAny2MzBvhPKs",
interval: "month",
limits: {
events: 500_000,
},
},
{
name: "pro1m",
priceId: "price_1R1fTMDFVprnAny2IdeB1bLV",
interval: "month",
limits: {
events: 1_000_000,
},
},
{
name: "pro2m",
priceId: "price_1R1fTXDFVprnAny2JBLVtkIU",
interval: "month",
limits: {
events: 2_000_000,
},
},
];