Refactor UsageService to streamline site and event count retrieval (#706)

- Updated the getSiteIdsForOrganization method to get all sites with their organization IDs, filtering out those without an organization.
- Introduced a new method, getAllSiteEventCounts, to retrieve monthly event counts for all sites in a single query, improving efficiency.
- Adjusted organization processing logic to utilize the new methods, enhancing clarity and performance in event usage checks.

Author: goldflag

server/src/services/usageService.ts

@@ -5,7 +5,7 @@ import { processResults } from "../api/analytics/utils.js";
 import { clickhouse } from "../db/clickhouse/clickhouse.js";
 import { db } from "../db/postgres/postgres.js";
 import { member, organization, sites, user } from "../db/postgres/schema.js";
-import { IS_CLOUD } from "../lib/const.js";
+import { DEFAULT_EVENT_LIMIT, IS_CLOUD } from "../lib/const.js";
 import { sendLimitExceededEmail } from "../lib/email/email.js";
 import { createServiceLogger } from "../lib/logger/logger.js";
 import { getBestSubscription } from "../lib/subscriptionUtils.js";
@@ -86,18 +86,21 @@ class UsageService {
   }
 
   /**
-   * Gets all site IDs for an organization
+   * Gets all sites with their organization IDs (excludes sites without an organization)
    */
-  private async getSiteIdsForOrganization(organizationId: string): Promise<number[]> {
+  private async getAllSites(): Promise<Array<{ siteId: number; organizationId: string }>> {
     try {
-      const siteRecords = await db
-        .select({ siteId: sites.siteId })
-        .from(sites)
-        .where(eq(sites.organizationId, organizationId));
+      const allSites = await db
+        .select({
+          siteId: sites.siteId,
+          organizationId: sites.organizationId,
+        })
+        .from(sites);
 
-      return siteRecords.map(record => record.siteId);
+      // Filter out sites without an organization ID
+      return allSites.filter(site => site.organizationId !== null) as Array<{ siteId: number; organizationId: string }>;
     } catch (error) {
-      this.logger.error(error as Error, `Error getting sites for organization ${organizationId}`);
+      this.logger.error(error as Error, `Error getting all sites`);
       return [];
     }
   }
@@ -114,7 +117,7 @@ class UsageService {
     name: string;
   }): Promise<[number, string | null]> {
     // Special case for specific organizations
-    if (orgData.name === "tomato 2" || orgData.name === "Zam") {
+    if (orgData.name === "rybbit" || orgData.name === "Zam") {
       return [Infinity, this.getStartOfMonth()];
     }
 
@@ -138,68 +141,40 @@ class UsageService {
   }
 
   /**
-   * Gets monthly event count from ClickHouse for the given site IDs
-   * Sites with ID < 2000 are grandfathered in and only count pageviews
-   * Sites with ID >= 2000 count all event types (pageview, custom_event, performance)
+   * Gets monthly event counts for all sites in a single query (for current month)
+   * Returns a map of site_id -> event count
    */
-  private async getMonthlyEventCount(siteIds: number[], startDate: string | null): Promise<number> {
-    if (!siteIds.length) {
-      return 0;
-    }
-
-    // If no startDate is provided (e.g., no subscription), default to start of month
-    const periodStart = startDate || this.getStartOfMonth();
-
-    // Split sites into grandfathered (< 2000) and new (>= 2000)
-    const grandfatheredSites = siteIds.filter(id => id < 2000);
-    const newSites = siteIds.filter(id => id >= 2000);
-
+  private async getAllSiteEventCounts(): Promise<Map<number, number>> {
     try {
-      let totalCount = 0;
-
-      // Count pageviews only for grandfathered sites
-      if (grandfatheredSites.length > 0) {
-        const grandfatheredResult = await clickhouse.query({
-          query: `
-            SELECT COUNT(*) as count
-            FROM events
-            WHERE site_id IN {siteIds:Array(Int32)} AND type = 'pageview'
+      const periodStart = this.getStartOfMonth();
+
+      const result = await clickhouse.query({
+        query: `
+          SELECT
+            site_id,
+            COUNT(*) as count
+          FROM events
+          WHERE type IN ('pageview', 'custom_event', 'performance')
             AND timestamp >= toDate({periodStart:String})
-          `,
-          format: "JSONEachRow",
-          query_params: {
-            siteIds: grandfatheredSites,
-            periodStart: periodStart,
-          },
-        });
-        const grandfatheredRows = await processResults<{ count: string }>(grandfatheredResult);
-        totalCount += parseInt(grandfatheredRows[0].count, 10);
-      }
+          GROUP BY site_id
+        `,
+        format: "JSONEachRow",
+        query_params: {
+          periodStart: periodStart,
+        },
+      });
 
-      // Count all events (pageview, custom_event, performance) for new sites
-      if (newSites.length > 0) {
-        const newSitesResult = await clickhouse.query({
-          query: `
-            SELECT COUNT(*) as count
-            FROM events
-            WHERE site_id IN {siteIds:Array(Int32)}
-            AND type IN ('pageview', 'custom_event', 'performance')
-            AND timestamp >= toDate({periodStart:String})
-          `,
-          format: "JSONEachRow",
-          query_params: {
-            siteIds: newSites,
-            periodStart: periodStart,
-          },
-        });
-        const newSitesRows = await processResults<{ count: string }>(newSitesResult);
-        totalCount += parseInt(newSitesRows[0].count, 10);
+      const rows = await processResults<{ site_id: number; count: string }>(result);
+
+      const eventCountMap = new Map<number, number>();
+      for (const row of rows) {
+        eventCountMap.set(row.site_id, parseInt(row.count, 10));
       }
 
-      return totalCount;
+      return eventCountMap;
     } catch (error) {
-      this.logger.error(error as Error, `Error querying ClickHouse for events for sites ${siteIds}`);
-      return 0;
+      this.logger.error(error as Error, "Error querying ClickHouse for event counts");
+      return new Map();
     }
   }
 
@@ -210,7 +185,22 @@ class UsageService {
     this.logger.info("Starting check of monthly event usage for organizations...");
 
     try {
-      // Get all organizations (both with and without Stripe customer IDs)
+      // Step 1: Get all sites with their organization IDs
+      const allSites = await this.getAllSites();
+
+      // Step 2: Get event counts for all sites in a single query (current month)
+      const eventCountMap = await this.getAllSiteEventCounts();
+
+      // Step 3: Build a map of organizationId -> { siteIds, eventCount }
+      const orgDataMap = new Map<string, { siteIds: number[]; eventCount: number }>();
+      for (const site of allSites) {
+        const orgData = orgDataMap.get(site.organizationId) || { siteIds: [], eventCount: 0 };
+        orgData.siteIds.push(site.siteId);
+        orgData.eventCount += eventCountMap.get(site.siteId) || 0;
+        orgDataMap.set(site.organizationId, orgData);
+      }
+
+      // Step 4: Get all organizations
       const organizations = await db
         .select({
           id: organization.id,
@@ -221,24 +211,30 @@ class UsageService {
         })
         .from(organization);
 
+      // Step 5: Process each organization
       for (const orgData of organizations) {
         try {
-          // Get site IDs for this organization
-          const siteIds = await this.getSiteIdsForOrganization(orgData.id);
-
-          // If organization has no sites, continue to next organization
-          if (!siteIds.length) {
-            continue;
+          const orgStats = orgDataMap.get(orgData.id);
+          const eventCount = orgStats?.eventCount || 0;
+          const siteIds = orgStats?.siteIds || [];
+
+          // Only fetch subscription info for organizations with > 3000 events
+          // This avoids slow Stripe API calls for low-usage orgs
+          let eventLimit: number;
+          let isOverLimit: boolean;
+
+          if (eventCount <= DEFAULT_EVENT_LIMIT) {
+            // Free tier limit is 3000, so they're definitely not over limit
+            eventLimit = DEFAULT_EVENT_LIMIT;
+            isOverLimit = false;
+            this.logger.debug(`Organization ${orgData.name} has ${eventCount} events, skipping subscription check`);
+          } else {
+            // High usage - need to check their actual subscription
+            const [fetchedLimit, periodStart] = await this.getOrganizationSubscriptionInfo(orgData);
+            eventLimit = fetchedLimit;
+            isOverLimit = eventCount > eventLimit;
           }
 
-          // Get organization's subscription information (limit and period start)
-          const [eventLimit, periodStart] = await this.getOrganizationSubscriptionInfo(orgData);
-
-          // Get monthly event count from ClickHouse using the billing period start date
-          const eventCount = await this.getMonthlyEventCount(siteIds, periodStart);
-
-          // Check if over limit and update global set
-          const isOverLimit = eventCount > eventLimit;
           const wasOverLimit = orgData.overMonthlyLimit ?? false;
 
           // Update organization's monthlyEventCount and overMonthlyLimit fields
@@ -286,13 +282,8 @@ class UsageService {
             }
           }
 
-          // Format additional date info for logging if available
-          const periodInfo = periodStart ? `period started ${periodStart}` : "this month";
-
           this.logger.info(
-            `Updated organization ${
-              orgData.name
-            }: ${eventCount.toLocaleString()} events, limit ${eventLimit.toLocaleString()}, ${periodInfo}`
+            `Updated organization ${orgData.name}: ${eventCount.toLocaleString()} events, limit ${eventLimit.toLocaleString()}`
           );
         } catch (error) {
           this.logger.error(error as Error, `Error processing organization ${orgData.id}`);