From 21658f1b781830c940bfcad875b2552cf7f38fe5 Mon Sep 17 00:00:00 2001 From: John Yeo Date: Wed, 22 Apr 2026 20:06:09 +0100 Subject: [PATCH] feat: optimize entities cache --- server/experiments/benchmarkV2CacheRedis.ts | 440 ++++++++++++++++++ server/experiments/explainEntityAggregate.ts | 49 +- server/experiments/explainGetFullSubject.ts | 56 ++- .../getEntityAggregateFragments.ts | 154 ++---- .../full-subject-aggregate-options.test.ts | 1 - .../full-subject-aggregates.test.ts | 2 - .../utils/buildComparableFullSubject.ts | 1 - .../full-subject-aggregate-balance.test.ts | 1 - .../cusEntModels/aggregatedCusEnt.ts | 1 - .../cusUtils/fullCusUtils/logFullCustomer.ts | 1 - .../utils/fullSubjectUtils/logFullSubject.ts | 1 - 11 files changed, 573 insertions(+), 134 deletions(-) create mode 100644 server/experiments/benchmarkV2CacheRedis.ts diff --git a/server/experiments/benchmarkV2CacheRedis.ts b/server/experiments/benchmarkV2CacheRedis.ts new file mode 100644 index 000000000..a1d62d444 --- /dev/null +++ b/server/experiments/benchmarkV2CacheRedis.ts @@ -0,0 +1,440 @@ +import { AppEnv } from "@autumn/shared"; +import type { Redis } from "ioredis"; +import { + initDrizzle, + prodTestCustomerId, + prodTestEntityId, + prodTestOrgId, +} from "./experimentEnv"; + +// --------------------------------------------------------------------------- +// EDIT ME — hardcode the customer / feature you want to diagnose +// --------------------------------------------------------------------------- +const CONFIG = { + orgId: prodTestOrgId as string, + env: AppEnv.Live as AppEnv, + customerId: prodTestCustomerId as string, + entityId: prodTestEntityId as string | undefined, // or undefined + featureId: undefined as string | undefined, // set to isolate to one feature + iterations: 5, + runLua: true, // lua writes back entitlement hashes even for amount=0 — disable if paranoid + forceRepopulate: true, // DEL the cached subject before priming, so it's rebuilt from the DB +}; +// --------------------------------------------------------------------------- + +const { resolveRedisV2 } = await import("../src/external/redis/resolveRedisV2"); +const { warmupRedisV2 } = await import("../src/external/redis/initRedisV2"); +const { buildFullSubjectKey } = await import( + "../src/internal/customers/cache/fullSubject/builders/buildFullSubjectKey" +); +const { buildSharedFullSubjectBalanceKey } = await import( + "../src/internal/customers/cache/fullSubject/builders/buildSharedFullSubjectBalanceKey" +); +const { buildDeductFromSubjectBalancesKeys } = await import( + "../src/internal/customers/cache/fullSubject/builders/buildDeductFromSubjectBalancesKeys" +); +const { getOrInitFullSubjectViewEpoch } = await import( + "../src/internal/customers/cache/fullSubject/actions/invalidate/getOrInitFullSubjectViewEpoch" +); +const { getCachedFeatureBalancesBatch } = await import( + "../src/internal/customers/cache/fullSubject/balances/getCachedFeatureBalances" +); +const { getOrSetCachedFullSubject } = await import( + "../src/internal/customers/cache/fullSubject/actions/getOrSetCachedFullSubject" +); +const { sanitizeCachedFullSubject } = await import( + "../src/internal/customers/cache/fullSubject/sanitize/index" +); +const { AGGREGATED_BALANCE_FIELD } = await import( + "../src/internal/customers/cache/fullSubject/config/fullSubjectCacheConfig" +); + +type Step = { name: string; ms: number; extra?: string }; +type Iter = { index: number; steps: Step[] }; + +const nowMs = () => performance.now(); + +const time = async ( + name: string, + fn: () => Promise, +): Promise<{ step: Step; value: T }> => { + const start = nowMs(); + const value = await fn(); + const ms = nowMs() - start; + return { step: { name, ms }, value }; +}; + +const makeCtx = ({ redisV2, db }: { redisV2: Redis; db: unknown }) => + ({ + org: { id: CONFIG.orgId }, + env: CONFIG.env, + redisV2, + db, + dbGeneral: db, + features: [], + skipCache: false, + isPublic: false, + extraLogs: {}, + logger: { + debug: (...a: unknown[]) => console.debug(...a), + info: (...a: unknown[]) => console.info(...a), + warn: (...a: unknown[]) => console.warn(...a), + error: (...a: unknown[]) => console.error(...a), + }, + }) as unknown as Parameters[0]["ctx"]; + +const pct = (sorted: number[], p: number): number => { + if (sorted.length === 0) return 0; + const idx = Math.min(sorted.length - 1, Math.floor((p / 100) * sorted.length)); + return sorted[idx]; +}; + +const summarize = (all: Iter[]) => { + const byStep = new Map(); + for (const iter of all) { + for (const step of iter.steps) { + const arr = byStep.get(step.name) ?? []; + arr.push(step.ms); + byStep.set(step.name, arr); + } + } + + console.log("\n=== Summary (ms) ==="); + console.log( + "step".padEnd(48), + "min".padStart(8), + "p50".padStart(8), + "p95".padStart(8), + "max".padStart(8), + ); + for (const [name, arr] of byStep) { + const sorted = [...arr].sort((a, b) => a - b); + console.log( + name.padEnd(48), + sorted[0].toFixed(2).padStart(8), + pct(sorted, 50).toFixed(2).padStart(8), + pct(sorted, 95).toFixed(2).padStart(8), + sorted[sorted.length - 1].toFixed(2).padStart(8), + ); + } +}; + +const runIteration = async ({ + i, + redisV2, + ctx, +}: { + i: number; + redisV2: Redis; + // biome-ignore lint/suspicious/noExplicitAny: stub ctx + ctx: any; +}): Promise => { + const steps: Step[] = []; + const subjectKey = buildFullSubjectKey({ + orgId: CONFIG.orgId, + env: CONFIG.env, + customerId: CONFIG.customerId, + entityId: CONFIG.entityId, + }); + + // A1. GET subject + const { step: s1, value: cachedRaw } = await time( + "GET full_subject", + () => redisV2.get(subjectKey), + ); + steps.push({ ...s1, extra: cachedRaw ? `${cachedRaw.length}b` : "null" }); + + if (!cachedRaw) { + console.log(` iter ${i}: no cached subject — stopping this iteration`); + return { index: i, steps }; + } + + // parse + const { step: sParse, value: cached } = await time( + "JSON.parse + sanitize", + async () => + sanitizeCachedFullSubject({ + cachedFullSubject: JSON.parse(cachedRaw), + }), + ); + steps.push(sParse); + + if (i === 1) { + const sizeOf = (v: unknown) => JSON.stringify(v ?? null).length; + const cachedAny = cached as unknown as Record; + const entries = Object.keys(cachedAny) + .map((k) => ({ k, bytes: sizeOf(cachedAny[k]) })) + .sort((a, b) => b.bytes - a.bytes); + const total = entries.reduce((a, e) => a + e.bytes, 0); + console.log(` blob anatomy (total ${(total / 1024).toFixed(1)}kb):`); + for (const e of entries.slice(0, 8)) { + const pct = ((e.bytes / total) * 100).toFixed(1); + console.log( + ` ${(e.bytes / 1024).toFixed(1).padStart(8)}kb ${pct.padStart(5)}% ${e.k}`, + ); + } + } + + // A2. epoch + const { step: s2 } = await time("getOrInitFullSubjectViewEpoch", () => + getOrInitFullSubjectViewEpoch({ ctx, customerId: CONFIG.customerId }), + ); + steps.push(s2); + + const meteredFeatures = CONFIG.featureId + ? cached.meteredFeatures.filter((f: string) => f === CONFIG.featureId) + : cached.meteredFeatures; + + const includeAggregated = !CONFIG.entityId; + + // A3. batch hmget pipeline + const { step: s3, value: batchOutcome } = await time( + `getCachedFeatureBalancesBatch (${meteredFeatures.length} features)`, + () => + getCachedFeatureBalancesBatch({ + ctx, + customerId: CONFIG.customerId, + featureIds: meteredFeatures, + customerEntitlementIdsByFeatureId: + cached.customerEntitlementIdsByFeatureId, + includeAggregated, + }), + ); + steps.push({ ...s3, extra: batchOutcome.kind }); + + // A4. sequential per-feature HMGET with per-feature timing (slow-key hunt) + const perFeature: Array<{ featureId: string; ms: number; bytes: number }> = + []; + const { step: s4 } = await time( + `sequential HMGET per feature (${meteredFeatures.length} calls)`, + async () => { + for (const featureId of meteredFeatures) { + const ids = + cached.customerEntitlementIdsByFeatureId[featureId] ?? []; + const fields = includeAggregated + ? [...ids, AGGREGATED_BALANCE_FIELD] + : ids; + const balanceKey = buildSharedFullSubjectBalanceKey({ + orgId: CONFIG.orgId, + env: CONFIG.env, + customerId: CONFIG.customerId, + featureId, + }); + const t = nowMs(); + let bytes = 0; + if (fields.length === 0) { + await redisV2.exists(balanceKey); + } else { + const vals = (await redisV2.hmget( + balanceKey, + ...fields, + )) as (string | null)[]; + bytes = vals.reduce((a, v) => a + (v?.length ?? 0), 0); + } + perFeature.push({ featureId, ms: nowMs() - t, bytes }); + } + }, + ); + steps.push(s4); + + // Only print on first iteration to keep output tidy + if (i === 1) { + const sorted = [...perFeature].sort((a, b) => b.ms - a.ms); + console.log(" top-5 slowest features (seq HMGET):"); + for (const f of sorted.slice(0, 5)) { + console.log( + ` ${f.ms.toFixed(2).padStart(7)}ms ${(f.bytes / 1024).toFixed(1).padStart(6)}kb ${f.featureId}`, + ); + } + } + + // C. dry lua invocation (amount_to_deduct=0) + if (CONFIG.runLua && meteredFeatures.length > 0) { + const customerEntitlementDeductions: Array<{ + customer_entitlement_id: string; + credit_cost: number; + feature_id: string; + entity_feature_id: string | null; + usage_allowed: boolean; + min_balance: number; + max_balance: number; + }> = []; + + for (const featureId of meteredFeatures) { + const ids = cached.customerEntitlementIdsByFeatureId[featureId] ?? []; + for (const id of ids) { + customerEntitlementDeductions.push({ + customer_entitlement_id: id, + credit_cost: 0, + feature_id: featureId, + entity_feature_id: null, + usage_allowed: false, + min_balance: 0, + max_balance: 0, + }); + } + } + + if (customerEntitlementDeductions.length === 0) { + console.log(" (skipping lua — no customer entitlements)"); + } else { + const routingKey = subjectKey; + const { keys, balanceKeyIndexByFeatureId } = + buildDeductFromSubjectBalancesKeys({ + orgId: CONFIG.orgId, + env: CONFIG.env, + customerId: CONFIG.customerId, + routingKey, + lockReceiptKey: null, + customerEntitlementDeductions, + fallbackFeatureId: + CONFIG.featureId ?? meteredFeatures[0] ?? "unknown", + }); + + const luaParams = { + org_id: CONFIG.orgId, + env: CONFIG.env, + customer_id: CONFIG.customerId, + customer_entitlement_deductions: customerEntitlementDeductions, + balance_key_index_by_feature_id: balanceKeyIndexByFeatureId, + spend_limit_by_feature_id: null, + usage_based_cus_ent_ids_by_feature_id: null, + amount_to_deduct: 0, + target_balance: null, + target_entity_id: CONFIG.entityId ?? null, + rollovers: null, + skip_additional_balance: false, + alter_granted_balance: false, + overage_behaviour: "cap", + feature_id: CONFIG.featureId ?? meteredFeatures[0], + lock: null, + unwind_value: null, + debug: true, + }; + + const { step: sLua, value: luaRaw } = await time( + `deductFromSubjectBalances lua (${keys.length} keys, ${customerEntitlementDeductions.length} ents)`, + () => + // biome-ignore lint/suspicious/noExplicitAny: ioredis custom command typed via module augmentation not available here + (redisV2 as any).deductFromSubjectBalances( + keys.length, + ...keys, + JSON.stringify(luaParams), + ), + ); + steps.push(sLua); + + try { + const luaResult = JSON.parse(luaRaw as string); + if (luaResult.error) { + console.log(` lua error: ${luaResult.error}`); + } + if (luaResult.logs?.length) { + console.log(` lua logs (${luaResult.logs.length} lines):`); + for (const line of luaResult.logs) console.log(` ${line}`); + } + } catch { + console.log(" (lua result unparseable)"); + } + } + } + + return { index: i, steps }; +}; + +const main = async () => { + console.log("=== V2 cache Redis benchmark ==="); + console.log("config:", { + orgId: CONFIG.orgId, + env: CONFIG.env, + customerId: CONFIG.customerId, + entityId: CONFIG.entityId, + featureId: CONFIG.featureId, + iterations: CONFIG.iterations, + runLua: CONFIG.runLua, + }); + + const { db } = initDrizzle(); + const redisV2 = resolveRedisV2(); + await warmupRedisV2(); + + const ctx = makeCtx({ redisV2, db }); + + // Prime the cache if cold, so the iteration benchmarks always measure a hit path. + const subjectKey = buildFullSubjectKey({ + orgId: CONFIG.orgId, + env: CONFIG.env, + customerId: CONFIG.customerId, + entityId: CONFIG.entityId, + }); + if (CONFIG.forceRepopulate) { + // Grab meteredFeatures from the existing subject so we know which + // balance hashes to drop. All keys share the {customerId} hash tag, + // so a single DEL across them is slot-safe. + const staleRaw = await redisV2.get(subjectKey); + const balanceKeys: string[] = []; + if (staleRaw) { + try { + const stale = JSON.parse(staleRaw) as { + meteredFeatures?: string[]; + }; + for (const featureId of stale.meteredFeatures ?? []) { + balanceKeys.push( + buildSharedFullSubjectBalanceKey({ + orgId: CONFIG.orgId, + env: CONFIG.env, + customerId: CONFIG.customerId, + featureId, + }), + ); + } + } catch { + console.warn(" (could not parse stale subject to extract balance keys)"); + } + } + const toDel = [subjectKey, ...balanceKeys]; + const deleted = await redisV2.del(...toDel); + console.log( + `\n--- forceRepopulate: DEL ${toDel.length} keys (subject + ${balanceKeys.length} balance hashes) → ${deleted} removed ---`, + ); + } + + const existing = await redisV2.get(subjectKey); + if (!existing) { + console.log("\n--- priming cache (cold) ---"); + const t0 = nowMs(); + try { + await getOrSetCachedFullSubject({ + ctx, + customerId: CONFIG.customerId, + entityId: CONFIG.entityId, + source: "benchmarkV2CacheRedis", + }); + console.log(` primed in ${(nowMs() - t0).toFixed(2)}ms`); + } catch (err) { + console.error(" prime failed:", err); + console.error( + " (cache is cold and cannot be populated from this script — aborting)", + ); + process.exit(1); + } + } else { + console.log(`\ncache already warm (${existing.length}b)`); + } + + const all: Iter[] = []; + for (let i = 1; i <= CONFIG.iterations; i++) { + console.log(`\n--- iteration ${i} ---`); + const iter = await runIteration({ i, redisV2, ctx }); + for (const step of iter.steps) { + const extra = step.extra ? ` (${step.extra})` : ""; + console.log(` ${step.ms.toFixed(2).padStart(8)}ms ${step.name}${extra}`); + } + all.push(iter); + } + + summarize(all); + process.exit(0); +}; + +await main(); diff --git a/server/experiments/explainEntityAggregate.ts b/server/experiments/explainEntityAggregate.ts index eb70cde9a..5178b6ca6 100644 --- a/server/experiments/explainEntityAggregate.ts +++ b/server/experiments/explainEntityAggregate.ts @@ -31,7 +31,7 @@ const main = async () => { console.log("--- Running entity aggregate query ---"); const start = performance.now(); - const result = await getEntityAggregateForSync({ + const parsedResult = await getEntityAggregateForSync({ db, orgId, env, @@ -39,23 +39,19 @@ const main = async () => { internalFeatureIds: features?.map((feature) => feature.internal_id), }); const elapsed = performance.now() - start; - console.log(`Rows returned: ${result.length}`); - console.log(`Wall-clock time: ${elapsed.toFixed(2)}ms`); - console.log("Result:", JSON.stringify(result, null, 2)); - console.log(); + console.log(`Rows (after schema parse): ${parsedResult.length}`); + console.log(`Wall-clock time: ${elapsed.toFixed(2)}ms\n`); - // Build the same query inline for EXPLAIN ANALYZE + // Also fetch raw rows so we can see the output even if schema parse fails const { getEntityAggregateFragments } = await import( "../src/internal/customers/repos/getFullSubject/getEntityAggregateFragments" ); - - const statusFilter = sql`AND cp.status = ANY(ARRAY['active', 'past_due', 'scheduled'])`; - const entityFragments = getEntityAggregateFragments({ - statusFilter, + const statusFilterRaw = sql`AND cp.status = ANY(ARRAY['active', 'past_due', 'scheduled'])`; + const entityFragmentsRaw = getEntityAggregateFragments({ + statusFilter: statusFilterRaw, internalFeatureIds: features?.map((feature) => feature.internal_id), }); - - const query = sql` + const rawQuery = sql` WITH subject_customer_records AS ( SELECT * FROM customers c @@ -65,15 +61,36 @@ const main = async () => { ORDER BY (c.id = ${customerId}) DESC LIMIT 1 ) - - ${entityFragments.ctes} - + ${entityFragmentsRaw.ctes} SELECT * FROM entity_aggregated_cus_entitlements `; + const rawResult = await db.execute(rawQuery); + const result = rawResult as unknown as Record[]; + console.log(`Raw rows: ${result.length}\n`); + + console.log("--- Entity aggregate output (per feature) ---\n"); + for (const row of result) { + const entities = (row as unknown as { entities?: Record }) + .entities; + const entityCount = (row as unknown as { entity_count?: number }) + .entity_count; + const entityKeys = entities ? Object.keys(entities) : []; + + console.log(`feature_id: ${row.feature_id}`); + console.log(` internal_feature_id: ${row.internal_feature_id}`); + console.log(` balance: ${row.balance}`); + console.log(` adjustment: ${row.adjustment}`); + console.log(` additional_balance: ${row.additional_balance}`); + console.log(` rollover_balance: ${row.rollover_balance}`); + console.log(` entity_count: ${entityCount}`); + console.log(` entities keys (${entityKeys.length}): ${JSON.stringify(entityKeys)}`); + console.log(` entities: ${JSON.stringify(entities, null, 2)}`); + console.log(); + } console.log("--- EXPLAIN (ANALYZE, BUFFERS) ---\n"); - const explainQuery = sql`EXPLAIN (ANALYZE, BUFFERS, FORMAT TEXT) ${query}`; + const explainQuery = sql`EXPLAIN (ANALYZE, BUFFERS, FORMAT TEXT) ${rawQuery}`; const explainResult = await db.execute(explainQuery); for (const row of explainResult) { diff --git a/server/experiments/explainGetFullSubject.ts b/server/experiments/explainGetFullSubject.ts index b4bcdaa2e..741fa0d17 100644 --- a/server/experiments/explainGetFullSubject.ts +++ b/server/experiments/explainGetFullSubject.ts @@ -1,3 +1,5 @@ +import { writeFileSync } from "node:fs"; +import { resolve } from "node:path"; import { AppEnv } from "@autumn/shared"; import { sql } from "drizzle-orm"; import { @@ -43,17 +45,67 @@ const main = async () => { const elapsed = performance.now() - start; console.log(`Rows returned: ${result.length}`); console.log(`Wall-clock time: ${elapsed.toFixed(2)}ms`); - console.log("Result:", JSON.stringify(result, null, 2)); + + // The full row payload is enormous — dump it to a file and print a summary. + const outDir = resolve(import.meta.dir, "out"); + try { + const { mkdirSync } = await import("node:fs"); + mkdirSync(outDir, { recursive: true }); + } catch { + // ignore + } + const resultPath = resolve(outDir, "getFullSubject-result.json"); + writeFileSync(resultPath, JSON.stringify(result, null, 2)); + console.log(`Result written to: ${resultPath}`); + + const row0 = (result[0] ?? {}) as Record; + const summarize = (value: unknown): unknown => { + if (Array.isArray(value)) return `Array(len=${value.length})`; + if (value && typeof value === "object") + return `Object(keys=${Object.keys(value).length})`; + return value; + }; + console.log("\n--- Top-level fields (summary) ---"); + for (const [key, value] of Object.entries(row0)) { + console.log(` ${key}: ${summarize(value)}`); + } + + const aggregated = (row0.aggregated_customer_entitlements ?? []) as Array< + Record + >; + if (aggregated.length > 0) { + console.log( + `\n--- aggregated_customer_entitlements (${aggregated.length}) ---`, + ); + for (const ae of aggregated) { + const entities = (ae.entities ?? {}) as Record; + console.log( + ` ${ae.feature_id}: balance=${ae.balance}, adj=${ae.adjustment}, add=${ae.additional_balance}, rollover=${ae.rollover_balance}, entities.keys=${Object.keys(entities).length}`, + ); + } + } console.log(); console.log("--- EXPLAIN (ANALYZE, BUFFERS) ---\n"); const explainQuery = sql`EXPLAIN (ANALYZE, BUFFERS, FORMAT TEXT) ${query}`; const explainResult = await db.execute(explainQuery); + const explainPath = resolve(outDir, "getFullSubject-explain.txt"); + const explainLines: string[] = []; for (const row of explainResult) { const line = (row as Record)["QUERY PLAN"]; - console.log(line); + if (typeof line === "string") explainLines.push(line); } + writeFileSync(explainPath, explainLines.join("\n")); + console.log(`EXPLAIN written to: ${explainPath}`); + + // Print just the top of the plan (total time, first ~20 lines) inline + const head = explainLines.slice(0, 20); + for (const line of head) console.log(line); + const execLine = explainLines.find((line) => /Execution Time:/.test(line)); + const planLine = explainLines.find((line) => /Planning Time:/.test(line)); + if (planLine) console.log(planLine); + if (execLine) console.log(execLine); process.exit(0); }; diff --git a/server/src/internal/customers/repos/getFullSubject/getEntityAggregateFragments.ts b/server/src/internal/customers/repos/getFullSubject/getEntityAggregateFragments.ts index cae7defa7..df2e707bd 100644 --- a/server/src/internal/customers/repos/getFullSubject/getEntityAggregateFragments.ts +++ b/server/src/internal/customers/repos/getFullSubject/getEntityAggregateFragments.ts @@ -2,60 +2,39 @@ import { type SQL, sql } from "drizzle-orm"; import { getEntityOptionsAggregateFragments } from "./getEntityOptionsAggregateFragments.js"; /** - * Rollover CTEs driven from the shared `entity_product_cus_ents` and - * `entity_loose_cus_ents` base CTEs — avoids re-scanning `customer_entitlements` - * and `customer_products` per branch. + * Rollover CTEs: + * - entity_rollover_keys: per-entity rollover balances, built ONLY from + * jsonb_each(r.entities). Keys only exist in the map when a rollover row + * explicitly carries a per-entity breakdown. + * - entity_rollover_feature: feature-level rollover totals, summed directly + * off r.balance / r.usage (one add per active rollover). */ const buildEntityRolloverCtes = () => sql` - entity_rollover_rows AS ( - -- Top-level rollover: one row per (rollover × cus_ent). - -- entity_key = cp.internal_entity_id for product-attached, ce.internal_entity_id for loose. - SELECT - ce.internal_feature_id, - ce.internal_customer_id, - COALESCE(ce.cp_entity_key, ce.internal_entity_id) AS entity_key, - r.balance::numeric AS rollover_balance, - COALESCE(r.usage, 0)::numeric AS rollover_usage - FROM rollovers r - JOIN entity_level_cus_ents ce ON r.cus_ent_id = ce.id - WHERE (r.expires_at IS NULL OR r.expires_at > EXTRACT(EPOCH FROM now()) * 1000) - - UNION ALL - - -- Per-entity rollover from jsonb_each(r.entities). + entity_rollover_keys AS ( SELECT ce.internal_feature_id, ce.internal_customer_id, kv.entity_key AS entity_key, - COALESCE((kv.entity_value->>'balance')::numeric, 0) AS rollover_balance, - COALESCE((kv.entity_value->>'usage')::numeric, 0) AS rollover_usage + SUM(COALESCE((kv.entity_value->>'balance')::numeric, 0)) AS rollover_balance, + SUM(COALESCE((kv.entity_value->>'usage')::numeric, 0)) AS rollover_usage FROM rollovers r JOIN entity_level_cus_ents ce ON r.cus_ent_id = ce.id CROSS JOIN LATERAL jsonb_each(r.entities) AS kv(entity_key, entity_value) WHERE jsonb_typeof(r.entities) = 'object' AND (r.expires_at IS NULL OR r.expires_at > EXTRACT(EPOCH FROM now()) * 1000) - ), - - entity_rollover_keys AS ( - SELECT - internal_feature_id, - internal_customer_id, - entity_key, - SUM(rollover_balance) AS rollover_balance, - SUM(rollover_usage) AS rollover_usage - FROM entity_rollover_rows - WHERE entity_key IS NOT NULL - GROUP BY internal_feature_id, internal_customer_id, entity_key + GROUP BY ce.internal_feature_id, ce.internal_customer_id, kv.entity_key ), entity_rollover_feature AS ( SELECT - internal_feature_id, - internal_customer_id, - SUM(rollover_balance) AS rollover_balance, - SUM(rollover_usage) AS rollover_usage - FROM entity_rollover_rows - GROUP BY internal_feature_id, internal_customer_id + ce.internal_feature_id, + ce.internal_customer_id, + SUM(r.balance::numeric) AS rollover_balance, + SUM(COALESCE(r.usage, 0)::numeric) AS rollover_usage + FROM rollovers r + JOIN entity_level_cus_ents ce ON r.cus_ent_id = ce.id + WHERE (r.expires_at IS NULL OR r.expires_at > EXTRACT(EPOCH FROM now()) * 1000) + GROUP BY ce.internal_feature_id, ce.internal_customer_id ) `; @@ -145,65 +124,23 @@ export const getEntityAggregateFragments = ({ ${entityOptionsAggregateFragments.ctes}, - entity_balance_rows AS ( - -- Top-level: one row per cus_ent (product-attached or loose). - -- entity_key = cp.internal_entity_id for product-attached, ce.internal_entity_id for loose. - SELECT - COALESCE(ce.external_id, ce.id) AS api_id, - ce.internal_feature_id, - ce.internal_customer_id, - ce.feature_id, - COALESCE(ent.allowance, 0)::numeric AS allowance, - ce.balance::numeric AS balance, - COALESCE(ce.adjustment, 0)::numeric AS adjustment, - COALESCE(ce.additional_balance, 0)::numeric AS additional_balance, - ce.unlimited, - ce.usage_allowed, - COALESCE(ce.cp_entity_key, ce.internal_entity_id) AS entity_key, - ce.balance::numeric AS entity_balance, - COALESCE(ce.adjustment, 0)::numeric AS entity_adjustment, - COALESCE(ce.additional_balance, 0)::numeric AS entity_additional_balance - FROM entity_level_cus_ents ce - JOIN entitlements ent ON ce.entitlement_id = ent.id - - UNION ALL - - -- Per-entity: N rows per cus_ent from jsonb_each(ce.entities). - -- balance/adj/additional = 0 to avoid double-counting at the aggregate level. - SELECT - COALESCE(ce.external_id, ce.id) AS api_id, - ce.internal_feature_id, - ce.internal_customer_id, - ce.feature_id, - COALESCE(ent.allowance, 0)::numeric AS allowance, - 0::numeric AS balance, - 0::numeric AS adjustment, - 0::numeric AS additional_balance, - ce.unlimited, - ce.usage_allowed, - kv.entity_key AS entity_key, - (kv.entity_value->>'balance')::numeric AS entity_balance, - COALESCE((kv.entity_value->>'adjustment')::numeric, 0) AS entity_adjustment, - COALESCE((kv.entity_value->>'additional_balance')::numeric, 0) AS entity_additional_balance - FROM entity_level_cus_ents ce - JOIN entitlements ent ON ce.entitlement_id = ent.id - CROSS JOIN LATERAL jsonb_each(ce.entities) AS kv(entity_key, entity_value) - WHERE jsonb_typeof(ce.entities) = 'object' - ), - ${buildEntityRolloverCtes()}, + -- Per-entity balance map: built ONLY from jsonb_each(ce.entities) so + -- that entities keys reflect real per-entity breakdowns, not every + -- product-attached cus_ent. Summed across cus_ents for the same key. entity_balance_keys AS ( SELECT - internal_feature_id, - internal_customer_id, - entity_key, - SUM(entity_balance) AS balance, - SUM(entity_adjustment) AS adjustment, - SUM(entity_additional_balance) AS additional_balance - FROM entity_balance_rows - WHERE entity_key IS NOT NULL - GROUP BY internal_feature_id, internal_customer_id, entity_key + ce.internal_feature_id, + ce.internal_customer_id, + kv.entity_key, + SUM((kv.entity_value->>'balance')::numeric) AS balance, + SUM(COALESCE((kv.entity_value->>'adjustment')::numeric, 0)) AS adjustment, + SUM(COALESCE((kv.entity_value->>'additional_balance')::numeric, 0)) AS additional_balance + FROM entity_level_cus_ents ce + CROSS JOIN LATERAL jsonb_each(ce.entities) AS kv(entity_key, entity_value) + WHERE jsonb_typeof(ce.entities) = 'object' + GROUP BY ce.internal_feature_id, ce.internal_customer_id, kv.entity_key ), entity_aggregate_map AS ( @@ -250,21 +187,23 @@ export const getEntityAggregateFragments = ({ GROUP BY ejk.internal_feature_id, ejk.internal_customer_id ), + -- Feature-level totals: summed directly off entity_level_cus_ents + -- (one add per cus_ent). entity_aggregate_totals AS ( SELECT - MIN(ebr.api_id) AS api_id, - ebr.internal_feature_id, - ebr.internal_customer_id, - MIN(ebr.feature_id) AS feature_id, - SUM(ebr.allowance) AS allowance_total, - SUM(ebr.balance) AS balance, - SUM(ebr.adjustment) AS adjustment, - SUM(ebr.additional_balance) AS additional_balance, - BOOL_OR(ebr.unlimited) AS unlimited, - BOOL_OR(ebr.usage_allowed) AS usage_allowed, - COUNT(DISTINCT ebr.entity_key) FILTER (WHERE ebr.entity_key IS NOT NULL) AS entity_count - FROM entity_balance_rows ebr - GROUP BY ebr.internal_feature_id, ebr.internal_customer_id + MIN(COALESCE(ce.external_id, ce.id)) AS api_id, + ce.internal_feature_id, + ce.internal_customer_id, + MIN(ce.feature_id) AS feature_id, + SUM(COALESCE(ent.allowance, 0)::numeric) AS allowance_total, + SUM(ce.balance::numeric) AS balance, + SUM(COALESCE(ce.adjustment, 0)::numeric) AS adjustment, + SUM(COALESCE(ce.additional_balance, 0)::numeric) AS additional_balance, + BOOL_OR(ce.unlimited) AS unlimited, + BOOL_OR(ce.usage_allowed) AS usage_allowed + FROM entity_level_cus_ents ce + JOIN entitlements ent ON ce.entitlement_id = ent.id + GROUP BY ce.internal_feature_id, ce.internal_customer_id ), entity_aggregated_cus_entitlements AS ( @@ -282,7 +221,6 @@ export const getEntityAggregateFragments = ({ COALESCE(erf.rollover_usage, 0) AS rollover_usage, eat.unlimited, eat.usage_allowed, - eat.entity_count, eam.entities FROM entity_aggregate_totals eat LEFT JOIN entity_aggregate_map eam diff --git a/server/tests/integration/db/full-subject/full-subject-aggregate-options.test.ts b/server/tests/integration/db/full-subject/full-subject-aggregate-options.test.ts index a34dff9de..4af94e744 100644 --- a/server/tests/integration/db/full-subject/full-subject-aggregate-options.test.ts +++ b/server/tests/integration/db/full-subject/full-subject-aggregate-options.test.ts @@ -236,7 +236,6 @@ describe(`${chalk.yellowBright("fullSubject aggregate options")}`, () => { expect(aggregateMessagesBalance?.allowance_total).toBe(600); expect(aggregateMessagesBalance?.prepaid_grant_from_options).toBe(1000); expect(aggregateMessagesBalance?.balance).toBe(1600); - expect(aggregateMessagesBalance?.entity_count).toBe(2); }, }); }); diff --git a/server/tests/integration/db/full-subject/full-subject-aggregates.test.ts b/server/tests/integration/db/full-subject/full-subject-aggregates.test.ts index f6d104110..ef282cf06 100644 --- a/server/tests/integration/db/full-subject/full-subject-aggregates.test.ts +++ b/server/tests/integration/db/full-subject/full-subject-aggregates.test.ts @@ -106,7 +106,6 @@ describe(`${chalk.yellowBright("fullSubject aggregates")}`, () => { expect(aggregate.balance).toBe(30); expect(aggregate.adjustment).toBe(0); expect(aggregate.additional_balance).toBe(0); - expect(aggregate.entity_count).toBe(2); expect(JSON.stringify(aggregate.feature.id)).toBe( JSON.stringify(scenario.customerEntitlements[1]!.feature_id), ); @@ -157,7 +156,6 @@ describe(`${chalk.yellowBright("fullSubject aggregates")}`, () => { const aggregate = comparable.aggregated_customer_entitlements[0]!; expect(aggregate.allowance_total).toBe(200); expect(aggregate.balance).toBe(30); - expect(aggregate.entity_count).toBe(2); expect(aggregate.entities).toMatchObject({ [scenario.ids.internalEntityIds[0]!]: { id: scenario.ids.internalEntityIds[0], diff --git a/server/tests/integration/db/full-subject/utils/buildComparableFullSubject.ts b/server/tests/integration/db/full-subject/utils/buildComparableFullSubject.ts index 8399f3bf8..20ee9d657 100644 --- a/server/tests/integration/db/full-subject/utils/buildComparableFullSubject.ts +++ b/server/tests/integration/db/full-subject/utils/buildComparableFullSubject.ts @@ -148,7 +148,6 @@ const comparableAggregatedEntitlement = ( additional_balance: customerEntitlement.additional_balance, unlimited: customerEntitlement.unlimited, usage_allowed: customerEntitlement.usage_allowed, - entity_count: customerEntitlement.entity_count, entities: customerEntitlement.entities, feature: { id: customerEntitlement.feature.id, diff --git a/server/tests/unit/full-subject-cache/full-subject-aggregate-balance.test.ts b/server/tests/unit/full-subject-cache/full-subject-aggregate-balance.test.ts index ce000e2ba..17951a41a 100644 --- a/server/tests/unit/full-subject-cache/full-subject-aggregate-balance.test.ts +++ b/server/tests/unit/full-subject-cache/full-subject-aggregate-balance.test.ts @@ -22,7 +22,6 @@ describe("fullSubject aggregate balance", () => { rollover_usage: 0, unlimited: false, usage_allowed: false, - entity_count: 2, entities: { ent1: { id: "ent1", diff --git a/shared/models/cusProductModels/cusEntModels/aggregatedCusEnt.ts b/shared/models/cusProductModels/cusEntModels/aggregatedCusEnt.ts index 6af85017a..f890a3865 100644 --- a/shared/models/cusProductModels/cusEntModels/aggregatedCusEnt.ts +++ b/shared/models/cusProductModels/cusEntModels/aggregatedCusEnt.ts @@ -24,7 +24,6 @@ export const AggregatedFeatureBalanceSchema = z.object({ rollover_usage: z.number().default(0), unlimited: z.boolean(), usage_allowed: z.boolean(), - entity_count: z.number(), entities: z.record(z.string(), AggregatedEntityBalanceSchema).nullish(), }); diff --git a/shared/utils/cusUtils/fullCusUtils/logFullCustomer.ts b/shared/utils/cusUtils/fullCusUtils/logFullCustomer.ts index d96ed7157..a9a405c8a 100644 --- a/shared/utils/cusUtils/fullCusUtils/logFullCustomer.ts +++ b/shared/utils/cusUtils/fullCusUtils/logFullCustomer.ts @@ -42,7 +42,6 @@ export const logFullCustomer = ({ balance: ae.balance, adjustment: ae.adjustment, unlimited: ae.unlimited, - entity_count: ae.entity_count, entitlement: "entitlement" in ae ? (ae as Record).entitlement diff --git a/shared/utils/fullSubjectUtils/logFullSubject.ts b/shared/utils/fullSubjectUtils/logFullSubject.ts index 74da8246d..008907c2e 100644 --- a/shared/utils/fullSubjectUtils/logFullSubject.ts +++ b/shared/utils/fullSubjectUtils/logFullSubject.ts @@ -62,7 +62,6 @@ export const logFullSubject = ({ feature_id: ae.feature_id, balance: ae.balance, unlimited: ae.unlimited, - entity_count: ae.entity_count, })) ?? "N/A", };