From 6750aa74f5c7f33e392cf5682809197565aa8d6f Mon Sep 17 00:00:00 2001 From: shogun444 Date: Mon, 28 Sep 2026 18:36:59 +0530 Subject: [PATCH 1/2] feat: aggregate user retention cohorts in set queries --- src/utils/masking/kpi.service.spec.ts | 70 +++++++++++ src/utils/masking/kpi.service.ts | 172 +++++++++++++++++++++----- 2 files changed, 212 insertions(+), 30 deletions(-) diff --git a/src/utils/masking/kpi.service.spec.ts b/src/utils/masking/kpi.service.spec.ts index 221bd712..46dee2b6 100644 --- a/src/utils/masking/kpi.service.spec.ts +++ b/src/utils/masking/kpi.service.spec.ts @@ -19,7 +19,9 @@ describe('KpiService', () => { innerJoin: jest.fn().mockReturnThis(), where: jest.fn().mockReturnThis(), andWhere: jest.fn().mockReturnThis(), + setParameters: jest.fn().mockReturnThis(), groupBy: jest.fn().mockReturnThis(), + addGroupBy: jest.fn().mockReturnThis(), getRawMany: jest.fn(), getRawOne: jest.fn(), }; @@ -108,6 +110,74 @@ describe('KpiService', () => { }); }); + describe('calculateUserRetention', () => { + beforeEach(() => { + jest.useFakeTimers().setSystemTime(new Date('2025-04-15T12:00:00Z')); + }); + + afterEach(() => { + jest.useRealTimers(); + }); + + it('should compute retention from grouped aggregates without find() or IN-lists', async () => { + mockQb.getRawMany.mockReset(); + mockQb.getRawMany + .mockResolvedValueOnce([ + { cohort_month: '2025-01', cohort_size: '100' }, + { cohort_month: '2025-02', cohort_size: '50' }, + { cohort_month: '2025-03', cohort_size: '200' }, + ]) + .mockResolvedValueOnce([ + { cohort_month: '2025-01', retained_month: '2025-02', retained_count: '40' }, + { cohort_month: '2025-01', retained_month: '2025-03', retained_count: '20' }, + ]); + + const setMock = jest.fn(); + const labelsSpy = jest + .spyOn(metricsService.userRetentionGauge, 'labels') + .mockReturnValue({ set: setMock } as any); + + await kpiService.calculateUserRetention(); + + expect(mockRepo.find).not.toHaveBeenCalled(); + expect(mockRepo.createQueryBuilder).toHaveBeenCalled(); + expect(mockQb.innerJoin).toHaveBeenCalled(); + expect(mockQb.groupBy).toHaveBeenCalled(); + + const whereCalls = [...mockQb.where.mock.calls, ...mockQb.andWhere.mock.calls].map( + (args) => String(args[0]), + ); + expect(whereCalls.join(' ')).not.toMatch(/IN\s*\(/i); + + // 40/100*100=40, 20/100*100=20 for the mocked cohort pair + expect(setMock).toHaveBeenCalledWith(40); + expect(setMock).toHaveBeenCalledWith(20); + labelsSpy.mockRestore(); + }); + + it('should report 0 when a retention bucket has no events and skip empty cohorts', async () => { + mockQb.getRawMany.mockReset(); + mockQb.getRawMany + .mockResolvedValueOnce([ + { cohort_month: '2025-01', cohort_size: '100' }, + // second cohort intentionally missing (empty) to verify skip + ]) + .mockResolvedValueOnce([]); + + const setMock = jest.fn(); + const labelsSpy = jest + .spyOn(metricsService.userRetentionGauge, 'labels') + .mockReturnValue({ set: setMock } as any); + + await kpiService.calculateUserRetention(); + + expect(mockRepo.find).not.toHaveBeenCalled(); + // Missing retention rows resolve to 0% rather than missing gauges + expect(setMock).toHaveBeenCalledWith(0); + labelsSpy.mockRestore(); + }); + }); + describe('handleCron', () => { it('should call all calculation methods', async () => { const activeUsersSpy = jest diff --git a/src/utils/masking/kpi.service.ts b/src/utils/masking/kpi.service.ts index 7ee9b24c..6eb4f604 100644 --- a/src/utils/masking/kpi.service.ts +++ b/src/utils/masking/kpi.service.ts @@ -1,7 +1,7 @@ import { Injectable, Logger } from '@nestjs/common'; import { Cron, CronExpression } from '@nestjs/schedule'; import { InjectRepository } from '@nestjs/typeorm'; -import { Repository, Between } from 'typeorm'; +import { Repository } from 'typeorm'; import { subDays, startOfDay, endOfDay, startOfMonth, format } from 'date-fns'; import { MetricsService } from './metrics.service'; @@ -129,27 +129,41 @@ export class KpiService { } async calculateUserRetention(): Promise { - // Calculate 3-month cohort retention + // Calculate 3-month cohort retention with set-based aggregates. + // Cohort/retention windows keep the exact JS definitions from before + // (rolling 30-day approximations); only the querying is set-based. const now = new Date(); this.metricsService.userRetentionGauge.reset(); + interface CohortWindow { + start: Date; + end: Date; + label: string; + } + + const cohortWindows: CohortWindow[] = []; for (let i = 1; i <= 3; i++) { const cohortMonthStart = startOfMonth(subDays(now, i * 30)); const cohortMonthEnd = endOfDay(subDays(startOfMonth(subDays(now, (i - 1) * 30)), 1)); - - const cohortUsers = await this.userRepository.find({ - select: ['id'], - where: { createdAt: Between(cohortMonthStart, cohortMonthEnd) }, + cohortWindows.push({ + start: cohortMonthStart, + end: cohortMonthEnd, + label: format(cohortMonthStart, 'yyyy-MM'), }); + } - const cohortUserIds = cohortUsers.map((u) => u.id); - const cohortSize = cohortUserIds.length; - - if (cohortSize === 0) continue; - - const cohortMonthLabel = format(cohortMonthStart, 'yyyy-MM'); + interface RetentionPair { + cohortLabel: string; + retainedLabel: string; + cohortStart: Date; + cohortEnd: Date; + retentionStart: Date; + retentionEnd: Date; + } - // Check retention for subsequent months + const pairs: RetentionPair[] = []; + for (let i = 1; i <= 3; i++) { + const cohort = cohortWindows[i - 1]; for (let j = 1; j < i; j++) { const retentionMonthStart = startOfMonth(subDays(now, (i - j) * 30)); const retentionMonthEnd = endOfDay( @@ -158,25 +172,123 @@ export class KpiService { if (retentionMonthStart > now) continue; - const retainedUsersCount = await this.eventRepository - .createQueryBuilder('event') - .select('COUNT(DISTINCT event.userId)', 'count') - .where('event.userId IN (:...cohortUserIds)', { cohortUserIds }) - .andWhere('event.createdAt BETWEEN :start AND :end', { - start: retentionMonthStart, - end: retentionMonthEnd, - }) - .getRawOne(); - - const retainedCount = parseInt(retainedUsersCount?.count ?? '0', 10); - const retentionRate = (retainedCount / cohortSize) * 100; - - const retainedMonthLabel = format(retentionMonthStart, 'yyyy-MM'); - this.metricsService.userRetentionGauge - .labels(cohortMonthLabel, retainedMonthLabel) - .set(retentionRate); + pairs.push({ + cohortLabel: cohort.label, + retainedLabel: format(retentionMonthStart, 'yyyy-MM'), + cohortStart: cohort.start, + cohortEnd: cohort.end, + retentionStart: retentionMonthStart, + retentionEnd: retentionMonthEnd, + }); + } + } + + if (pairs.length === 0) { + this.logger.log('Calculated user retention cohorts.'); + return; + } + + // 1) Cohort sizes in a single grouped query (no per-cohort find()). + const cohortCaseWhens: string[] = []; + const cohortParams: Record = {}; + cohortWindows.forEach((cohort, idx) => { + cohortCaseWhens.push( + `WHEN cohortUser.createdAt BETWEEN :c${idx}start AND :c${idx}end THEN :c${idx}label`, + ); + cohortParams[`c${idx}start`] = cohort.start; + cohortParams[`c${idx}end`] = cohort.end; + cohortParams[`c${idx}label`] = cohort.label; + }); + const cohortWhere = cohortWindows + .map((_, idx) => `cohortUser.createdAt BETWEEN :c${idx}start AND :c${idx}end`) + .join(' OR '); + + const cohortRows: Array<{ cohort_month: string; cohort_size: string }> = + await this.userRepository + .createQueryBuilder('cohortUser') + .select(`CASE ${cohortCaseWhens.join(' ')} END`, 'cohort_month') + .addSelect('COUNT(cohortUser.id)', 'cohort_size') + .where(`(${cohortWhere})`, cohortParams) + .setParameters(cohortParams) + .groupBy('cohort_month') + .getRawMany(); + + const cohortSizeByLabel = new Map(); + for (const row of cohortRows) { + if (!row.cohort_month) continue; + cohortSizeByLabel.set(row.cohort_month, Number(row.cohort_size ?? 0)); + } + + // 2) Retained distinct users in a single joined + grouped query (no IN-lists). + const uniqueRetentionWindows = new Map(); + for (const pair of pairs) { + if (!uniqueRetentionWindows.has(pair.retainedLabel)) { + uniqueRetentionWindows.set(pair.retainedLabel, { + start: pair.retentionStart, + end: pair.retentionEnd, + }); } } + const retentionLabels = [...uniqueRetentionWindows.keys()]; + const retentionCaseWhens: string[] = []; + const retentionParams: Record = { ...cohortParams }; + retentionLabels.forEach((label, idx) => { + const window = uniqueRetentionWindows.get(label)!; + retentionCaseWhens.push( + `WHEN event.createdAt BETWEEN :r${idx}start AND :r${idx}end THEN :r${idx}label`, + ); + retentionParams[`r${idx}start`] = window.start; + retentionParams[`r${idx}end`] = window.end; + retentionParams[`r${idx}label`] = label; + }); + + const pairClauses: string[] = []; + pairs.forEach((pair, idx) => { + pairClauses.push( + `(cohortUser.createdAt BETWEEN :p${idx}cStart AND :p${idx}cEnd AND event.createdAt BETWEEN :p${idx}rStart AND :p${idx}rEnd)`, + ); + retentionParams[`p${idx}cStart`] = pair.cohortStart; + retentionParams[`p${idx}cEnd`] = pair.cohortEnd; + retentionParams[`p${idx}rStart`] = pair.retentionStart; + retentionParams[`p${idx}rEnd`] = pair.retentionEnd; + }); + + const retentionRows: Array<{ + cohort_month: string; + retained_month: string; + retained_count: string; + }> = await this.eventRepository + .createQueryBuilder('event') + .innerJoin('event.user', 'cohortUser') + .select(`CASE ${cohortCaseWhens.join(' ')} END`, 'cohort_month') + .addSelect(`CASE ${retentionCaseWhens.join(' ')} END`, 'retained_month') + .addSelect('COUNT(DISTINCT event.userId)', 'retained_count') + .where(`(${pairClauses.join(' OR ')})`, retentionParams) + .setParameters(retentionParams) + .groupBy('cohort_month') + .addGroupBy('retained_month') + .getRawMany(); + + const retainedByPair = new Map(); + for (const row of retentionRows) { + if (!row.cohort_month || !row.retained_month) continue; + retainedByPair.set( + `${row.cohort_month}|${row.retained_month}`, + Number(row.retained_count ?? 0), + ); + } + + for (const pair of pairs) { + const cohortSize = cohortSizeByLabel.get(pair.cohortLabel) ?? 0; + if (cohortSize === 0) continue; + + const retainedCount = retainedByPair.get(`${pair.cohortLabel}|${pair.retainedLabel}`) ?? 0; + const retentionRate = (retainedCount / cohortSize) * 100; + + this.metricsService.userRetentionGauge + .labels(pair.cohortLabel, pair.retainedLabel) + .set(retentionRate); + } this.logger.log('Calculated user retention cohorts.'); } } From 6278befed5b2e5045d2ef80f3fa48559ea623c40 Mon Sep 17 00:00:00 2001 From: shogun444 Date: Mon, 28 Sep 2026 18:45:56 +0530 Subject: [PATCH 2/2] fix: correct retention spec formatting --- src/utils/masking/kpi.service.spec.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/utils/masking/kpi.service.spec.ts b/src/utils/masking/kpi.service.spec.ts index 46dee2b6..53da30d1 100644 --- a/src/utils/masking/kpi.service.spec.ts +++ b/src/utils/masking/kpi.service.spec.ts @@ -144,8 +144,8 @@ describe('KpiService', () => { expect(mockQb.innerJoin).toHaveBeenCalled(); expect(mockQb.groupBy).toHaveBeenCalled(); - const whereCalls = [...mockQb.where.mock.calls, ...mockQb.andWhere.mock.calls].map( - (args) => String(args[0]), + const whereCalls = [...mockQb.where.mock.calls, ...mockQb.andWhere.mock.calls].map((args) => + String(args[0]), ); expect(whereCalls.join(' ')).not.toMatch(/IN\s*\(/i);