comi-logic/cfw-autumn
8
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

finished postgres functions

Browse Source
This commit is contained in:
John Yeo
2026-03-09 11:15:29 +00:00
parent 6de5947e0a
commit b506847ded
13 changed files with 560 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
server/src/db/initializeDatabaseFunctions.ts
View File

@@ -23,6 +23,7 @@ export const initializeDatabaseFunctions = async () => {
// Helper functions
"deductFromRollovers.sql",
"deductFromMainBalance.sql",
"unwindFromLockReceipt.sql",
"getTotalBalance.sql",
"deductFromAdditionalBalance.sql",
"performDeduction.sql",

16
server/src/external/autumn/autumnCli.ts vendored
View File

@@ -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<string, string>;
} = {},
) => {
const data = await this.post(`/balances.finalize`, params, {
...(skipCache ? { "x-skip-cache": "true" } : {}),
...headers,
});
return data;
},
};

2
server/src/honoMiddlewares/baseMiddleware.ts
View File

@@ -65,7 +65,7 @@ export const baseMiddleware = async (c: Context<HonoEnv>, next: Next) => {
// Query params
expand: [],
skipCache: false,
skipCache: c.req.header("x-skip-cache") === "true",
// Test params:
extraLogs: {},

2
server/src/honoMiddlewares/expandMiddleware.ts
View File

@@ -64,7 +64,7 @@ export const expandMiddleware = (): MiddlewareHandler<HonoEnv> => {
c.set("ctx", {
...ctx,
expand,
skipCache: skipCacheValue,
skipCache: ctx.skipCache || skipCacheValue,
});
await next();

56
server/src/internal/balances/finalizeLock/finalizeLock.ts
View File

@@ -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,

45
server/src/internal/balances/utils/deduction/executePostgresDeduction.ts
View File

@@ -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<string, DeductionUpdate>
> => {
const executeDeduction = async (): Promise<{
updates: Record<string, DeductionUpdate>;
mutationLogs: MutationLogItem[];
}> => {
let allUpdates: Record<string, DeductionUpdate> = {};
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<string, DeductionUpdate>;
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,
};
};

50
server/src/internal/balances/utils/lock/saveLockReceipt.ts Normal file
View File

@@ -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}`,
});
};

58
server/src/internal/balances/utils/sql/deductFromMainBalance.sql
View File

@@ -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;
$$;

53
server/src/internal/balances/utils/sql/deductFromRollovers.sql
View File

@@ -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;
$$;

95
server/src/internal/balances/utils/sql/performDeduction.sql
View File

@@ -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;
$$;

208
server/src/internal/balances/utils/sql/unwindFromLockReceipt.sql Normal file
View File

@@ -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;
$$;

2
server/src/internal/balances/utils/types/featureDeduction.ts
View File

@@ -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;
};

22
server/tests/integration/balances/lock/check-with-lock-postgres.test.ts
View File

@@ -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<ApiCustomerV5>(customerId);
const customer = await autumnV2.customers.get<ApiCustomerV5>(customerId);
// console.log("Message balance:", customer.balances[TestFeature.Messages]);
console.log("Message balance:", customer.balances[TestFeature.Messages]);
});