From 69acd664d7f7457e4c9685bcb466da49ae019314 Mon Sep 17 00:00:00 2001 From: diegosouzapw Date: Sat, 30 May 2026 22:32:02 -0300 Subject: [PATCH] =?UTF-8?q?feat(quota):=20pool=20can=20span=20N=20connecti?= =?UTF-8?q?ons=20=E2=80=94=20quota=5Fpool=5Fconnections=20join=20table=20(?= =?UTF-8?q?Phase=20D1)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds migration 086 to create the `quota_pool_connections` join table with a backfill that seeds every existing pool's single connection_id as its first member. Updates QuotaPool type with `connectionIds: string[]`, wires createPool/updatePool/deletePool to maintain the join table transactionally, and keeps `connection_id` as the primary back-compat column synced to `connectionIds[0]`. --- .../migrations/086_quota_pool_connections.sql | 21 ++ src/lib/db/quotaPools.ts | 147 ++++++++++-- tests/unit/quota-pool-connections.test.ts | 211 ++++++++++++++++++ 3 files changed, 361 insertions(+), 18 deletions(-) create mode 100644 src/lib/db/migrations/086_quota_pool_connections.sql create mode 100644 tests/unit/quota-pool-connections.test.ts diff --git a/src/lib/db/migrations/086_quota_pool_connections.sql b/src/lib/db/migrations/086_quota_pool_connections.sql new file mode 100644 index 0000000000..e823ff68db --- /dev/null +++ b/src/lib/db/migrations/086_quota_pool_connections.sql @@ -0,0 +1,21 @@ +-- 086: Multi-provider quota pools — quota_pool_connections join table (Phase D1). +-- +-- A quota pool can now span N provider connections (multi-provider). The legacy +-- quota_pools.connection_id column stays as the "primary" connection for +-- backwards compatibility. This table is the authoritative membership list. +-- Idempotent via CREATE TABLE/INDEX IF NOT EXISTS + INSERT OR IGNORE. + +CREATE TABLE IF NOT EXISTS quota_pool_connections ( + pool_id TEXT NOT NULL, + connection_id TEXT NOT NULL, + created_at TEXT NOT NULL DEFAULT (datetime('now')), + PRIMARY KEY (pool_id, connection_id) +); + +CREATE INDEX IF NOT EXISTS idx_quota_pool_connections_pool + ON quota_pool_connections(pool_id); + +-- Backfill: every existing pool's single connection_id becomes its first member. +INSERT OR IGNORE INTO quota_pool_connections (pool_id, connection_id) + SELECT id, connection_id FROM quota_pools + WHERE connection_id IS NOT NULL AND connection_id <> ''; diff --git a/src/lib/db/quotaPools.ts b/src/lib/db/quotaPools.ts index c4f70ec52b..88ddc45742 100644 --- a/src/lib/db/quotaPools.ts +++ b/src/lib/db/quotaPools.ts @@ -49,7 +49,10 @@ export interface PoolAllocation { export interface QuotaPool { id: string; + /** Primary / legacy single connection. Kept for back-compat. */ connectionId: string; + /** All member connections (≥1 after backfill). Primary is always connectionIds[0]. */ + connectionIds: string[]; name: string; createdAt: string; allocations: PoolAllocation[]; @@ -59,11 +62,22 @@ export interface PoolCreate { connectionId: string; name: string; allocations?: PoolAllocation[]; + /** + * Full member list. When provided, connectionId is ignored for the join table + * and connectionIds[0] is used as the primary. When omitted, defaults to + * [connectionId]. + */ + connectionIds?: string[]; } export interface PoolUpdate { name?: string; allocations?: PoolAllocation[]; + /** + * When provided, replaces the entire join-table membership for this pool. + * connection_id column is synced to connectionIds[0]. + */ + connectionIds?: string[]; } // --------------------------------------------------------------------------- @@ -112,10 +126,28 @@ function rowToAllocation(row: AllocationRow): PoolAllocation { return alloc; } +interface PoolConnectionRow { + connection_id: string; +} + +function getConnectionIds(poolId: string, fallbackConnectionId: string): string[] { + const rows = getDb() + .prepare( + "SELECT connection_id FROM quota_pool_connections WHERE pool_id = ? ORDER BY created_at ASC" + ) + .all(poolId); + if (rows.length > 0) { + return rows.map((r) => r.connection_id); + } + // Defensive fallback: join table empty (shouldn't happen post-backfill). + return fallbackConnectionId ? [fallbackConnectionId] : []; +} + function rowToPool(row: PoolRow, allocations: PoolAllocation[]): QuotaPool { return { id: row.id, connectionId: row.connection_id, + connectionIds: getConnectionIds(row.id, row.connection_id), name: row.name, createdAt: row.created_at, allocations, @@ -168,22 +200,55 @@ export function getPool(id: string): QuotaPool | null { } /** - * Create a new quota pool, optionally with initial allocations. + * Create a new quota pool, optionally with initial allocations and member connections. + * When `connectionIds` is provided, its first element becomes the primary connection_id. + * When omitted, defaults to [connectionId]. */ export function createPool(input: PoolCreate): QuotaPool { const id = makeId(); const now = new Date().toISOString(); - getDb() - .prepare("INSERT INTO quota_pools (id, connection_id, name, created_at) VALUES (?, ?, ?, ?)") - .run(id, input.connectionId, input.name, now); + // Resolve effective member list and primary connection. + const members: string[] = + input.connectionIds && input.connectionIds.length > 0 + ? input.connectionIds + : [input.connectionId]; + const primaryConnectionId = members[0]; - if (input.allocations && input.allocations.length > 0) { - upsertAllocations(id, input.allocations); - } + const database = getDb(); + const doCreate = database.transaction(() => { + database + .prepare("INSERT INTO quota_pools (id, connection_id, name, created_at) VALUES (?, ?, ?, ?)") + .run(id, primaryConnectionId, input.name, now); + + const insertConn = database.prepare( + "INSERT OR IGNORE INTO quota_pool_connections (pool_id, connection_id) VALUES (?, ?)" + ); + for (const connId of members) { + insertConn.run(id, connId); + } + + if (input.allocations && input.allocations.length > 0) { + const insertAlloc = database.prepare( + `INSERT INTO quota_allocations (pool_id, api_key_id, weight, cap_value, cap_unit, policy) + VALUES (?, ?, ?, ?, ?, ?)` + ); + for (const alloc of input.allocations) { + insertAlloc.run( + id, + alloc.apiKeyId, + alloc.weight, + alloc.capValue ?? null, + alloc.capUnit ?? null, + alloc.policy + ); + } + } + }); + doCreate(); const result = rowToPool( - { id, connection_id: input.connectionId, name: input.name, created_at: now }, + { id, connection_id: primaryConnectionId, name: input.name, created_at: now }, getAllocations(id) ); @@ -194,23 +259,63 @@ export function createPool(input: PoolCreate): QuotaPool { } /** - * Update an existing pool's name and/or allocations. + * Update an existing pool's name, allocations, and/or member connections. * Returns updated pool, or null if pool not found. + * When `connectionIds` is provided, the join table is replaced atomically and + * connection_id (primary) is synced to connectionIds[0]. */ export function updatePool(id: string, input: PoolUpdate): QuotaPool | null { - const existing = getDb() + const database = getDb(); + const existing = database .prepare("SELECT id, connection_id, name, created_at FROM quota_pools WHERE id = ?") .get(id); if (!existing) return null; - if (input.name !== undefined) { - getDb().prepare("UPDATE quota_pools SET name = ? WHERE id = ?").run(input.name, id); - existing.name = input.name; - } + const doUpdate = database.transaction(() => { + if (input.name !== undefined) { + database.prepare("UPDATE quota_pools SET name = ? WHERE id = ?").run(input.name, id); + existing.name = input.name; + } - if (input.allocations !== undefined) { - upsertAllocations(id, input.allocations); - } + if (input.connectionIds !== undefined && input.connectionIds.length > 0) { + const newPrimary = input.connectionIds[0]; + // Replace join rows. + database + .prepare("DELETE FROM quota_pool_connections WHERE pool_id = ?") + .run(id); + const insertConn = database.prepare( + "INSERT OR IGNORE INTO quota_pool_connections (pool_id, connection_id) VALUES (?, ?)" + ); + for (const connId of input.connectionIds) { + insertConn.run(id, connId); + } + // Sync primary column. + database + .prepare("UPDATE quota_pools SET connection_id = ? WHERE id = ?") + .run(newPrimary, id); + existing.connection_id = newPrimary; + } + + if (input.allocations !== undefined) { + // Inline the allocation upsert inside the transaction (avoids nested transaction). + database.prepare("DELETE FROM quota_allocations WHERE pool_id = ?").run(id); + const insertAlloc = database.prepare( + `INSERT INTO quota_allocations (pool_id, api_key_id, weight, cap_value, cap_unit, policy) + VALUES (?, ?, ?, ?, ?, ?)` + ); + for (const alloc of input.allocations) { + insertAlloc.run( + id, + alloc.apiKeyId, + alloc.weight, + alloc.capValue ?? null, + alloc.capUnit ?? null, + alloc.policy + ); + } + } + }); + doUpdate(); const result = rowToPool(existing, getAllocations(id)); @@ -222,6 +327,7 @@ export function updatePool(id: string, input: PoolUpdate): QuotaPool | null { /** * Delete a pool by id. CASCADE removes associated allocations. + * Also removes join rows in quota_pool_connections. * Returns true if a row was deleted, false if not found. */ export function deletePool(id: string): boolean { @@ -229,7 +335,12 @@ export function deletePool(id: string): boolean { // removeQuotaCombosForPool can still resolve the pool name → slug. void removeQuotaCombosGuarded(id); - const result = getDb().prepare("DELETE FROM quota_pools WHERE id = ?").run(id); + const database = getDb(); + const doDelete = database.transaction(() => { + database.prepare("DELETE FROM quota_pool_connections WHERE pool_id = ?").run(id); + return database.prepare("DELETE FROM quota_pools WHERE id = ?").run(id); + }); + const result = doDelete(); return result.changes > 0; } diff --git a/tests/unit/quota-pool-connections.test.ts b/tests/unit/quota-pool-connections.test.ts new file mode 100644 index 0000000000..6a5f93aab6 --- /dev/null +++ b/tests/unit/quota-pool-connections.test.ts @@ -0,0 +1,211 @@ +/** + * tests/unit/quota-pool-connections.test.ts + * + * Phase D1 — Multi-provider quota pools. + * + * Coverage: + * - Migration file exists and contains expected SQL. + * - createPool with connectionIds: [a, b] → getPool returns connectionIds.length === 2 + * and connectionId === a (primary). + * - updatePool replacing connectionIds → getPool reflects new set. + * - deletePool removes join rows. + * - Pool created the old way (single connectionId, no connectionIds arg) → + * connectionIds === [connectionId] (back-compat). + */ + +import test from "node:test"; +import assert from "node:assert/strict"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; + +// ── DB harness (same pattern as db-quota-pools.test.ts) ──────────────────── +const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-pool-conn-")); +process.env.DATA_DIR = TEST_DATA_DIR; + +const core = await import("../../src/lib/db/core.ts"); +const poolsDb = await import("../../src/lib/db/quotaPools.ts"); + +async function resetStorage() { + core.resetDbInstance(); + for (let attempt = 0; attempt < 10; attempt++) { + try { + if (fs.existsSync(TEST_DATA_DIR)) { + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); + } + break; + } catch (err: any) { + if ((err?.code === "EBUSY" || err?.code === "EPERM") && attempt < 9) { + await new Promise((resolve) => setTimeout(resolve, 50 * (attempt + 1))); + } else { + throw err; + } + } + } + fs.mkdirSync(TEST_DATA_DIR, { recursive: true }); +} + +test.beforeEach(async () => { + await resetStorage(); +}); + +test.after(async () => { + core.resetDbInstance(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); +}); + +// ── D1.1: Migration file ──────────────────────────────────────────────────── + +test("migration 086 file exists and contains quota_pool_connections DDL", () => { + const migrationPath = path.resolve( + "src/lib/db/migrations/086_quota_pool_connections.sql" + ); + assert.ok(fs.existsSync(migrationPath), `migration file not found: ${migrationPath}`); + + const sql = fs.readFileSync(migrationPath, "utf8"); + assert.ok( + sql.includes("quota_pool_connections"), + "migration SQL should reference quota_pool_connections" + ); + assert.ok( + sql.includes("INSERT OR IGNORE INTO quota_pool_connections"), + "migration SQL should contain the backfill INSERT" + ); + assert.ok( + sql.includes("SELECT id, connection_id FROM quota_pools"), + "backfill SELECT should reference quota_pools columns" + ); +}); + +// ── D1.2: createPool with connectionIds ──────────────────────────────────── + +test("createPool with connectionIds: [a, b] → connectionIds.length === 2, connectionId === a", () => { + const pool = poolsDb.createPool({ + connectionId: "conn-a", + name: "Multi-conn Pool", + connectionIds: ["conn-a", "conn-b"], + }); + + assert.ok(pool.id, "pool should have an id"); + assert.equal(pool.connectionId, "conn-a", "primary connectionId should be conn-a"); + assert.equal(pool.connectionIds.length, 2, "connectionIds should have 2 members"); + assert.ok(pool.connectionIds.includes("conn-a"), "conn-a should be a member"); + assert.ok(pool.connectionIds.includes("conn-b"), "conn-b should be a member"); + assert.equal(pool.connectionIds[0], "conn-a", "first connectionId should be the primary"); +}); + +test("getPool reflects both connectionIds after multi-connection create", () => { + const created = poolsDb.createPool({ + connectionId: "p-a", + name: "Pool ABC", + connectionIds: ["p-a", "p-b", "p-c"], + }); + + const found = poolsDb.getPool(created.id)!; + assert.ok(found, "pool should be found"); + assert.equal(found.connectionId, "p-a"); + assert.equal(found.connectionIds.length, 3); + assert.deepEqual([...found.connectionIds].sort(), ["p-a", "p-b", "p-c"].sort()); +}); + +// ── D1.3: updatePool replacing connectionIds ──────────────────────────────── + +test("updatePool with new connectionIds replaces the join rows", () => { + const pool = poolsDb.createPool({ + connectionId: "old-a", + name: "Updatable Pool", + connectionIds: ["old-a", "old-b"], + }); + + const updated = poolsDb.updatePool(pool.id, { + connectionIds: ["new-x", "new-y"], + }); + + assert.ok(updated, "updatePool should return the updated pool"); + assert.equal(updated!.connectionId, "new-x", "primary should be updated to new-x"); + assert.equal(updated!.connectionIds.length, 2); + assert.ok(updated!.connectionIds.includes("new-x"), "new-x should be a member"); + assert.ok(updated!.connectionIds.includes("new-y"), "new-y should be a member"); + assert.ok(!updated!.connectionIds.includes("old-a"), "old-a should be removed"); + assert.ok(!updated!.connectionIds.includes("old-b"), "old-b should be removed"); + + // Re-read from DB to confirm persistence. + const reread = poolsDb.getPool(pool.id)!; + assert.equal(reread.connectionId, "new-x"); + assert.deepEqual([...reread.connectionIds].sort(), ["new-x", "new-y"].sort()); +}); + +test("updatePool without connectionIds leaves join rows untouched", () => { + const pool = poolsDb.createPool({ + connectionId: "stable-a", + name: "Stable Pool", + connectionIds: ["stable-a", "stable-b"], + }); + + poolsDb.updatePool(pool.id, { name: "Renamed Pool" }); + + const reread = poolsDb.getPool(pool.id)!; + assert.equal(reread.name, "Renamed Pool"); + assert.equal(reread.connectionIds.length, 2, "connectionIds should be unchanged"); + assert.ok(reread.connectionIds.includes("stable-a")); + assert.ok(reread.connectionIds.includes("stable-b")); +}); + +// ── D1.4: deletePool removes join rows ──────────────────────────────────── + +test("deletePool removes quota_pool_connections rows", () => { + const pool = poolsDb.createPool({ + connectionId: "del-a", + name: "To Delete", + connectionIds: ["del-a", "del-b"], + }); + + const deleted = poolsDb.deletePool(pool.id); + assert.equal(deleted, true, "deletePool should return true"); + + // Pool should be gone. + assert.equal(poolsDb.getPool(pool.id), null, "pool should be null after deletion"); + + // The join rows are cleaned up — no ghost references. + // We verify indirectly: creating a new pool with the same connection IDs should work + // without PK conflicts in quota_pool_connections. + const newPool = poolsDb.createPool({ + connectionId: "del-a", + name: "Reused conn", + connectionIds: ["del-a", "del-b"], + }); + assert.ok(newPool.id, "new pool with same conn IDs should be created without conflict"); +}); + +// ── D1.5: Back-compat — single connectionId, no connectionIds arg ────────── + +test("pool created with single connectionId (legacy) returns connectionIds === [connectionId]", () => { + const pool = poolsDb.createPool({ + connectionId: "legacy-conn", + name: "Legacy Pool", + }); + + assert.equal(pool.connectionId, "legacy-conn"); + assert.deepEqual(pool.connectionIds, ["legacy-conn"]); + + const found = poolsDb.getPool(pool.id)!; + assert.deepEqual(found.connectionIds, ["legacy-conn"]); +}); + +test("listPools returns connectionIds on every pool", () => { + poolsDb.createPool({ connectionId: "lc-1", name: "Pool 1" }); + poolsDb.createPool({ + connectionId: "lc-2", + name: "Pool 2", + connectionIds: ["lc-2", "lc-3"], + }); + + const pools = poolsDb.listPools(); + assert.equal(pools.length, 2); + + const p1 = pools.find((p) => p.name === "Pool 1")!; + assert.deepEqual(p1.connectionIds, ["lc-1"]); + + const p2 = pools.find((p) => p.name === "Pool 2")!; + assert.equal(p2.connectionIds.length, 2); +});