From 1abf7865720769088e8fcd1e2ee3d3d456372263 Mon Sep 17 00:00:00 2001 From: Bo Lu Date: Sat, 25 Apr 2026 13:47:04 +0800 Subject: [PATCH] feat(db): add SQLite module with migrations (Task #1) Introduce bun:sqlite-backed db module with WAL mode, foreign keys, and a meta-table-driven migration runner. Initial schema (001) creates accounts, model_pricing, model_pricing_versions, pricing_sync_log, usage_events, usage_daily and supporting indexes. - src/lib/db.ts: initDb/getDb/withTransaction + test-only reset helper - src/lib/migrations/001_initial.sql: full schema per design doc 02 - src/lib/paths.ts: USAGE_DB_PATH under APP_DIR - src/start.ts: --db-path flag, initDb after ensurePaths - src/types.d.ts: declare *.sql text imports - tests/db.test.ts: 8 tests covering init, idempotency, tx, WAL, indexes Closes #1 Co-Authored-By: Claude Opus 4 --- src/lib/db.ts | 99 +++++++++++++++++ src/lib/migrations/001_initial.sql | 113 +++++++++++++++++++ src/lib/paths.ts | 2 + src/start.ts | 12 +- src/types.d.ts | 4 + tests/db.test.ts | 170 +++++++++++++++++++++++++++++ 6 files changed, 399 insertions(+), 1 deletion(-) create mode 100644 src/lib/db.ts create mode 100644 src/lib/migrations/001_initial.sql create mode 100644 src/types.d.ts create mode 100644 tests/db.test.ts diff --git a/src/lib/db.ts b/src/lib/db.ts new file mode 100644 index 000000000..8df341dea --- /dev/null +++ b/src/lib/db.ts @@ -0,0 +1,99 @@ +import { Database } from "bun:sqlite" +import fs from "node:fs" +import path from "node:path" + +import migration001 from "./migrations/001_initial.sql" with { type: "text" } + +export const CURRENT_SCHEMA_VERSION = 1 + +const MIGRATIONS: Array<{ version: number; sql: string }> = [ + { version: 1, sql: migration001 }, +] + +let dbInstance: Database | undefined + +export function initDb(dbPath: string): Database { + if (dbInstance) return dbInstance + + if (dbPath !== ":memory:") { + fs.mkdirSync(path.dirname(dbPath), { recursive: true }) + } + + const db = new Database(dbPath, { create: true }) + + // Pragmas — set before any schema work. + db.run("PRAGMA journal_mode = WAL") + db.run("PRAGMA synchronous = NORMAL") + db.run("PRAGMA foreign_keys = ON") + + runMigrations(db) + + dbInstance = db + return db +} + +export function getDb(): Database { + if (!dbInstance) { + throw new Error( + "Database not initialized. Call initDb(path) before getDb().", + ) + } + return dbInstance +} + +export function withTransaction(fn: (db: Database) => T): T { + const db = getDb() + const tx = db.transaction((arg: () => T) => arg()) + return tx(() => fn(db)) +} + +/** + * Test-only helper. Closes any current instance and clears the singleton so + * the next initDb call starts from scratch. Production code must never call + * this — it exists to keep tests isolated. + */ +export function __resetDbForTests(): void { + if (dbInstance) { + try { + dbInstance.close() + } catch { + // ignore + } + dbInstance = undefined + } +} + +function runMigrations(db: Database): void { + // Bootstrap meta table so we can read schema_version. + db.run( + "CREATE TABLE IF NOT EXISTS meta (key TEXT PRIMARY KEY, value TEXT NOT NULL)", + ) + + const row = db + .query< + { value: string }, + [] + >("SELECT value FROM meta WHERE key='schema_version'") + .get() + const currentVersion = row ? Number.parseInt(row.value, 10) : 0 + + const pending = MIGRATIONS.filter((m) => m.version > currentVersion).sort( + (a, b) => a.version - b.version, + ) + + if (pending.length === 0) return + + const apply = db.transaction(() => { + for (const m of pending) { + // Migration SQL contains multiple statements; only exec() handles that. + // eslint-disable-next-line @typescript-eslint/no-deprecated + db.exec(m.sql) + } + db.run( + "INSERT INTO meta (key, value) VALUES ('schema_version', ?) " + + "ON CONFLICT(key) DO UPDATE SET value=excluded.value", + [String(CURRENT_SCHEMA_VERSION)], + ) + }) + apply() +} diff --git a/src/lib/migrations/001_initial.sql b/src/lib/migrations/001_initial.sql new file mode 100644 index 000000000..3f91480e5 --- /dev/null +++ b/src/lib/migrations/001_initial.sql @@ -0,0 +1,113 @@ +-- Initial schema for copilot-api multi-account & usage billing. +-- Owned by tasks #1 (this migration), #2 (accounts loader), #6 (usage recorder), +-- #13 (pricing version writer). See docs/design/. + +CREATE TABLE IF NOT EXISTS accounts ( + name TEXT PRIMARY KEY, + account_type TEXT NOT NULL, + created_at INTEGER NOT NULL +); + +CREATE TABLE IF NOT EXISTS model_pricing ( + model_id TEXT PRIMARY KEY, + input_per_mtok REAL, + cached_input_per_mtok REAL, + output_per_mtok REAL, + reasoning_per_mtok REAL, + premium_multiplier REAL, + premium_unit_price REAL, + currency TEXT NOT NULL DEFAULT 'USD', + source TEXT, + source_skus TEXT, + updated_at INTEGER NOT NULL +); + +CREATE TABLE IF NOT EXISTS pricing_sync_log ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + ts INTEGER NOT NULL, + status TEXT NOT NULL, + source_count INTEGER, + llm_model TEXT, + models_updated INTEGER, + models_rejected INTEGER, + error TEXT, + raw_request_json TEXT, + raw_response_json TEXT, + diff_json TEXT +); + +CREATE TABLE IF NOT EXISTS model_pricing_versions ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + model_id TEXT NOT NULL, + effective_from INTEGER NOT NULL, + effective_to INTEGER, + input_per_mtok REAL, + cached_input_per_mtok REAL, + output_per_mtok REAL, + reasoning_per_mtok REAL, + premium_multiplier REAL, + premium_unit_price REAL, + currency TEXT NOT NULL DEFAULT 'USD', + source TEXT, + source_skus TEXT, + sync_log_id INTEGER, + created_at INTEGER NOT NULL, + FOREIGN KEY (sync_log_id) REFERENCES pricing_sync_log(id) +); + +CREATE INDEX IF NOT EXISTS idx_pricing_versions_model_time + ON model_pricing_versions(model_id, effective_from); + +CREATE INDEX IF NOT EXISTS idx_pricing_versions_current + ON model_pricing_versions(model_id) WHERE effective_to IS NULL; + +CREATE TABLE IF NOT EXISTS usage_events ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + ts INTEGER NOT NULL, + account_name TEXT NOT NULL, + model_id TEXT NOT NULL, + endpoint TEXT NOT NULL, + upstream_format TEXT NOT NULL, + is_streaming INTEGER NOT NULL, + input_tokens INTEGER DEFAULT 0, + cached_input_tokens INTEGER DEFAULT 0, + output_tokens INTEGER DEFAULT 0, + reasoning_tokens INTEGER DEFAULT 0, + total_tokens INTEGER DEFAULT 0, + premium_request_count REAL DEFAULT 0, + input_price_snapshot REAL, + cached_input_price_snapshot REAL, + output_price_snapshot REAL, + reasoning_price_snapshot REAL, + premium_unit_price_snapshot REAL, + premium_multiplier_snapshot REAL, + request_id TEXT, + status TEXT NOT NULL, + duration_ms INTEGER, + FOREIGN KEY (account_name) REFERENCES accounts(name) +); + +CREATE INDEX IF NOT EXISTS idx_usage_account_model_ts + ON usage_events(account_name, model_id, ts); + +CREATE INDEX IF NOT EXISTS idx_usage_ts ON usage_events(ts); + +CREATE TABLE IF NOT EXISTS usage_daily ( + day TEXT NOT NULL, + account_name TEXT NOT NULL, + model_id TEXT NOT NULL, + endpoint TEXT NOT NULL, + req_count INTEGER NOT NULL DEFAULT 0, + input_tokens INTEGER NOT NULL DEFAULT 0, + cached_input_tokens INTEGER NOT NULL DEFAULT 0, + output_tokens INTEGER NOT NULL DEFAULT 0, + reasoning_tokens INTEGER NOT NULL DEFAULT 0, + total_tokens INTEGER NOT NULL DEFAULT 0, + premium_requests REAL NOT NULL DEFAULT 0, + PRIMARY KEY (day, account_name, model_id, endpoint) +); + +CREATE TABLE IF NOT EXISTS meta ( + key TEXT PRIMARY KEY, + value TEXT NOT NULL +); diff --git a/src/lib/paths.ts b/src/lib/paths.ts index 8d0a9f02b..231560d12 100644 --- a/src/lib/paths.ts +++ b/src/lib/paths.ts @@ -5,10 +5,12 @@ import path from "node:path" const APP_DIR = path.join(os.homedir(), ".local", "share", "copilot-api") const GITHUB_TOKEN_PATH = path.join(APP_DIR, "github_token") +const USAGE_DB_PATH = path.join(APP_DIR, "usage.sqlite") export const PATHS = { APP_DIR, GITHUB_TOKEN_PATH, + USAGE_DB_PATH, } export async function ensurePaths(): Promise { diff --git a/src/start.ts b/src/start.ts index 14abbbdff..d3539362a 100644 --- a/src/start.ts +++ b/src/start.ts @@ -6,7 +6,8 @@ import consola from "consola" import { serve, type ServerHandler } from "srvx" import invariant from "tiny-invariant" -import { ensurePaths } from "./lib/paths" +import { initDb } from "./lib/db" +import { ensurePaths, PATHS } from "./lib/paths" import { initProxyFromEnv } from "./lib/proxy" import { generateEnvScript } from "./lib/shell" import { state } from "./lib/state" @@ -25,6 +26,7 @@ interface RunServerOptions { claudeCode: boolean showToken: boolean proxyEnv: boolean + dbPath: string } export async function runServer(options: RunServerOptions): Promise { @@ -48,6 +50,7 @@ export async function runServer(options: RunServerOptions): Promise { state.showToken = options.showToken await ensurePaths() + initDb(options.dbPath) await cacheVSCodeVersion() if (options.githubToken) { @@ -184,6 +187,12 @@ export const start = defineCommand({ default: false, description: "Initialize proxy from environment variables", }, + "db-path": { + type: "string", + default: PATHS.USAGE_DB_PATH, + description: + "Path to the usage SQLite database (defaults to ~/.local/share/copilot-api/usage.sqlite)", + }, }, run({ args }) { const rateLimitRaw = args["rate-limit"] @@ -202,6 +211,7 @@ export const start = defineCommand({ claudeCode: args["claude-code"], showToken: args["show-token"], proxyEnv: args["proxy-env"], + dbPath: args["db-path"], }) }, }) diff --git a/src/types.d.ts b/src/types.d.ts new file mode 100644 index 000000000..28cb1e264 --- /dev/null +++ b/src/types.d.ts @@ -0,0 +1,4 @@ +declare module "*.sql" { + const content: string + export default content +} diff --git a/tests/db.test.ts b/tests/db.test.ts new file mode 100644 index 000000000..b61512c41 --- /dev/null +++ b/tests/db.test.ts @@ -0,0 +1,170 @@ +import { test, expect, describe, beforeEach } from "bun:test" +import fs from "node:fs" +import os from "node:os" +import path from "node:path" + +import { + initDb, + getDb, + withTransaction, + CURRENT_SCHEMA_VERSION, + __resetDbForTests, +} from "../src/lib/db" + +const tmpDbPath = () => + path.join( + os.tmpdir(), + `copilot-api-test-${Date.now()}-${Math.random().toString(36).slice(2)}.sqlite`, + ) + +describe("db module", () => { + beforeEach(() => { + __resetDbForTests() + }) + + test("initDb on a fresh path creates all tables and sets schema_version", () => { + const p = tmpDbPath() + const db = initDb(p) + + const tables = db + .query<{ name: string }, []>( + "SELECT name FROM sqlite_master WHERE type='table' ORDER BY name", + ) + .all() + .map((r) => r.name) + + for (const t of [ + "accounts", + "model_pricing", + "model_pricing_versions", + "pricing_sync_log", + "usage_daily", + "usage_events", + "meta", + ]) { + expect(tables).toContain(t) + } + + const ver = db + .query< + { value: string }, + [] + >("SELECT value FROM meta WHERE key='schema_version'") + .get() + expect(ver?.value).toBe(String(CURRENT_SCHEMA_VERSION)) + + db.close() + fs.unlinkSync(p) + }) + + test("initDb is idempotent: running twice leaves schema_version unchanged and does not duplicate rows", () => { + const p = tmpDbPath() + const db1 = initDb(p) + db1.run( + "INSERT INTO meta (key, value) VALUES ('marker', 'persisted') " + + "ON CONFLICT(key) DO UPDATE SET value=excluded.value", + ) + db1.close() + + __resetDbForTests() + const db2 = initDb(p) + const marker = db2 + .query<{ value: string }, []>("SELECT value FROM meta WHERE key='marker'") + .get() + expect(marker?.value).toBe("persisted") + + const ver = db2 + .query< + { value: string }, + [] + >("SELECT value FROM meta WHERE key='schema_version'") + .get() + expect(ver?.value).toBe(String(CURRENT_SCHEMA_VERSION)) + + db2.close() + fs.unlinkSync(p) + }) + + test("getDb throws before initDb is called", () => { + expect(() => getDb()).toThrow() + }) + + test("getDb returns the initialized instance", () => { + const p = tmpDbPath() + const db = initDb(p) + expect(getDb()).toBe(db) + db.close() + fs.unlinkSync(p) + }) + + test("withTransaction commits on success", () => { + const p = tmpDbPath() + const db = initDb(p) + + withTransaction((d) => { + d.run( + "INSERT INTO accounts (name, account_type, created_at) " + + "VALUES ('a', 'individual', 1)", + ) + }) + + const row = db + .query<{ name: string }, []>("SELECT name FROM accounts WHERE name='a'") + .get() + expect(row?.name).toBe("a") + + db.close() + fs.unlinkSync(p) + }) + + test("withTransaction rolls back on throw", () => { + const p = tmpDbPath() + const db = initDb(p) + + expect(() => + withTransaction((d) => { + d.run( + "INSERT INTO accounts (name, account_type, created_at) " + + "VALUES ('b', 'individual', 1)", + ) + throw new Error("boom") + }), + ).toThrow("boom") + + const row = db + .query<{ name: string }, []>("SELECT name FROM accounts WHERE name='b'") + .get() + expect(row).toBeNull() + + db.close() + fs.unlinkSync(p) + }) + + test("WAL mode is enabled", () => { + const p = tmpDbPath() + const db = initDb(p) + const mode = db + .query<{ journal_mode: string }, []>("PRAGMA journal_mode") + .get() + expect(mode?.journal_mode.toLowerCase()).toBe("wal") + db.close() + fs.unlinkSync(p) + }) + + test("schema includes expected indexes", () => { + const p = tmpDbPath() + const db = initDb(p) + const idxs = db + .query<{ name: string }, []>( + "SELECT name FROM sqlite_master WHERE type='index'", + ) + .all() + .map((r) => r.name) + expect(idxs).toContain("idx_usage_account_model_ts") + expect(idxs).toContain("idx_usage_ts") + expect(idxs).toContain("idx_pricing_versions_model_time") + expect(idxs).toContain("idx_pricing_versions_current") + db.close() + fs.unlinkSync(p) + }) +})