diff --git a/server/src/db/initializeDatabaseFunctions.ts b/server/src/db/initializeDatabaseFunctions.ts index d6033c655..bef7109e8 100644 --- a/server/src/db/initializeDatabaseFunctions.ts +++ b/server/src/db/initializeDatabaseFunctions.ts @@ -23,6 +23,7 @@ export const initializeDatabaseFunctions = async () => { // Helper functions "deductFromRollovers.sql", "deductFromMainBalance.sql", + "unwindFromLockReceipt.sql", "getTotalBalance.sql", "deductFromAdditionalBalance.sql", "performDeduction.sql", diff --git a/server/src/external/autumn/autumnCli.ts b/server/src/external/autumn/autumnCli.ts index 0d8415d24..76f8f6b65 100644 --- a/server/src/external/autumn/autumnCli.ts +++ b/server/src/external/autumn/autumnCli.ts @@ -850,8 +850,20 @@ export class AutumnInt { const data = await this.post(`/balances.delete`, params); return data; }, - finalize: async (params: FinalizeLockParamsV0) => { - const data = await this.post(`/balances.finalize`, params); + finalize: async ( + params: FinalizeLockParamsV0, + { + skipCache = false, + headers, + }: { + skipCache?: boolean; + headers?: Record; + } = {}, + ) => { + const data = await this.post(`/balances.finalize`, params, { + ...(skipCache ? { "x-skip-cache": "true" } : {}), + ...headers, + }); return data; }, }; diff --git a/server/src/honoMiddlewares/baseMiddleware.ts b/server/src/honoMiddlewares/baseMiddleware.ts index a597fc2f5..7b72baec3 100644 --- a/server/src/honoMiddlewares/baseMiddleware.ts +++ b/server/src/honoMiddlewares/baseMiddleware.ts @@ -65,7 +65,7 @@ export const baseMiddleware = async (c: Context, next: Next) => { // Query params expand: [], - skipCache: false, + skipCache: c.req.header("x-skip-cache") === "true", // Test params: extraLogs: {}, diff --git a/server/src/honoMiddlewares/expandMiddleware.ts b/server/src/honoMiddlewares/expandMiddleware.ts index 4ceb31cee..f9f3b2785 100644 --- a/server/src/honoMiddlewares/expandMiddleware.ts +++ b/server/src/honoMiddlewares/expandMiddleware.ts @@ -64,7 +64,7 @@ export const expandMiddleware = (): MiddlewareHandler => { c.set("ctx", { ...ctx, expand, - skipCache: skipCacheValue, + skipCache: ctx.skipCache || skipCacheValue, }); await next(); diff --git a/server/src/internal/balances/finalizeLock/finalizeLock.ts b/server/src/internal/balances/finalizeLock/finalizeLock.ts index 6d4f69f43..e9ab4acbe 100644 --- a/server/src/internal/balances/finalizeLock/finalizeLock.ts +++ b/server/src/internal/balances/finalizeLock/finalizeLock.ts @@ -1,6 +1,11 @@ -import { type FinalizeLockParamsV0, findFeatureById } from "@autumn/shared"; +import { + type FinalizeLockParamsV0, + findFeatureById, + tryCatch, +} from "@autumn/shared"; import { currentRegion } from "@/external/redis/initRedis.js"; import type { AutumnContext } from "@/honoUtils/HonoEnv"; +import { executePostgresDeduction } from "@/internal/balances/utils/deduction/executePostgresDeduction.js"; import { executeRedisDeduction } from "@/internal/balances/utils/deduction/executeRedisDeduction.js"; import { fetchLockReceipt } from "@/internal/balances/utils/lock/fetchLockReceipt.js"; import { calculateUnwindValue } from "@/internal/balances/utils/lock/unwindLockUtils.js"; @@ -8,6 +13,7 @@ import { deductionUpdatesToModifiedIds } from "@/internal/balances/utils/sync/de import { globalSyncBatchingManagerV2 } from "@/internal/balances/utils/sync/SyncBatchingManagerV2.js"; import type { DeductionUpdate } from "@/internal/balances/utils/types/deductionUpdate.js"; import type { FeatureDeduction } from "@/internal/balances/utils/types/featureDeduction.js"; +import { RedisDeductionError } from "@/internal/balances/utils/types/redisDeductionError.js"; import type { RolloverUpdate } from "@/internal/balances/utils/types/rolloverUpdate.js"; import { getOrSetCachedFullCustomer } from "@/internal/customers/cusUtils/fullCustomerCacheUtils/getOrSetCachedFullCustomer.js"; @@ -74,21 +80,51 @@ export const finalizeLock = async ({ const deduction: FeatureDeduction = { feature, deduction: additionalValue, + lockReceipt: receipt, // For unwinding when finalizing a lock unwindValue, lockReceiptKey, }; - const { updates, rolloverUpdates } = await executeRedisDeduction({ - ctx, - fullCustomer, - entityId: receipt.entity_id ?? undefined, - deductions: [deduction], - deductionOptions: { - triggerAutoTopUp: true, - }, - }); + const deductionOptions = { + triggerAutoTopUp: true, + }; + + const { data: redisResult, error } = await tryCatch( + executeRedisDeduction({ + ctx, + fullCustomer, + entityId: receipt.entity_id ?? undefined, + deductions: [deduction], + deductionOptions, + }), + ); + + if (error) { + if (error instanceof RedisDeductionError && error.shouldFallback()) { + ctx.logger.warn( + `Falling back to Postgres for finalize lock: ${error.code}`, + ); + + await executePostgresDeduction({ + ctx, + fullCustomer, + customerId: receipt.customer_id, + entityId: receipt.entity_id ?? undefined, + deductions: [deduction], + options: deductionOptions, + }); + + return { + success: true, + }; + } + + throw error; + } + + const { updates, rolloverUpdates } = redisResult; queueSyncItem({ ctx, diff --git a/server/src/internal/balances/utils/deduction/executePostgresDeduction.ts b/server/src/internal/balances/utils/deduction/executePostgresDeduction.ts index 9dc209eac..06bd307cc 100644 --- a/server/src/internal/balances/utils/deduction/executePostgresDeduction.ts +++ b/server/src/internal/balances/utils/deduction/executePostgresDeduction.ts @@ -11,6 +11,7 @@ import type { AutumnContext } from "../../../../honoUtils/HonoEnv.js"; import { CusService } from "../../../customers/CusService.js"; import type { EventInfo } from "../../events/initEvent.js"; import { applyDeductionUpdateToFullCustomer } from "../../utils/deduction/applyDeductionUpdateToFullCustomer.js"; +import { saveLockReceipt } from "../../utils/lock/saveLockReceipt.js"; import type { DeductionUpdate } from "../../utils/types/deductionUpdate.js"; import type { FeatureDeduction } from "../../utils/types/featureDeduction.js"; import type { MutationLogItem } from "../../utils/types/mutationLogItem.js"; @@ -76,21 +77,30 @@ export const executePostgresDeduction = async ({ deductions, }); - const executeDeduction = async (): Promise< - Record - > => { + const executeDeduction = async (): Promise<{ + updates: Record; + mutationLogs: MutationLogItem[]; + }> => { let allUpdates: Record = {}; let allRolloverOverwrites: RolloverOverwrite[] = []; + let allMutationLogs: MutationLogItem[] = []; // Need to deduct from customer entitlement... for (const deduction of deductions) { - const { feature, deduction: toDeduct, targetBalance } = deduction; + const { + feature, + deduction: toDeduct, + targetBalance, + lockReceipt, + unwindValue, + } = deduction; const { customerEntitlementDeductions, rollovers, customerEntitlements, unlimitedFeatureIds, + lock: preparedLock, } = prepareFeatureDeduction({ ctx, fullCustomer, @@ -108,6 +118,8 @@ export const executePostgresDeduction = async ({ sorted_entitlements: customerEntitlementDeductions, amount_to_deduct: toDeduct ?? null, target_balance: targetBalance ?? null, + lock_receipt: lockReceipt ?? null, + unwind_value: unwindValue ?? null, target_entity_id: entityId || null, rollovers: rollovers.length > 0 ? rollovers : null, cus_ent_ids: customerEntitlements.map((ce) => ce.id), @@ -124,6 +136,7 @@ export const executePostgresDeduction = async ({ updates: Record; remaining: number; rollover_updates: RolloverOverwrite[]; + mutation_logs: MutationLogItem[]; }; if (!resultJson) { @@ -132,7 +145,7 @@ export const executePostgresDeduction = async ({ }); } - const { updates, rollover_updates } = resultJson; + const { updates, rollover_updates, mutation_logs } = resultJson; logDeductionUpdates({ ctx, fullCustomer, @@ -140,6 +153,7 @@ export const executePostgresDeduction = async ({ source: "executePostgresDeduction", }); allUpdates = { ...allUpdates, ...updates }; + allMutationLogs = [...allMutationLogs, ...(mutation_logs ?? [])]; if (rollover_updates?.length > 0) { allRolloverOverwrites = [...allRolloverOverwrites, ...rollover_updates]; } @@ -171,6 +185,16 @@ export const executePostgresDeduction = async ({ update, }); } + + if (preparedLock?.enabled) { + await saveLockReceipt({ + lock: preparedLock, + customerId: fullCustomer.id || customerId, + featureId: feature.id, + entityId, + items: mutation_logs ?? [], + }); + } } catch (error) { if (error instanceof Error && !error?.message?.includes("declined")) { ctx.logger.error( @@ -220,10 +244,13 @@ export const executePostgresDeduction = async ({ rolloverOverwrites: allRolloverOverwrites, }); - return allUpdates; + return { + updates: allUpdates, + mutationLogs: allMutationLogs, + }; }; - const allUpdates = resolvedOptions.paidAllocated + const deductionResult = resolvedOptions.paidAllocated ? await withLock({ lockKey: `lock:deduction:${org.id}:${env}:${customerId}`, ttlMs: 60000, @@ -235,7 +262,7 @@ export const executePostgresDeduction = async ({ return { oldFullCus, fullCus: fullCustomer, - updates: allUpdates, - mutationLogs: [], + updates: deductionResult.updates, + mutationLogs: deductionResult.mutationLogs, }; }; diff --git a/server/src/internal/balances/utils/lock/saveLockReceipt.ts b/server/src/internal/balances/utils/lock/saveLockReceipt.ts new file mode 100644 index 000000000..a54ff92ff --- /dev/null +++ b/server/src/internal/balances/utils/lock/saveLockReceipt.ts @@ -0,0 +1,50 @@ +import { InternalError } from "@autumn/shared"; +import { redis } from "@/external/redis/initRedis.js"; +import type { MutationLogItem } from "@/internal/balances/utils/types/mutationLogItem.js"; +import { tryRedisWrite } from "@/utils/cacheUtils/cacheUtils.js"; + +export const saveLockReceipt = async ({ + lock, + customerId, + featureId, + entityId, + items, +}: { + lock: { + key?: string; + hashed_key?: string; + expires_at?: string; + redis_receipt_key: string; + created_at: number; + }; + customerId: string; + featureId: string; + entityId?: string; + items: MutationLogItem[]; +}) => { + const result = await tryRedisWrite( + () => + redis.call( + "JSON.SET", + lock.redis_receipt_key, + "$", + JSON.stringify({ + lock_key: lock.key ?? null, + hashed_key: lock.hashed_key ?? null, + status: "pending", + customer_id: customerId, + feature_id: featureId, + entity_id: entityId ?? null, + expires_at: lock.expires_at ?? null, + created_at: lock.created_at, + items, + }), + ) as Promise<"OK" | null>, + ); + + if (result === "OK") return; + + throw new InternalError({ + message: `Failed to save lock receipt for key: ${lock.key ?? lock.hashed_key}`, + }); +}; diff --git a/server/src/internal/balances/utils/sql/deductFromMainBalance.sql b/server/src/internal/balances/utils/sql/deductFromMainBalance.sql index 80bc5b018..4a571d08e 100644 --- a/server/src/internal/balances/utils/sql/deductFromMainBalance.sql +++ b/server/src/internal/balances/utils/sql/deductFromMainBalance.sql @@ -7,12 +7,14 @@ RETURNS TABLE ( deducted numeric, new_balance numeric, new_entities jsonb, - new_adjustment numeric + new_adjustment numeric, + mutation_logs jsonb ) LANGUAGE plpgsql AS $$ DECLARE -- Extract parameters from JSONB + customer_entitlement_id text := NULLIF(params->>'customer_entitlement_id', ''); current_balance numeric := (params->>'current_balance')::numeric; current_entities jsonb := COALESCE(params->'current_entities', '{}'::jsonb); current_adjustment numeric := COALESCE((params->>'current_adjustment')::numeric, 0); @@ -48,6 +50,7 @@ DECLARE entity_adjustment numeric; ceiling numeric; max_addable numeric; + mutation_logs_json jsonb := '[]'::jsonb; BEGIN -- Initialize return values @@ -112,6 +115,20 @@ BEGIN to_jsonb(COALESCE((result_entities->entity_key->>'adjustment')::numeric, 0) - deduct_amount) ); END IF; + + mutation_logs_json := mutation_logs_json || jsonb_build_array( + jsonb_build_object( + 'target_type', 'customer_entitlement', + 'customer_entitlement_id', customer_entitlement_id, + 'rollover_id', NULL, + 'entity_id', entity_key, + 'credit_cost', credit_cost, + 'balance_delta', -deduct_amount, + 'adjustment_delta', CASE WHEN alter_granted_balance THEN -deduct_amount ELSE 0 END, + 'usage_delta', 0, + 'value_delta', deduct_amount / credit_cost + ) + ); remaining := remaining - deduct_amount; deducted_amount := deducted_amount + deduct_amount; @@ -170,6 +187,20 @@ BEGIN to_jsonb(COALESCE((result_entities->target_entity_id->>'adjustment')::numeric, 0) - deducted_amount) ); END IF; + + mutation_logs_json := mutation_logs_json || jsonb_build_array( + jsonb_build_object( + 'target_type', 'customer_entitlement', + 'customer_entitlement_id', customer_entitlement_id, + 'rollover_id', NULL, + 'entity_id', target_entity_id, + 'credit_cost', credit_cost, + 'balance_delta', -deducted_amount, + 'adjustment_delta', CASE WHEN alter_granted_balance THEN -deducted_amount ELSE 0 END, + 'usage_delta', 0, + 'value_delta', deducted_amount / credit_cost + ) + ); ELSE result_entities := current_entities; END IF; @@ -215,10 +246,31 @@ BEGIN IF alter_granted_balance THEN result_adjustment := result_adjustment - deducted_amount; END IF; + + IF deducted_amount != 0 THEN + mutation_logs_json := mutation_logs_json || jsonb_build_array( + jsonb_build_object( + 'target_type', 'customer_entitlement', + 'customer_entitlement_id', customer_entitlement_id, + 'rollover_id', NULL, + 'entity_id', NULL, + 'credit_cost', credit_cost, + 'balance_delta', -deducted_amount, + 'adjustment_delta', CASE WHEN alter_granted_balance THEN -deducted_amount ELSE 0 END, + 'usage_delta', 0, + 'value_delta', deducted_amount / credit_cost + ) + ); + END IF; END IF; -- Return results - RETURN QUERY SELECT deducted_amount, result_balance, result_entities, result_adjustment; + RETURN QUERY + SELECT + deducted_amount, + result_balance, + result_entities, + result_adjustment, + mutation_logs_json; END; $$; - diff --git a/server/src/internal/balances/utils/sql/deductFromRollovers.sql b/server/src/internal/balances/utils/sql/deductFromRollovers.sql index fda0ace7b..088153bd1 100644 --- a/server/src/internal/balances/utils/sql/deductFromRollovers.sql +++ b/server/src/internal/balances/utils/sql/deductFromRollovers.sql @@ -4,7 +4,7 @@ DROP FUNCTION IF EXISTS deduct_from_rollovers(jsonb); CREATE FUNCTION deduct_from_rollovers(params jsonb) -RETURNS TABLE(total_deducted numeric) +RETURNS TABLE(total_deducted numeric, mutation_logs jsonb) LANGUAGE plpgsql AS $$ DECLARE @@ -23,6 +23,7 @@ DECLARE current_balance numeric; current_usage numeric; current_entities jsonb; + current_cus_ent_id text; entity_key text; entity_balance numeric; @@ -33,6 +34,7 @@ DECLARE new_usage numeric; new_entities jsonb; rollover_total_deducted_features numeric := 0; + mutation_logs_json jsonb := '[]'::jsonb; BEGIN -- Normalize input: if rollovers array provided use it, otherwise convert rollover_ids to same format IF params->'rollovers' IS NOT NULL AND jsonb_typeof(params->'rollovers') = 'array' AND jsonb_array_length(params->'rollovers') > 0 THEN @@ -45,13 +47,13 @@ BEGIN FROM unnest(rollover_ids) AS id; ELSE -- No rollovers to process - RETURN QUERY SELECT 0::numeric; + RETURN QUERY SELECT 0::numeric, '[]'::jsonb; RETURN; END IF; -- Early return if no amount IF remaining_amount <= 0 THEN - RETURN QUERY SELECT 0::numeric; + RETURN QUERY SELECT 0::numeric, '[]'::jsonb; RETURN; END IF; @@ -64,8 +66,8 @@ BEGIN credit_cost := COALESCE((rollover_obj->>'credit_cost')::numeric, 1); -- Lock and fetch rollover data - SELECT r.balance, COALESCE(r.usage, 0), r.entities - INTO current_balance, current_usage, current_entities + SELECT r.balance, COALESCE(r.usage, 0), r.entities, r.cus_ent_id + INTO current_balance, current_usage, current_entities, current_cus_ent_id FROM rollovers r WHERE r.id = rollover_id FOR UPDATE; @@ -99,6 +101,19 @@ BEGIN WHERE r.id = rollover_id; feature_deduct_amount := credit_deduct_amount / credit_cost; + mutation_logs_json := mutation_logs_json || jsonb_build_array( + jsonb_build_object( + 'target_type', 'rollover', + 'customer_entitlement_id', current_cus_ent_id, + 'rollover_id', rollover_id, + 'entity_id', target_entity_id, + 'credit_cost', credit_cost, + 'balance_delta', -credit_deduct_amount, + 'adjustment_delta', 0, + 'usage_delta', credit_deduct_amount, + 'value_delta', feature_deduct_amount + ) + ); remaining_amount := remaining_amount - feature_deduct_amount; rollover_total_deducted_features := rollover_total_deducted_features + feature_deduct_amount; END IF; @@ -132,6 +147,19 @@ BEGIN ); feature_deduct_amount := credit_deduct_amount / credit_cost; + mutation_logs_json := mutation_logs_json || jsonb_build_array( + jsonb_build_object( + 'target_type', 'rollover', + 'customer_entitlement_id', current_cus_ent_id, + 'rollover_id', rollover_id, + 'entity_id', entity_key, + 'credit_cost', credit_cost, + 'balance_delta', -credit_deduct_amount, + 'adjustment_delta', 0, + 'usage_delta', credit_deduct_amount, + 'value_delta', feature_deduct_amount + ) + ); remaining_amount := remaining_amount - feature_deduct_amount; rollover_total_deducted_features := rollover_total_deducted_features + feature_deduct_amount; END IF; @@ -153,12 +181,25 @@ BEGIN WHERE r.id = rollover_id; feature_deduct_amount := credit_deduct_amount / credit_cost; + mutation_logs_json := mutation_logs_json || jsonb_build_array( + jsonb_build_object( + 'target_type', 'rollover', + 'customer_entitlement_id', current_cus_ent_id, + 'rollover_id', rollover_id, + 'entity_id', NULL, + 'credit_cost', credit_cost, + 'balance_delta', -credit_deduct_amount, + 'adjustment_delta', 0, + 'usage_delta', credit_deduct_amount, + 'value_delta', feature_deduct_amount + ) + ); remaining_amount := remaining_amount - feature_deduct_amount; rollover_total_deducted_features := rollover_total_deducted_features + feature_deduct_amount; END IF; END IF; END LOOP; - RETURN QUERY SELECT rollover_total_deducted_features; + RETURN QUERY SELECT rollover_total_deducted_features, mutation_logs_json; END; $$; diff --git a/server/src/internal/balances/utils/sql/performDeduction.sql b/server/src/internal/balances/utils/sql/performDeduction.sql index 717e9ac2a..a370ed57f 100644 --- a/server/src/internal/balances/utils/sql/performDeduction.sql +++ b/server/src/internal/balances/utils/sql/performDeduction.sql @@ -21,10 +21,15 @@ DECLARE END; -- New: rollovers array with {id, credit_cost} objects for credit cost support rollovers_arr jsonb := params->'rollovers'; + lock_receipt jsonb := params->'lock_receipt'; cus_ent_ids text[] := CASE WHEN params->'cus_ent_ids' IS NULL OR jsonb_typeof(params->'cus_ent_ids') != 'array' THEN NULL ELSE ARRAY(SELECT jsonb_array_elements_text(params->'cus_ent_ids')) END; + unwind_value numeric := CASE + WHEN params->>'unwind_value' IS NULL THEN NULL + ELSE (params->>'unwind_value')::numeric + END; skip_additional_balance boolean := COALESCE((params->>'skip_additional_balance')::boolean, false); alter_granted_balance boolean := COALESCE((params->>'alter_granted_balance')::boolean, false); overage_behaviour text := NULLIF(params->>'overage_behaviour', ''); @@ -61,13 +66,19 @@ DECLARE new_balance numeric; new_entities jsonb; new_adjustment numeric; + step_mutation_logs jsonb := '[]'::jsonb; + unwind_updates_json jsonb := '{}'::jsonb; + unwind_modified_rollover_ids text[] := ARRAY[]::text[]; + unwind_remaining_value numeric := 0; -- Tracking updates_json jsonb := '{}'::jsonb; rollover_updates_json jsonb := '[]'::jsonb; + mutation_logs_json jsonb := '[]'::jsonb; result_json jsonb; -- For calculating total balance total_balance numeric; + final_rollover_ids text[] := ARRAY[]::text[]; BEGIN -- Compute overage_behavior_is_allow once (used for cap bypass in deduct_from_main_balance) overage_behavior_is_allow := alter_granted_balance OR overage_behaviour = 'allow'; @@ -88,6 +99,45 @@ BEGIN PERFORM 1 FROM rollovers r WHERE r.id IN (SELECT jsonb_array_elements(rollovers_arr)->>'id') FOR UPDATE; END IF; + -- Lock rows referenced by the lock receipt too. + IF lock_receipt IS NOT NULL + AND jsonb_typeof(lock_receipt->'items') = 'array' + AND jsonb_array_length(lock_receipt->'items') > 0 + THEN + PERFORM 1 + FROM customer_entitlements ce + WHERE ce.id IN ( + SELECT DISTINCT item->>'customer_entitlement_id' + FROM jsonb_array_elements(lock_receipt->'items') item + WHERE NULLIF(item->>'customer_entitlement_id', '') IS NOT NULL + ) + FOR UPDATE; + + PERFORM 1 + FROM rollovers r + WHERE r.id IN ( + SELECT DISTINCT item->>'rollover_id' + FROM jsonb_array_elements(lock_receipt->'items') item + WHERE NULLIF(item->>'rollover_id', '') IS NOT NULL + ) + FOR UPDATE; + END IF; + + -- ============================================================================ + -- STEP 0.5: Unwind lock receipt before running the regular deduction flow + -- ============================================================================ + IF lock_receipt IS NOT NULL AND COALESCE(unwind_value, 0) > 0 THEN + SELECT * + INTO unwind_remaining_value, unwind_updates_json, unwind_modified_rollover_ids, step_mutation_logs + FROM unwind_from_lock_receipt(jsonb_build_object( + 'lock_receipt', lock_receipt, + 'unwind_value', unwind_value + )); + + updates_json := updates_json || COALESCE(unwind_updates_json, '{}'::jsonb); + mutation_logs_json := mutation_logs_json || COALESCE(step_mutation_logs, '[]'::jsonb); + END IF; + -- ============================================================================ -- STEP 1: Calculate amount_to_deduct if target_balance is provided -- ============================================================================ @@ -126,7 +176,7 @@ BEGIN -- Use new rollovers array (with credit_cost) if present, otherwise fall back to rollover_ids IF ((rollovers_arr IS NOT NULL AND jsonb_typeof(rollovers_arr) = 'array' AND jsonb_array_length(rollovers_arr) > 0) OR (rollover_ids IS NOT NULL AND array_length(rollover_ids, 1) > 0)) AND rollover_deducted = 0 THEN - SELECT * INTO rollover_deducted + SELECT * INTO rollover_deducted, step_mutation_logs FROM deduct_from_rollovers(jsonb_build_object( 'rollover_ids', rollover_ids, 'rollovers', rollovers_arr, @@ -134,6 +184,7 @@ BEGIN 'target_entity_id', target_entity_id, 'has_entity_scope', has_entity_scope )); + mutation_logs_json := mutation_logs_json || COALESCE(step_mutation_logs, '[]'::jsonb); -- rollover_deducted is returned in feature units, so can subtract directly remaining_amount := remaining_amount - rollover_deducted; END IF; @@ -170,8 +221,9 @@ BEGIN remaining_amount := remaining_amount - additional_deducted; -- STEP 3: Perform deduction from main balance (Pass 1: allow_negative = false) - SELECT * INTO deducted, new_balance, new_entities, new_adjustment + SELECT * INTO deducted, new_balance, new_entities, new_adjustment, step_mutation_logs FROM deduct_from_main_balance(jsonb_build_object( + 'customer_entitlement_id', ent_id, 'current_balance', current_balance, 'current_entities', current_entities, 'current_adjustment', new_adjustment, @@ -188,6 +240,7 @@ BEGIN -- STEP 4: Update database if any deduction occurred IF deducted != 0 OR additional_deducted != 0 THEN + mutation_logs_json := mutation_logs_json || COALESCE(step_mutation_logs, '[]'::jsonb); IF has_entity_scope THEN UPDATE customer_entitlements ce SET @@ -206,7 +259,7 @@ BEGIN WHERE ce.id = ent_id; END IF; - -- Track in updates_json (deducted is inclusive of additional_deducted) + -- Track in updates_json (merge with any unwind-side update already recorded) updates_json := jsonb_set( updates_json, ARRAY[ent_id], @@ -215,8 +268,8 @@ BEGIN 'additional_balance', new_additional_balance, 'adjustment', new_adjustment, 'entities', new_entities, - 'deducted', deducted + additional_deducted, - 'additional_deducted', additional_deducted + 'deducted', COALESCE((updates_json->ent_id->>'deducted')::numeric, 0) + deducted + additional_deducted, + 'additional_deducted', COALESCE((updates_json->ent_id->>'additional_deducted')::numeric, 0) + additional_deducted ) ); @@ -263,8 +316,9 @@ BEGIN new_additional_balance := current_additional_balance; -- Perform deduction (Pass 2: allow_negative = true) - SELECT * INTO deducted, new_balance, new_entities, new_adjustment + SELECT * INTO deducted, new_balance, new_entities, new_adjustment, step_mutation_logs FROM deduct_from_main_balance(jsonb_build_object( + 'customer_entitlement_id', ent_id, 'current_balance', current_balance, 'current_entities', current_entities, 'current_adjustment', current_adjustment, @@ -281,6 +335,7 @@ BEGIN -- Update database if deduction occurred (or addition with negative amount) IF deducted != 0 THEN + mutation_logs_json := mutation_logs_json || COALESCE(step_mutation_logs, '[]'::jsonb); IF has_entity_scope THEN UPDATE customer_entitlements ce SET @@ -345,7 +400,14 @@ BEGIN END IF; -- Read back updated rollovers (if any were involved in this deduction) + final_rollover_ids := COALESCE(unwind_modified_rollover_ids, ARRAY[]::text[]); + IF rollovers_arr IS NOT NULL AND jsonb_typeof(rollovers_arr) = 'array' AND jsonb_array_length(rollovers_arr) > 0 THEN + final_rollover_ids := array_cat( + final_rollover_ids, + ARRAY(SELECT (jsonb_array_elements(rollovers_arr))->>'id') + ); + SELECT COALESCE(jsonb_agg(jsonb_build_object( 'id', r.id, 'cus_ent_id', r.cus_ent_id, @@ -355,8 +417,10 @@ BEGIN )), '[]'::jsonb) INTO rollover_updates_json FROM rollovers r - WHERE r.id IN (SELECT (jsonb_array_elements(rollovers_arr))->>'id'); + WHERE r.id = ANY(final_rollover_ids); ELSIF rollover_ids IS NOT NULL AND array_length(rollover_ids, 1) > 0 THEN + final_rollover_ids := array_cat(final_rollover_ids, rollover_ids); + SELECT COALESCE(jsonb_agg(jsonb_build_object( 'id', r.id, 'cus_ent_id', r.cus_ent_id, @@ -366,17 +430,28 @@ BEGIN )), '[]'::jsonb) INTO rollover_updates_json FROM rollovers r - WHERE r.id = ANY(rollover_ids); + WHERE r.id = ANY(final_rollover_ids); + ELSIF array_length(final_rollover_ids, 1) > 0 THEN + SELECT COALESCE(jsonb_agg(jsonb_build_object( + 'id', r.id, + 'cus_ent_id', r.cus_ent_id, + 'balance', r.balance, + 'usage', r.usage, + 'entities', COALESCE(r.entities, '{}'::jsonb) + )), '[]'::jsonb) + INTO rollover_updates_json + FROM rollovers r + WHERE r.id = ANY(final_rollover_ids); END IF; -- Build final result result_json := jsonb_build_object( 'updates', updates_json, 'remaining', remaining_amount, - 'rollover_updates', rollover_updates_json + 'rollover_updates', rollover_updates_json, + 'mutation_logs', mutation_logs_json ); RETURN result_json; END; $$; - diff --git a/server/src/internal/balances/utils/sql/unwindFromLockReceipt.sql b/server/src/internal/balances/utils/sql/unwindFromLockReceipt.sql new file mode 100644 index 000000000..2bcf005c7 --- /dev/null +++ b/server/src/internal/balances/utils/sql/unwindFromLockReceipt.sql @@ -0,0 +1,208 @@ +DROP FUNCTION IF EXISTS unwind_from_lock_receipt(jsonb); + +CREATE FUNCTION unwind_from_lock_receipt(params jsonb) +RETURNS TABLE ( + remaining_unwind_value numeric, + updates jsonb, + modified_rollover_ids text[], + mutation_logs jsonb +) +LANGUAGE plpgsql +AS $$ +DECLARE + lock_receipt jsonb := params->'lock_receipt'; + receipt_items jsonb := COALESCE(lock_receipt->'items', '[]'::jsonb); + requested_unwind_value numeric := COALESCE((params->>'unwind_value')::numeric, 0); + + item_index integer; + item jsonb; + item_target_type text; + customer_entitlement_id text; + rollover_id text; + entity_id text; + credit_cost numeric; + + item_value_delta numeric; + item_value_magnitude numeric; + unwind_iteration_value numeric; + credits_to_unwind numeric; + + inverse_balance_delta numeric; + inverse_adjustment_delta numeric; + inverse_usage_delta numeric; + inverse_value_delta numeric; + + updated_balance numeric; + updated_additional_balance numeric; + updated_adjustment numeric; + updated_entities jsonb; + + remaining_value numeric := requested_unwind_value; + updates_json jsonb := '{}'::jsonb; + mutation_logs_json jsonb := '[]'::jsonb; + modified_rollover_ids_array text[] := ARRAY[]::text[]; +BEGIN + IF requested_unwind_value <= 0 THEN + RETURN QUERY SELECT 0::numeric, '{}'::jsonb, ARRAY[]::text[], '[]'::jsonb; + RETURN; + END IF; + + IF jsonb_typeof(receipt_items) != 'array' OR jsonb_array_length(receipt_items) = 0 THEN + RAISE EXCEPTION 'LOCK_RECEIPT_ITEMS_MISSING'; + END IF; + + FOR item_index IN REVERSE jsonb_array_length(receipt_items) - 1..0 + LOOP + EXIT WHEN remaining_value <= 0; + + item := receipt_items->item_index; + item_target_type := item->>'target_type'; + customer_entitlement_id := NULLIF(item->>'customer_entitlement_id', ''); + rollover_id := NULLIF(item->>'rollover_id', ''); + entity_id := NULLIF(item->>'entity_id', ''); + credit_cost := COALESCE((item->>'credit_cost')::numeric, 1); + + item_value_delta := COALESCE((item->>'value_delta')::numeric, 0); + item_value_magnitude := ABS(item_value_delta); + unwind_iteration_value := LEAST(item_value_magnitude, remaining_value); + + IF unwind_iteration_value <= 0 THEN + CONTINUE; + END IF; + + credits_to_unwind := unwind_iteration_value * credit_cost; + + inverse_balance_delta := CASE + WHEN COALESCE((item->>'balance_delta')::numeric, 0) > 0 THEN -credits_to_unwind + WHEN COALESCE((item->>'balance_delta')::numeric, 0) < 0 THEN credits_to_unwind + ELSE 0 + END; + inverse_adjustment_delta := CASE + WHEN COALESCE((item->>'adjustment_delta')::numeric, 0) > 0 THEN -credits_to_unwind + WHEN COALESCE((item->>'adjustment_delta')::numeric, 0) < 0 THEN credits_to_unwind + ELSE 0 + END; + inverse_usage_delta := CASE + WHEN COALESCE((item->>'usage_delta')::numeric, 0) > 0 THEN -credits_to_unwind + WHEN COALESCE((item->>'usage_delta')::numeric, 0) < 0 THEN credits_to_unwind + ELSE 0 + END; + inverse_value_delta := CASE + WHEN item_value_delta > 0 THEN -unwind_iteration_value + WHEN item_value_delta < 0 THEN unwind_iteration_value + ELSE 0 + END; + + IF item_target_type = 'customer_entitlement' THEN + IF customer_entitlement_id IS NULL THEN + RAISE EXCEPTION 'LOCK_CUSTOMER_ENTITLEMENT_ID_MISSING'; + END IF; + + IF entity_id IS NULL THEN + UPDATE customer_entitlements ce + SET + balance = ce.balance + inverse_balance_delta, + adjustment = COALESCE(ce.adjustment, 0) + inverse_adjustment_delta + WHERE ce.id = customer_entitlement_id + RETURNING + ce.balance, + COALESCE(ce.additional_balance, 0), + COALESCE(ce.adjustment, 0), + COALESCE(ce.entities, '{}'::jsonb) + INTO updated_balance, updated_additional_balance, updated_adjustment, updated_entities; + ELSE + UPDATE customer_entitlements ce + SET entities = jsonb_set( + jsonb_set( + COALESCE(ce.entities, '{}'::jsonb), + ARRAY[entity_id, 'balance'], + to_jsonb(COALESCE((COALESCE(ce.entities, '{}'::jsonb)->entity_id->>'balance')::numeric, 0) + inverse_balance_delta), + true + ), + ARRAY[entity_id, 'adjustment'], + to_jsonb(COALESCE((COALESCE(ce.entities, '{}'::jsonb)->entity_id->>'adjustment')::numeric, 0) + inverse_adjustment_delta), + true + ) + WHERE ce.id = customer_entitlement_id + RETURNING + ce.balance, + COALESCE(ce.additional_balance, 0), + COALESCE(ce.adjustment, 0), + COALESCE(ce.entities, '{}'::jsonb) + INTO updated_balance, updated_additional_balance, updated_adjustment, updated_entities; + END IF; + + updates_json := jsonb_set( + updates_json, + ARRAY[customer_entitlement_id], + jsonb_build_object( + 'balance', updated_balance, + 'additional_balance', updated_additional_balance, + 'adjustment', updated_adjustment, + 'entities', updated_entities, + 'deducted', COALESCE((updates_json->customer_entitlement_id->>'deducted')::numeric, 0) + inverse_value_delta, + 'additional_deducted', COALESCE((updates_json->customer_entitlement_id->>'additional_deducted')::numeric, 0) + ), + true + ); + ELSIF item_target_type = 'rollover' THEN + IF rollover_id IS NULL THEN + RAISE EXCEPTION 'LOCK_ROLLOVER_ID_MISSING'; + END IF; + + IF entity_id IS NULL THEN + UPDATE rollovers r + SET + balance = r.balance + inverse_balance_delta, + usage = COALESCE(r.usage, 0) + inverse_usage_delta + WHERE r.id = rollover_id; + ELSE + UPDATE rollovers r + SET entities = jsonb_set( + jsonb_set( + COALESCE(r.entities, '{}'::jsonb), + ARRAY[entity_id, 'balance'], + to_jsonb(COALESCE((COALESCE(r.entities, '{}'::jsonb)->entity_id->>'balance')::numeric, 0) + inverse_balance_delta), + true + ), + ARRAY[entity_id, 'usage'], + to_jsonb(COALESCE((COALESCE(r.entities, '{}'::jsonb)->entity_id->>'usage')::numeric, 0) + inverse_usage_delta), + true + ) + WHERE r.id = rollover_id; + END IF; + + modified_rollover_ids_array := array_append(modified_rollover_ids_array, rollover_id); + ELSE + RAISE EXCEPTION 'INVALID_LOCK_ITEM_TARGET_TYPE|targetType:%', item_target_type; + END IF; + + mutation_logs_json := mutation_logs_json || jsonb_build_array( + jsonb_build_object( + 'target_type', item_target_type, + 'customer_entitlement_id', customer_entitlement_id, + 'rollover_id', rollover_id, + 'entity_id', entity_id, + 'credit_cost', credit_cost, + 'balance_delta', inverse_balance_delta, + 'adjustment_delta', inverse_adjustment_delta, + 'usage_delta', inverse_usage_delta, + 'value_delta', inverse_value_delta + ) + ); + + remaining_value := remaining_value - unwind_iteration_value; + END LOOP; + + IF remaining_value > 0 THEN + RAISE EXCEPTION 'LOCK_UNWIND_INCOMPLETE|remaining:%', remaining_value; + END IF; + + RETURN QUERY + SELECT + remaining_value, + updates_json, + modified_rollover_ids_array, + mutation_logs_json; +END; +$$; diff --git a/server/src/internal/balances/utils/types/featureDeduction.ts b/server/src/internal/balances/utils/types/featureDeduction.ts index 4d09401ef..b440158d1 100644 --- a/server/src/internal/balances/utils/types/featureDeduction.ts +++ b/server/src/internal/balances/utils/types/featureDeduction.ts @@ -1,10 +1,12 @@ import type { Feature, LockParams } from "@autumn/shared"; +import type { LockReceipt } from "../lock/fetchLockReceipt.js"; export type FeatureDeduction = { feature: Feature; deduction: number; targetBalance?: number; lock?: LockParams; + lockReceipt?: LockReceipt; lockReceiptKey?: string; unwindValue?: number; }; diff --git a/server/tests/integration/balances/lock/check-with-lock-postgres.test.ts b/server/tests/integration/balances/lock/check-with-lock-postgres.test.ts index b9d2a7f99..08fab12a2 100644 --- a/server/tests/integration/balances/lock/check-with-lock-postgres.test.ts +++ b/server/tests/integration/balances/lock/check-with-lock-postgres.test.ts @@ -1,4 +1,5 @@ import { test } from "bun:test"; +import type { ApiCustomerV5 } from "@autumn/shared"; import { deleteLock } from "@tests/integration/balances/utils/lockUtils/deleteLock.js"; import { TestFeature } from "@tests/setup/v2Features.js"; import { items } from "@tests/utils/fixtures/items.js"; @@ -42,14 +43,19 @@ test.concurrent(`${chalk.yellowBright("check-with-lock-postgres: /check with loc skip_cache: true, }); - // // Release lock - // await autumnV2.balances.finalize({ - // finalize_action: "confirm", - // overwrite_value: 4, - // lock_key: lockKey, - // }); + // Release lock + await autumnV2.balances.finalize( + { + finalize_action: "confirm", + overwrite_value: 12, + lock_key: lockKey, + }, + { + skipCache: true, + }, + ); - // const customer = await autumnV2.customers.get(customerId); + const customer = await autumnV2.customers.get(customerId); - // console.log("Message balance:", customer.balances[TestFeature.Messages]); + console.log("Message balance:", customer.balances[TestFeature.Messages]); });