From bd66f1177a68e5a37c1bf730f98b0103de07c180 Mon Sep 17 00:00:00 2001 From: Joseph-tech10 <243628124+joseph-tech10@users.noreply.github.com> Date: Tue, 20 Oct 2026 16:01:04 +0100 Subject: [PATCH] fix(db): add global request ID registry for partitioned usage events --- docs/schema.md | 11 + .../0027_usage_event_request_ids.down.sql | 14 ++ migrations/0027_usage_event_request_ids.sql | 53 +++++ scripts/backfill-usage-partitions.ts | 76 +++++-- src/repositories/usageEventsRepository.pg.ts | 75 ++++--- src/routes/billing/deduct.ts | 4 + src/routes/billing/deduct/bulk.ts | 6 +- src/services/billing.ts | 116 +++++++--- src/services/usageStore.ts | 29 ++- tests/integration/usageEventRegistry.test.ts | 199 ++++++++++++++++++ 10 files changed, 497 insertions(+), 86 deletions(-) create mode 100644 migrations/0027_usage_event_request_ids.down.sql create mode 100644 migrations/0027_usage_event_request_ids.sql create mode 100644 tests/integration/usageEventRegistry.test.ts diff --git a/docs/schema.md b/docs/schema.md index d90a7c64..7543cae7 100644 --- a/docs/schema.md +++ b/docs/schema.md @@ -8,6 +8,17 @@ This document describes the core tables that power the Callora API marketplace, ## Tables +### `usage_events` and `usage_event_request_ids` + +`usage_events` is hash partitioned by `developer_id`; its per-partition primary +key includes that partition key, so it cannot enforce global request ID +uniqueness. Every writer reserves `request_id` in `usage_event_request_ids` +within the same transaction before inserting the event. The registry primary +key provides the global idempotency guarantee and its unique `usage_event_id` +is the target for `revenue_ledger.usage_event_id`. The partitioned parent also +contains `api_key VARCHAR(255)` and `status_code INTEGER NOT NULL DEFAULT 200` +for persistent usage-store reads and writes. + ### `apis` Stores the top-level API products created by developers. diff --git a/migrations/0027_usage_event_request_ids.down.sql b/migrations/0027_usage_event_request_ids.down.sql new file mode 100644 index 00000000..2ca98f6a --- /dev/null +++ b/migrations/0027_usage_event_request_ids.down.sql @@ -0,0 +1,14 @@ +BEGIN; +ALTER TABLE revenue_ledger DROP CONSTRAINT IF EXISTS revenue_ledger_usage_event_id_fkey; +DO $$ +BEGIN + IF to_regclass('usage_events_old') IS NOT NULL THEN + ALTER TABLE revenue_ledger + ADD CONSTRAINT revenue_ledger_usage_event_id_fkey + FOREIGN KEY (usage_event_id) REFERENCES usage_events_old(id); + END IF; +END$$; +DROP TABLE IF EXISTS usage_event_request_ids; +ALTER TABLE usage_events DROP COLUMN IF EXISTS status_code; +ALTER TABLE usage_events DROP COLUMN IF EXISTS api_key; +COMMIT; diff --git a/migrations/0027_usage_event_request_ids.sql b/migrations/0027_usage_event_request_ids.sql new file mode 100644 index 00000000..d0b73d23 --- /dev/null +++ b/migrations/0027_usage_event_request_ids.sql @@ -0,0 +1,53 @@ +-- Global idempotency registry for hash-partitioned usage events. PostgreSQL +-- cannot enforce request_id uniqueness on a partitioned table without the +-- partition key, so request IDs are reserved here in the same transaction. +-- destructive-approved: #1486 +BEGIN; + +CREATE TABLE IF NOT EXISTS usage_event_request_ids ( + request_id VARCHAR(255) PRIMARY KEY, + usage_event_id BIGINT NOT NULL UNIQUE, + developer_id VARCHAR(255) NOT NULL, + created_at TIMESTAMP NOT NULL DEFAULT NOW() +); + +-- Preserve the legacy persistent-store fields on the new partitioned parent. +ALTER TABLE usage_events + ADD COLUMN IF NOT EXISTS api_key VARCHAR(255), + ADD COLUMN IF NOT EXISTS status_code INTEGER NOT NULL DEFAULT 200; + +-- Carry existing IDs forward before attaching the ledger FK. The standalone +-- backfill script performs a duplicate audit and writes any collisions to a +-- report; this SQL migration remains safe to rerun on a partially copied set. +DO $$ +BEGIN + IF to_regclass('usage_events_old') IS NOT NULL THEN + INSERT INTO usage_event_request_ids (request_id, usage_event_id, developer_id) + SELECT o.request_id, o.id, COALESCE(NULLIF(o.developer_id, ''), a.developer_id::text, '') + FROM usage_events_old o LEFT JOIN apis a ON a.id::text = o.api_id + ON CONFLICT (request_id) DO NOTHING; + END IF; +END$$; + +INSERT INTO usage_event_request_ids (request_id, usage_event_id, developer_id) +SELECT ue.request_id, ue.id, ue.developer_id + FROM usage_events ue +ON CONFLICT (request_id) DO NOTHING; + +-- Existing IDs were allocated by the pre-partitioned sequence. Ensure the new +-- parent sequence starts above every ID reserved in the global registry. +SELECT setval( + pg_get_serial_sequence('usage_events', 'id'), + GREATEST(COALESCE(MAX(usage_event_id), 1), 1), + COALESCE(MAX(usage_event_id), 0) > 0 +) +FROM usage_event_request_ids; + +-- Replace the stale FK (which remained attached to usage_events_old after the +-- table swap) with a link to the global registry's unique event id. +ALTER TABLE revenue_ledger DROP CONSTRAINT IF EXISTS revenue_ledger_usage_event_id_fkey; +ALTER TABLE revenue_ledger + ADD CONSTRAINT revenue_ledger_usage_event_id_fkey + FOREIGN KEY (usage_event_id) REFERENCES usage_event_request_ids(usage_event_id); + +COMMIT; diff --git a/scripts/backfill-usage-partitions.ts b/scripts/backfill-usage-partitions.ts index f080e075..7f329bfb 100644 --- a/scripts/backfill-usage-partitions.ts +++ b/scripts/backfill-usage-partitions.ts @@ -3,8 +3,7 @@ * Backfill script: copy rows from usage_events_old into the new * hash-partitioned usage_events table. * - * Idempotent: uses ON CONFLICT (request_id, developer_id) DO NOTHING so - * re-running after a partial copy is safe. + * Idempotent: each batch reserves request IDs and copies events transactionally. * * Usage: * DATABASE_URL=postgres://... tsx scripts/backfill-usage-partitions.ts @@ -15,6 +14,7 @@ */ import pg from 'pg'; +import { writeFile } from 'node:fs/promises'; import { logger } from '../src/logger.js'; const { Pool } = pg; @@ -47,6 +47,13 @@ async function main(): Promise { ); process.exit(1); } + const { rows: registryCheck } = await pool.query<{ exists: boolean }>(` + SELECT to_regclass('usage_event_request_ids') IS NOT NULL AS exists + `); + if (!registryCheck[0]?.exists) { + logger.error('usage_event_request_ids not found. Run migrations/0027_usage_event_request_ids.sql first.'); + process.exit(1); + } const { rows: totalRows } = await pool.query<{ count: string }>( 'SELECT COUNT(*)::text AS count FROM usage_events_old', @@ -59,14 +66,50 @@ async function main(): Promise { return; } + const { rows: duplicates } = await pool.query<{ request_id: string; developers: string }>( + `WITH source AS ( + SELECT o.request_id, o.id, + COALESCE(NULLIF(o.developer_id, ''), a.developer_id::text, '') AS developer_id + FROM usage_events_old o LEFT JOIN apis a ON a.id::text = o.api_id + ), duplicate_ids AS ( + SELECT request_id, + string_agg(DISTINCT developer_id, ',' ORDER BY developer_id) AS developers + FROM source GROUP BY request_id HAVING COUNT(*) > 1 + ), registry_conflicts AS ( + SELECT s.request_id, + string_agg(DISTINCT s.developer_id || ' <> ' || r.developer_id, ',' ORDER BY s.developer_id || ' <> ' || r.developer_id) AS developers + FROM source s JOIN usage_event_request_ids r USING (request_id) + WHERE r.usage_event_id <> s.id OR r.developer_id <> s.developer_id + GROUP BY s.request_id + ) + SELECT request_id, developers FROM duplicate_ids + UNION + SELECT request_id, developers FROM registry_conflicts + ORDER BY request_id`, + ); + if (duplicates.length) { + const reportPath = process.env['DUPLICATE_REPORT'] ?? 'usage-event-request-id-duplicates.csv'; + await writeFile(reportPath, `request_id,developers\n${duplicates.map((row) => `"${row.request_id.replaceAll('"', '""')}","${row.developers.replaceAll('"', '""')}"`).join('\n')}\n`); + throw new Error(`Found ${duplicates.length} duplicate request_id values; report written to ${reportPath}`); + } + let offset = 0; let copied = 0; let skipped = 0; - while (true) { + while (offset < total) { // Fetch a batch ordered by id for deterministic cursor progress - const { rows: batch } = await pool.query<{ id: string }>( - `INSERT INTO usage_events ( + const client = await pool.connect(); + let batch: { id: string }[]; + try { + await client.query('BEGIN'); + await client.query( + `INSERT INTO usage_event_request_ids (request_id, usage_event_id, developer_id) + SELECT o.request_id, o.id, COALESCE(NULLIF(o.developer_id, ''), a.developer_id::text, '') + FROM (SELECT * FROM usage_events_old ORDER BY id LIMIT $1 OFFSET $2) o + LEFT JOIN apis a ON a.id::text = o.api_id + ON CONFLICT (request_id) DO NOTHING`, [BATCH_SIZE, offset]); + const inserted = await client.query<{ id: string }>(`INSERT INTO usage_events ( id, user_id, api_id, @@ -84,7 +127,7 @@ async function main(): Promise { o.api_id, o.endpoint_id, o.api_key_id, - COALESCE(o.developer_id, COALESCE(a.developer_id::text, '')), + COALESCE(NULLIF(o.developer_id, ''), COALESCE(a.developer_id::text, '')), o.amount_usdc, o.request_id, o.stellar_tx_hash, @@ -95,12 +138,13 @@ async function main(): Promise { LIMIT $1 OFFSET $2 ) o LEFT JOIN apis a ON a.id::text = o.api_id - ON CONFLICT (request_id, developer_id) DO NOTHING - RETURNING id`, - [BATCH_SIZE, offset], - ); - - if (batch.length === 0) break; + ON CONFLICT (id, developer_id) DO NOTHING RETURNING id`, [BATCH_SIZE, offset]); + batch = inserted.rows; + await client.query('COMMIT'); + } catch (error) { + await client.query('ROLLBACK'); + throw error; + } finally { client.release(); } const batchInserted = batch.length; // batch from the SELECT could be up to BATCH_SIZE; inserted may be less due to conflicts @@ -112,6 +156,14 @@ async function main(): Promise { skipped = total - copied; + await pool.query(` + SELECT setval( + pg_get_serial_sequence('usage_events', 'id'), + GREATEST(COALESCE(MAX(usage_event_id), 1), 1), + COALESCE(MAX(usage_event_id), 0) > 0 + ) FROM usage_event_request_ids + `); + logger.info(`Backfill complete: ${copied} inserted, ${skipped} skipped (already present).`); } finally { await pool.end(); diff --git a/src/repositories/usageEventsRepository.pg.ts b/src/repositories/usageEventsRepository.pg.ts index 982c44ee..9abd95ad 100644 --- a/src/repositories/usageEventsRepository.pg.ts +++ b/src/repositories/usageEventsRepository.pg.ts @@ -249,55 +249,60 @@ export class PgUsageEventsRepository implements UsageEventsPgRepository { const apiId = assertNonEmpty(event.apiId, 'apiId'); const endpointId = assertNonEmpty(event.endpointId, 'endpointId'); const apiKeyId = assertNonEmpty(event.apiKeyId, 'apiKeyId'); - const developerId = (event.developerId ?? '').trim(); const requestId = assertNonEmpty(event.requestId, 'requestId'); const amount = assertAmount(event.amount).toString(); const result = await this.writeDb.query( ` - INSERT INTO usage_events ( - user_id, - api_id, - endpoint_id, - api_key_id, - developer_id, - amount_usdc, - request_id, - stellar_tx_hash, - created_at + WITH owner AS ( + SELECT developer_id::text AS developer_id FROM apis WHERE id::text = $2 LIMIT 1 + ), reservation AS ( + INSERT INTO usage_event_request_ids (request_id, usage_event_id, developer_id) + SELECT $6, nextval(pg_get_serial_sequence('usage_events', 'id')), developer_id FROM owner + ON CONFLICT (request_id) DO NOTHING + RETURNING usage_event_id, developer_id + ), inserted AS ( + INSERT INTO usage_events ( + id, user_id, api_id, endpoint_id, api_key_id, developer_id, + amount_usdc, request_id, stellar_tx_hash, created_at + ) + SELECT r.usage_event_id, $1, $2, $3, $4, r.developer_id, + $5, $6, $7, COALESCE($8, NOW()) + FROM reservation r + RETURNING id, user_id, api_id, endpoint_id, api_key_id, developer_id, + amount_usdc, request_id, stellar_tx_hash, created_at ) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, COALESCE($9, NOW())) - ON CONFLICT (request_id, developer_id) - DO UPDATE SET request_id = EXCLUDED.request_id - RETURNING - id, - user_id, - api_id, - endpoint_id, - api_key_id, - developer_id, - amount_usdc, - request_id, - stellar_tx_hash, - created_at + SELECT * FROM inserted + UNION ALL + SELECT ue.id, ue.user_id, ue.api_id, ue.endpoint_id, ue.api_key_id, + ue.developer_id, ue.amount_usdc, ue.request_id, ue.stellar_tx_hash, ue.created_at + FROM usage_event_request_ids r + JOIN usage_events ue ON ue.id = r.usage_event_id + WHERE r.request_id = $6 AND NOT EXISTS (SELECT 1 FROM inserted) `, [ - userId, - apiId, - endpointId, - apiKeyId, - developerId, - amount, - requestId, - event.stellarTxHash ?? null, - event.createdAt ?? null, + userId, apiId, endpointId, apiKeyId, amount, requestId, + event.stellarTxHash ?? null, event.createdAt ?? null, ], ); const row = result.rows[0]; if (!row) { - throw new Error(`Failed to create or retrieve usage event for requestId "${requestId}".`); + const api = await this.writeDb.query('SELECT 1 FROM apis WHERE id::text = $1 LIMIT 1', [apiId]); + if (!api.rows[0]) { + throw Object.assign(new Error(`Unknown API: ${apiId}`), { statusCode: 404, code: 'API_NOT_FOUND' }); + } + // A concurrent writer may have committed the reservation after this + // statement's snapshot was taken. A fresh statement sees its complete row. + const retry = await this.writeDb.query( + `SELECT ue.id, ue.user_id, ue.api_id, ue.endpoint_id, ue.api_key_id, + ue.developer_id, ue.amount_usdc, ue.request_id, ue.stellar_tx_hash, ue.created_at + FROM usage_event_request_ids r JOIN usage_events ue ON ue.id = r.usage_event_id + WHERE r.request_id = $1`, [requestId], + ); + if (!retry.rows[0]) throw new Error(`Reserved request id has no usage event: ${requestId}`); + return mapUsageEventRow(retry.rows[0]); } return mapUsageEventRow(row); diff --git a/src/routes/billing/deduct.ts b/src/routes/billing/deduct.ts index c6a5cb8f..e29a87e7 100644 --- a/src/routes/billing/deduct.ts +++ b/src/routes/billing/deduct.ts @@ -141,6 +141,10 @@ router.post( }); if (!result.success) { + if (result.errorCode === "API_NOT_FOUND") { + next(new NotFoundError(result.error ?? "API not found")); + return; + } if (result.simulationDetails) { next( simulationFailureError( diff --git a/src/routes/billing/deduct/bulk.ts b/src/routes/billing/deduct/bulk.ts index 2c162b07..eb49a8ae 100644 --- a/src/routes/billing/deduct/bulk.ts +++ b/src/routes/billing/deduct/bulk.ts @@ -1,6 +1,6 @@ import { Router, type Response, type NextFunction, type Request } from 'express'; import { z } from 'zod'; -import { UnauthorizedError, PaymentRequiredError, InternalServerError } from '../../../errors/index.js'; +import { UnauthorizedError, PaymentRequiredError, InternalServerError, NotFoundError } from '../../../errors/index.js'; import { requireAuth, type AuthenticatedLocals } from '../../../middleware/requireAuth.js'; import { validate } from '../../../middleware/validate.js'; import { idempotencyMiddleware } from '../../../middleware/idempotency.js'; @@ -78,6 +78,10 @@ router.post( ); if (!result.success) { + if (result.errorCode === 'API_NOT_FOUND') { + next(new NotFoundError(result.error ?? 'API not found')); + return; + } if (result.error?.toLowerCase().includes('insufficient balance')) { next(new PaymentRequiredError(result.error, 'INSUFFICIENT_BALANCE')); return; diff --git a/src/services/billing.ts b/src/services/billing.ts index b14c1717..17d9b119 100644 --- a/src/services/billing.ts +++ b/src/services/billing.ts @@ -72,6 +72,7 @@ export interface BillingDeductResult { reconciliationRequired: boolean; deductionStatus?: DeductionStatus; error?: string; + errorCode?: string; simulationDetails?: SimulationDetails; } @@ -93,6 +94,7 @@ export interface BillingBulkDeductResult { totalDeductedAmountUsdc: string; stellarTxHash?: string; error?: string; + errorCode?: string; simulationDetails?: SimulationDetails; } @@ -327,17 +329,37 @@ async function runPhase1( _amountInContractUnits: bigint, ): Promise { await client.query("BEGIN"); + // Resolve the partition key before reserving the globally unique request id. + const owner = await client.query<{ developer_id: string }>( + `SELECT developer_id::text AS developer_id FROM apis WHERE id::text = $1 LIMIT 1`, + [request.apiId], + ); + if (!owner.rows[0]) { + const error = new Error(`Unknown API: ${request.apiId}`) as Error & { statusCode: number; code: string }; + error.statusCode = 404; + error.code = "API_NOT_FOUND"; + throw error; + } + const developerId = owner.rows[0].developer_id; + + // Reserve the request id first. The parent table cannot provide global + // uniqueness because it is partitioned by developer_id. + const reserved = await client.query<{ usage_event_id: string }>( + `INSERT INTO usage_event_request_ids (request_id, usage_event_id, developer_id) + VALUES ($1, nextval(pg_get_serial_sequence('usage_events', 'id')), $2) + ON CONFLICT (request_id) DO NOTHING RETURNING usage_event_id`, + [request.requestId, developerId], + ); - // Lock the row (if it exists) so concurrent requests with the same - // request_id serialise here rather than racing to INSERT. + if (!reserved.rows[0]) { const existingEvent = await client.query<{ id: string; stellar_tx_hash: string | null; deduction_status: string | null; }>( - `SELECT id, stellar_tx_hash, deduction_status - FROM usage_events - WHERE request_id = $1 + `SELECT ue.id, ue.stellar_tx_hash, ue.deduction_status + FROM usage_event_request_ids r JOIN usage_events ue ON ue.id = r.usage_event_id + WHERE r.request_id = $1 FOR UPDATE`, [request.requestId], ); @@ -377,6 +399,11 @@ async function runPhase1( deductionStatus: status, }; } + // A registry row without its usage row indicates a broken invariant. Never + // charge again or silently manufacture a replacement. + await client.query("ROLLBACK"); + throw new Error(`Request id ${request.requestId} is reserved without a usage event`); + } // Balance check — still inside the transaction so the connection is held, // but the Soroban call itself is NOT inside the transaction (see Phase 2). @@ -389,14 +416,16 @@ async function runPhase1( // Committed immediately so the record survives even if Phase 2 fails. const insertResult = await client.query<{ id: string }>( `INSERT INTO usage_events - (user_id, api_id, endpoint_id, api_key_id, amount_usdc, request_id, created_at, deduction_status) - VALUES ($1, $2, $3, $4, $5, $6, NOW(), 'pending') + (id, user_id, api_id, endpoint_id, api_key_id, developer_id, amount_usdc, request_id, created_at, deduction_status) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, NOW(), 'pending') RETURNING id`, [ + reserved.rows[0].usage_event_id, request.userId, request.apiId, request.endpointId, request.apiKeyId, + developerId, request.amountUsdc, request.requestId, ], @@ -460,9 +489,9 @@ async function getUsageEventsByRequestIds( } const result = await pool.query( - `SELECT request_id, id, stellar_tx_hash, deduction_status - FROM usage_events - WHERE request_id = ANY($1::text[])`, + `SELECT r.request_id, ue.id, ue.stellar_tx_hash, ue.deduction_status + FROM usage_event_request_ids r JOIN usage_events ue ON ue.id = r.usage_event_id + WHERE r.request_id = ANY($1::text[])`, [requestIds], ); @@ -477,25 +506,32 @@ async function runPhase1Bulk( ): Promise { await client.query("BEGIN"); - const requestIds = requests.map((request) => request.requestId); - const existingEvents = await client.query( - `SELECT request_id, id, stellar_tx_hash, deduction_status - FROM usage_events - WHERE request_id = ANY($1::text[]) - FOR UPDATE`, - [requestIds], - ); - - const rowsByRequestId = new Map( - existingEvents.rows.map((row: ExistingUsageEventRow) => [ - row.request_id, - row, - ]), - ); + const rowsByRequestId = new Map(); const inserted: BulkInsertedUsageEvent[] = []; for (const request of requests) { - const existingRow = rowsByRequestId.get(request.requestId); + const owner = await client.query<{ developer_id: string }>( + `SELECT developer_id::text AS developer_id FROM apis WHERE id::text = $1 LIMIT 1`, [request.apiId], + ); + if (!owner.rows[0]) { + const error = new Error(`Unknown API: ${request.apiId}`) as Error & { statusCode: number; code: string }; + error.statusCode = 404; + error.code = "API_NOT_FOUND"; + throw error; + } + const developerId = owner.rows[0].developer_id; + const reservation = await client.query<{ usage_event_id: string }>( + `INSERT INTO usage_event_request_ids (request_id, usage_event_id, developer_id) + VALUES ($1, nextval(pg_get_serial_sequence('usage_events', 'id')), $2) + ON CONFLICT (request_id) DO NOTHING RETURNING usage_event_id`, + [request.requestId, developerId], + ); + const existingResult = reservation.rows[0] ? { rows: [] as ExistingUsageEventRow[] } : await client.query( + `SELECT ue.request_id, ue.id, ue.stellar_tx_hash, ue.deduction_status + FROM usage_event_request_ids r JOIN usage_events ue ON ue.id = r.usage_event_id + WHERE r.request_id = $1 FOR UPDATE`, [request.requestId], + ); + const existingRow = rowsByRequestId.get(request.requestId) ?? existingResult.rows[0]; if (existingRow) { const status = (existingRow.deduction_status as DeductionStatus) ?? @@ -525,14 +561,16 @@ async function runPhase1Bulk( const insertResult = await client.query<{ id: string }>( `INSERT INTO usage_events - (user_id, api_id, endpoint_id, api_key_id, amount_usdc, request_id, created_at, deduction_status) - VALUES ($1, $2, $3, $4, $5, $6, NOW(), 'pending') + (id, user_id, api_id, endpoint_id, api_key_id, developer_id, amount_usdc, request_id, created_at, deduction_status) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, NOW(), 'pending') RETURNING id`, [ + reservation.rows[0].usage_event_id, request.userId, request.apiId, request.endpointId, request.apiKeyId, + developerId, request.amountUsdc, request.requestId, ], @@ -635,6 +673,9 @@ export class BillingService { deductionApplied: false, reconciliationRequired: false, error: normalizeErrorMessage(error), + errorCode: error && typeof error === 'object' && 'code' in error + ? String((error as { code?: unknown }).code) + : undefined, }; } @@ -702,7 +743,9 @@ export class BillingService { stellar_tx_hash: string | null; deduction_status: string | null; }>( - `SELECT id, stellar_tx_hash, deduction_status FROM usage_events WHERE request_id = $1`, + `SELECT ue.id, ue.stellar_tx_hash, ue.deduction_status + FROM usage_event_request_ids r JOIN usage_events ue ON ue.id = r.usage_event_id + WHERE r.request_id = $1`, [request.requestId], ); if (existing.rows.length > 0) { @@ -729,6 +772,9 @@ export class BillingService { deductionApplied: false, reconciliationRequired: false, error: normalizeErrorMessage(error), + errorCode: error && typeof error === 'object' && 'code' in error + ? String((error as { code?: unknown }).code) + : undefined, }; } finally { client.release(); @@ -864,6 +910,9 @@ export class BillingService { deductedCount: 0, totalDeductedAmountUsdc: "0", error: normalizeErrorMessage(error), + errorCode: error && typeof error === 'object' && 'code' in error + ? String((error as { code?: unknown }).code) + : undefined, }; } @@ -964,6 +1013,9 @@ export class BillingService { deductedCount: 0, totalDeductedAmountUsdc: "0", error: normalizeErrorMessage(error), + errorCode: error && typeof error === 'object' && 'code' in error + ? String((error as { code?: unknown }).code) + : undefined, }; } finally { client.release(); @@ -1123,9 +1175,9 @@ export class BillingService { stellar_tx_hash: string | null; deduction_status: string | null; }>( - `SELECT id, stellar_tx_hash, deduction_status - FROM usage_events - WHERE request_id = $1`, + `SELECT ue.id, ue.stellar_tx_hash, ue.deduction_status + FROM usage_event_request_ids r JOIN usage_events ue ON ue.id = r.usage_event_id + WHERE r.request_id = $1`, [requestId], ); diff --git a/src/services/usageStore.ts b/src/services/usageStore.ts index 008fb7ea..31210fa7 100644 --- a/src/services/usageStore.ts +++ b/src/services/usageStore.ts @@ -232,6 +232,24 @@ export class PostgresUsageStore implements UsageAdminStore { try { await client.query('BEGIN'); + const owner = await client.query<{ developer_id: string }>( + `SELECT developer_id::text AS developer_id FROM apis WHERE id::text = $1 LIMIT 1`, + [event.apiId], + ); + if (!owner.rows[0]) { + throw Object.assign(new Error(`Unknown API: ${event.apiId}`), { statusCode: 404, code: 'API_NOT_FOUND' }); + } + const reservation = await client.query<{ usage_event_id: string }>( + `INSERT INTO usage_event_request_ids (request_id, usage_event_id, developer_id) + VALUES ($1, nextval(pg_get_serial_sequence('usage_events', 'id')), $2) + ON CONFLICT (request_id) DO NOTHING RETURNING usage_event_id`, + [event.requestId, owner.rows[0].developer_id], + ); + if (!reservation.rows[0]) { + await client.query('ROLLBACK'); + return false; + } + const insertResult = await client.query<{ id: string | number }>( ` INSERT INTO usage_events ( @@ -242,14 +260,11 @@ export class PostgresUsageStore implements UsageAdminStore { developer_id, amount_usdc, request_id, + api_key, status_code, created_at ) - VALUES ($1, $2, $3, $4, ( - SELECT COALESCE(a.developer_id::text, '') - FROM apis a WHERE a.id = $2 LIMIT 1 - ), $5, $6, $7, $8) - ON CONFLICT (request_id, developer_id) DO NOTHING + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10) RETURNING id `, [ @@ -257,8 +272,10 @@ export class PostgresUsageStore implements UsageAdminStore { event.apiId, event.endpointId, event.apiKeyId, + owner.rows[0].developer_id, event.amountUsdc, event.requestId, + event.apiKey, event.statusCode, event.timestamp, ], @@ -309,7 +326,7 @@ export class PostgresUsageStore implements UsageAdminStore { async hasEvent(requestId: string): Promise { const result = await this.db.query<{ exists: number }>( - 'SELECT 1 AS exists FROM usage_events WHERE request_id = $1 LIMIT 1', + 'SELECT 1 AS exists FROM usage_event_request_ids WHERE request_id = $1 LIMIT 1', [requestId], ); return Boolean(result.rows[0]); diff --git a/tests/integration/usageEventRegistry.test.ts b/tests/integration/usageEventRegistry.test.ts new file mode 100644 index 00000000..0fcd2df8 --- /dev/null +++ b/tests/integration/usageEventRegistry.test.ts @@ -0,0 +1,199 @@ +import assert from 'node:assert/strict'; +import { mkdtemp, readFile, rm } from 'node:fs/promises'; +import { spawnSync } from 'node:child_process'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; +import { GenericContainer, StartedTestContainer, Wait } from 'testcontainers'; +import pg from 'pg'; +import { PgUsageEventsRepository } from '../../src/repositories/usageEventsRepository.pg.js'; +import { BillingService } from '../../src/services/billing.js'; +import { PostgresUsageStore } from '../../src/services/usageStore.js'; + +const { Pool } = pg; + +describe('partitioned usage event global request ID registry', () => { + let container: StartedTestContainer; + let pool: InstanceType; + let repository: PgUsageEventsRepository; + let databaseUrl: string; + + beforeAll(async () => { + container = await new GenericContainer('postgres:16-alpine') + .withExposedPorts(5432) + .withEnvironment({ POSTGRES_USER: 'postgres', POSTGRES_PASSWORD: 'postgres', POSTGRES_DB: 'callora_test' }) + .withWaitStrategy(Wait.forLogMessage('database system is ready to accept connections', 2)) + .start(); + databaseUrl = `postgresql://postgres:postgres@${container.getHost()}:${container.getMappedPort(5432)}/callora_test`; + pool = new Pool({ connectionString: databaseUrl }); + + await pool.query(` + CREATE TABLE apis (id VARCHAR(255) PRIMARY KEY, developer_id VARCHAR(255) NOT NULL); + INSERT INTO apis VALUES ('api-a', 'developer-a'), ('api-b', 'developer-b'); + CREATE TABLE usage_events_old ( + id BIGSERIAL PRIMARY KEY, + user_id VARCHAR(255) NOT NULL, + api_id VARCHAR(255) NOT NULL, + endpoint_id VARCHAR(255) NOT NULL, + api_key_id VARCHAR(255) NOT NULL, + developer_id VARCHAR(255) NOT NULL DEFAULT '', + amount_usdc NUMERIC(20, 0) NOT NULL, + request_id VARCHAR(255) NOT NULL, + stellar_tx_hash VARCHAR(64), + created_at TIMESTAMP NOT NULL DEFAULT NOW() + ); + INSERT INTO usage_events_old (user_id, api_id, endpoint_id, api_key_id, developer_id, amount_usdc, request_id) + VALUES ('user-old', 'api-a', 'endpoint-old', 'key-old', 'developer-a', 7, 'legacy-request'); + CREATE TABLE usage_events_partitioned ( + id BIGSERIAL, + user_id VARCHAR(255) NOT NULL, + api_id VARCHAR(255) NOT NULL, + endpoint_id VARCHAR(255) NOT NULL, + api_key_id VARCHAR(255) NOT NULL, + developer_id VARCHAR(255) NOT NULL DEFAULT '', + amount_usdc NUMERIC(20, 0) NOT NULL, + request_id VARCHAR(255) NOT NULL, + stellar_tx_hash VARCHAR(64), + created_at TIMESTAMP NOT NULL DEFAULT NOW(), + PRIMARY KEY (id, developer_id), + UNIQUE (request_id, developer_id) + ) PARTITION BY HASH (developer_id); + CREATE TABLE usage_events_p0 PARTITION OF usage_events_partitioned FOR VALUES WITH (MODULUS 2, REMAINDER 0); + CREATE TABLE usage_events_p1 PARTITION OF usage_events_partitioned FOR VALUES WITH (MODULUS 2, REMAINDER 1); + ALTER TABLE usage_events_partitioned RENAME TO usage_events; + CREATE TABLE revenue_ledger ( + id BIGSERIAL PRIMARY KEY, + api_id VARCHAR(255) NOT NULL DEFAULT '', + developer_id VARCHAR(255) NOT NULL DEFAULT '', + amount_usdc NUMERIC(20, 7) NOT NULL DEFAULT 0, + usage_event_id BIGINT UNIQUE REFERENCES usage_events_old(id), + settlement_id BIGINT, + created_at TIMESTAMP NOT NULL DEFAULT NOW() + ); + CREATE TABLE settlements (id BIGSERIAL PRIMARY KEY, external_id VARCHAR(255)); + `); + + const migration = await readFile(join(process.cwd(), 'migrations/0027_usage_event_request_ids.sql'), 'utf8'); + await pool.query(migration); + await pool.query(migration); + repository = new PgUsageEventsRepository(pool); + }, 60000); + + afterAll(async () => { + if (pool) await pool.end(); + if (container) await container.stop(); + }); + + test('migration backfills legacy IDs, restores store columns, and repoints the ledger FK', async () => { + const migrated = await pool.query<{ request_id: string; developer_id: string }>( + 'SELECT request_id, developer_id FROM usage_event_request_ids WHERE request_id = $1', ['legacy-request'], + ); + assert.deepEqual(migrated.rows, [{ request_id: 'legacy-request', developer_id: 'developer-a' }]); + + const columns = await pool.query<{ column_name: string }>( + `SELECT column_name FROM information_schema.columns + WHERE table_name = 'usage_events' AND column_name IN ('api_key', 'status_code') ORDER BY column_name`, + ); + assert.deepEqual(columns.rows.map((row) => row.column_name), ['api_key', 'status_code']); + + const legacyId = await pool.query<{ usage_event_id: string }>( + 'SELECT usage_event_id FROM usage_event_request_ids WHERE request_id = $1', ['legacy-request'], + ); + await pool.query('INSERT INTO revenue_ledger (usage_event_id) VALUES ($1)', [legacyId.rows[0]?.usage_event_id]); + await assert.rejects( + pool.query('INSERT INTO revenue_ledger (usage_event_id) VALUES (987654321)'), + /foreign key constraint/i, + ); + }); + + test('same request ID across different API owners commits exactly one event', async () => { + const args = (apiId: string, userId: string) => ({ + apiId, userId, endpointId: 'endpoint', apiKeyId: 'key', amount: 100n, requestId: 'concurrent-global-request', + }); + const [first, second] = await Promise.all([ + repository.create(args('api-a', 'user-a')), + repository.create(args('api-b', 'user-b')), + ]); + assert.equal(first.id, second.id); + const stored = await pool.query<{ count: string; developer_id: string }>( + `SELECT COUNT(*)::text AS count, MIN(developer_id) AS developer_id + FROM usage_events WHERE request_id = $1`, ['concurrent-global-request'], + ); + assert.equal(stored.rows[0]?.count, '1'); + assert.ok(['developer-a', 'developer-b'].includes(stored.rows[0]?.developer_id ?? '')); + const registry = await pool.query<{ count: string }>( + 'SELECT COUNT(*)::text AS count FROM usage_event_request_ids WHERE request_id = $1', ['concurrent-global-request'], + ); + assert.equal(registry.rows[0]?.count, '1'); + const retry = await repository.create(args('api-a', 'retry-user')); + assert.equal(retry.id, first.id); + assert.equal(retry.developerId, first.developerId); + }); + + test('unknown API is a typed 404 and writes no registry or usage row', async () => { + await assert.rejects( + repository.create({ + apiId: 'missing-api', userId: 'user', endpointId: 'endpoint', apiKeyId: 'key', + amount: 1n, requestId: 'unknown-api-request', + }), + (error: unknown) => error instanceof Error && (error as Error & { statusCode?: number; code?: string }).statusCode === 404 + && (error as Error & { code?: string }).code === 'API_NOT_FOUND', + ); + const billing = new BillingService(pool, { + async getBalance() { return { balance: '1000000000' }; }, + async deductBalance() { return { txHash: 'unused' }; }, + }); + const request = { + requestId: 'unknown-billing-api', userId: 'user', apiId: 'missing-api', + endpointId: 'endpoint', apiKeyId: 'key', amountUsdc: '1', + }; + const single = await billing.deduct(request); + assert.equal(single.errorCode, 'API_NOT_FOUND'); + assert.equal(single.success, false); + const bulk = await billing.deductBulk([request]); + assert.equal(bulk.errorCode, 'API_NOT_FOUND'); + assert.equal(bulk.success, false); + const result = await pool.query( + `SELECT 1 FROM usage_event_request_ids WHERE request_id IN ('unknown-api-request', 'unknown-billing-api') + UNION ALL SELECT 1 FROM usage_events WHERE request_id IN ('unknown-api-request', 'unknown-billing-api')`, + ); + assert.equal(result.rowCount, 0); + }); + + test('PostgresUsageStore persists and reads api_key and status_code on the partitioned table', async () => { + const store = new PostgresUsageStore(pool); + const didRecord = await store.record({ + id: '', requestId: 'store-roundtrip', apiKey: 'stored-api-key', apiKeyId: 'key', + apiId: 'api-a', endpointId: 'endpoint', userId: 'caller', amountUsdc: 30, + statusCode: 429, timestamp: new Date().toISOString(), + }); + assert.equal(didRecord, true); + const events = await store.getEvents('stored-api-key'); + assert.equal(events.length, 1); + assert.equal(events[0]?.statusCode, 429); + assert.equal(events[0]?.apiKey, 'stored-api-key'); + }); + + test('backfill rejects duplicate request IDs and writes a report', async () => { + await pool.query(` + INSERT INTO usage_events_old (user_id, api_id, endpoint_id, api_key_id, developer_id, amount_usdc, request_id) + VALUES ('user-b', 'api-b', 'endpoint-b', 'key-b', 'developer-b', 9, 'legacy-request') + `); + const reportDirectory = await mkdtemp(join(tmpdir(), 'usage-event-duplicates-')); + const reportPath = join(reportDirectory, 'duplicates.csv'); + try { + const script = spawnSync(process.execPath, ['--import', 'tsx', 'scripts/backfill-usage-partitions.ts'], { + cwd: process.cwd(), + encoding: 'utf8', + timeout: 60000, + env: { ...process.env, DATABASE_URL: databaseUrl, DUPLICATE_REPORT: reportPath }, + }); + assert.notEqual(script.status, 0, script.stdout + script.stderr); + const report = await readFile(reportPath, 'utf8'); + assert.match(report, /legacy-request/); + assert.match(report, /developer-a/); + assert.match(report, /developer-b/); + } finally { + await rm(reportDirectory, { recursive: true, force: true }); + } + }); +});