diff --git a/src/__tests__/keys-db-fallback-helper.test.ts b/src/__tests__/keys-db-fallback-helper.test.ts new file mode 100644 index 0000000..18754d5 --- /dev/null +++ b/src/__tests__/keys-db-fallback-helper.test.ts @@ -0,0 +1,68 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; + +vi.unmock("../services/keys.js"); + +// The DB mock is set up in setup.ts — we need to control queryWithRetry +const mockQueryWithRetry = vi.fn(); +vi.mock("../services/db.js", () => ({ + default: { query: vi.fn(), end: vi.fn() }, + pool: { query: vi.fn(), end: vi.fn() }, + queryWithRetry: (...args: unknown[]) => mockQueryWithRetry(...args), + connectWithRetry: vi.fn(), + initDatabase: vi.fn(), + cleanupStaleData: vi.fn(), +})); + +import { findKeyInCacheOrDb } from "../services/keys.js"; + +describe("findKeyInCacheOrDb", () => { + beforeEach(() => { + mockQueryWithRetry.mockReset(); + }); + + it("returns null when DB finds no row", async () => { + mockQueryWithRetry.mockResolvedValue({ rows: [] }); + const result = await findKeyInCacheOrDb("stripe_customer_id", "cus_nonexistent"); + expect(result).toBeNull(); + expect(mockQueryWithRetry).toHaveBeenCalledWith( + expect.stringContaining("WHERE stripe_customer_id = $1"), + ["cus_nonexistent"] + ); + }); + + it("returns ApiKey when DB finds a row", async () => { + mockQueryWithRetry.mockResolvedValue({ + rows: [{ + key: "df_pro_abc", + tier: "pro", + email: "test@example.com", + created_at: "2026-01-01T00:00:00.000Z", + stripe_customer_id: "cus_123", + }], + }); + const result = await findKeyInCacheOrDb("stripe_customer_id", "cus_123"); + expect(result).toEqual({ + key: "df_pro_abc", + tier: "pro", + email: "test@example.com", + createdAt: "2026-01-01T00:00:00.000Z", + stripeCustomerId: "cus_123", + }); + }); + + it("handles Date objects in created_at", async () => { + mockQueryWithRetry.mockResolvedValue({ + rows: [{ + key: "df_pro_abc", + tier: "pro", + email: "test@example.com", + created_at: new Date("2026-01-01T00:00:00.000Z"), + stripe_customer_id: null, + }], + }); + const result = await findKeyInCacheOrDb("key", "df_pro_abc"); + expect(result).not.toBeNull(); + expect(result!.createdAt).toBe("2026-01-01T00:00:00.000Z"); + expect(result!.stripeCustomerId).toBeUndefined(); + }); +}); diff --git a/src/services/keys.ts b/src/services/keys.ts index 123bc7b..df62d5f 100644 --- a/src/services/keys.ts +++ b/src/services/keys.ts @@ -14,6 +14,26 @@ export interface ApiKey { // In-memory cache for fast lookups, synced with PostgreSQL let keysCache: ApiKey[] = []; +/** Look up a key row in the DB by a given column. Returns null if not found. */ +export async function findKeyInCacheOrDb( + column: "key" | "stripe_customer_id" | "email", + value: string +): Promise { + const result = await queryWithRetry( + `SELECT key, tier, email, created_at, stripe_customer_id FROM api_keys WHERE ${column} = $1 LIMIT 1`, + [value] + ); + if (result.rows.length === 0) return null; + const r = result.rows[0]; + return { + key: r.key, + tier: r.tier as "free" | "pro", + email: r.email, + createdAt: r.created_at instanceof Date ? r.created_at.toISOString() : r.created_at, + stripeCustomerId: r.stripe_customer_id || undefined, + }; +} + export async function loadKeys(): Promise { try { const result = await queryWithRetry( @@ -137,47 +157,22 @@ export async function downgradeByCustomer(stripeCustomerId: string): Promise { - // Check DB directly — survives pod restarts unlike in-memory cache - const result = await queryWithRetry( - "SELECT key, tier, email, created_at, stripe_customer_id FROM api_keys WHERE stripe_customer_id = $1 LIMIT 1", - [stripeCustomerId] - ); - if (result.rows.length === 0) return null; - const r = result.rows[0]; - return { - key: r.key, - tier: r.tier as "free" | "pro", - email: r.email, - createdAt: r.created_at instanceof Date ? r.created_at.toISOString() : r.created_at, - stripeCustomerId: r.stripe_customer_id || undefined, - }; + return findKeyInCacheOrDb("stripe_customer_id", stripeCustomerId); } export function getAllKeys(): ApiKey[] { @@ -194,27 +189,15 @@ export async function updateKeyEmail(apiKey: string, newEmail: string): Promise< // DB fallback: key may exist on another pod's cache or after a restart logger.info({ apiKey: apiKey.slice(0, 10) + "..." }, "updateKeyEmail: cache miss, falling back to DB"); - const result = await queryWithRetry( - "SELECT key, tier, email, created_at, stripe_customer_id FROM api_keys WHERE key = $1 LIMIT 1", - [apiKey] - ); - if (result.rows.length === 0) { + const dbKey = await findKeyInCacheOrDb("key", apiKey); + if (!dbKey) { logger.warn({ apiKey: apiKey.slice(0, 10) + "..." }, "updateKeyEmail: key not found in cache or DB"); return false; } - const row = result.rows[0]; await queryWithRetry("UPDATE api_keys SET email = $1 WHERE key = $2", [newEmail, apiKey]); - - // Hydrate local cache - const cached: ApiKey = { - key: row.key, - tier: row.tier as "free" | "pro", - email: newEmail, - createdAt: row.created_at instanceof Date ? row.created_at.toISOString() : row.created_at, - stripeCustomerId: row.stripe_customer_id || undefined, - }; - keysCache.push(cached); + dbKey.email = newEmail; + keysCache.push(dbKey); logger.info({ apiKey: apiKey.slice(0, 10) + "..." }, "updateKeyEmail: updated via DB fallback"); return true; @@ -230,28 +213,16 @@ export async function updateEmailByCustomer(stripeCustomerId: string, newEmail: // DB fallback: key may exist on another pod's cache or after a restart logger.info({ stripeCustomerId }, "updateEmailByCustomer: cache miss, falling back to DB"); - const result = await queryWithRetry( - "SELECT key, tier, email, created_at, stripe_customer_id FROM api_keys WHERE stripe_customer_id = $1 LIMIT 1", - [stripeCustomerId] - ); - if (result.rows.length === 0) { + const dbKey = await findKeyInCacheOrDb("stripe_customer_id", stripeCustomerId); + if (!dbKey) { logger.warn({ stripeCustomerId }, "updateEmailByCustomer: customer not found in cache or DB"); return false; } - const row = result.rows[0]; await queryWithRetry("UPDATE api_keys SET email = $1 WHERE stripe_customer_id = $2", [newEmail, stripeCustomerId]); + dbKey.email = newEmail; + keysCache.push(dbKey); - // Hydrate local cache - const cached: ApiKey = { - key: row.key, - tier: row.tier as "free" | "pro", - email: newEmail, - createdAt: row.created_at instanceof Date ? row.created_at.toISOString() : row.created_at, - stripeCustomerId: row.stripe_customer_id || undefined, - }; - keysCache.push(cached); - - logger.info({ stripeCustomerId, key: row.key }, "updateEmailByCustomer: updated via DB fallback"); + logger.info({ stripeCustomerId, key: dbKey.key }, "updateEmailByCustomer: updated via DB fallback"); return true; }