import fs from 'node:fs'; import path from 'node:path'; import pg, { type PoolClient } from 'pg'; import type Database from 'better-sqlite3'; import dotenv from 'dotenv'; import { randomBytes, scryptSync } from 'node:crypto'; dotenv.config(); type SqliteDatabase = Database.Database; type Connection = PoolClient | SqliteDatabase; const databaseUrl = process.env.DATABASE_URL?.trim(); export const databaseDialect = databaseUrl ? 'postgres' : 'sqlite'; let sqlite: SqliteDatabase | undefined; let pool: pg.Pool | undefined; if (databaseUrl) { if (databaseUrl === 'pg-mem://' && process.env.NODE_ENV === 'test') { const { newDb } = await import('pg-mem'); const adapter = newDb({ autoCreateForeignKeyIndices: true }).adapters.createPg(); pool = new adapter.Pool() as unknown as pg.Pool; } else { pool = new pg.Pool({ connectionString: databaseUrl, ssl: process.env.PGSSL === 'disable' ? false : undefined, max: Number(process.env.PG_POOL_MAX || 10) }); } let schema = fs.readFileSync(path.resolve('db/postgres/schema.sql'), 'utf8'); if (databaseUrl === 'pg-mem://') schema = schema.replace(/CREATE UNIQUE INDEX IF NOT EXISTS one_active_storage_config[^;]+;/, ''); await pool.query(schema); const userCount = Number((await pool.query('SELECT COUNT(*)::int AS count FROM users')).rows[0].count); if (userCount === 0) { const initialPassword = process.env.INITIAL_ADMIN_PASSWORD; if (!initialPassword) throw new Error('空 PostgreSQL 数据库需要配置 INITIAL_ADMIN_PASSWORD'); const salt = randomBytes(16).toString('hex'); const passwordHash = `${salt}:${scryptSync(initialPassword, salt, 64).toString('hex')}`; await pool.query('INSERT INTO users (group_id, username, display_name, password_hash, role, must_change_password) VALUES (NULL, $1, $2, $3, $4, TRUE)', [process.env.INITIAL_ADMIN_USERNAME || 'admin', process.env.INITIAL_ADMIN_DISPLAY_NAME || '平台管理员', passwordHash, 'platform_admin']); } } else { sqlite = (await import('./db.js')).db; } function postgresSql(sql: string): string { let index = 0; return sql.replace(/\?/g, () => `$${++index}`); } async function connectionQuery(connection: Connection, sql: string, params: unknown[]): Promise<{ rows: T[]; rowCount: number }> { if (databaseDialect === 'postgres') { const result = await (connection as PoolClient).query(postgresSql(sql), params); return { rows: result.rows as T[], rowCount: result.rowCount ?? 0 }; } const statement = (connection as SqliteDatabase).prepare(sql); const sqliteParams = params.map((value) => typeof value === 'boolean' ? Number(value) : value); if (/^\s*(SELECT|WITH|PRAGMA)/i.test(sql) || /\bRETURNING\b/i.test(sql)) { return { rows: statement.all(...sqliteParams) as T[], rowCount: 0 }; } const result = statement.run(...sqliteParams); return { rows: [], rowCount: result.changes }; } export interface QueryContext { all(sql: string, params?: unknown[]): Promise; one(sql: string, params?: unknown[]): Promise; execute(sql: string, params?: unknown[]): Promise<{ changes: number }>; insertId(sql: string, params?: unknown[]): Promise; } function context(connection: Connection): QueryContext { return { async all(sql, params = []) { return (await connectionQuery(connection, sql, params)).rows; }, async one(sql, params = []) { return (await connectionQuery(connection, sql, params)).rows[0]; }, async execute(sql, params = []) { return { changes: (await connectionQuery(connection, sql, params)).rowCount }; }, async insertId(sql, params = []) { if (databaseDialect === 'postgres') { const result = await connectionQuery<{ id: number }>(connection, `${sql.replace(/;\s*$/, '')} RETURNING id`, params); return Number(result.rows[0]?.id); } const sqliteParams = params.map((value) => typeof value === 'boolean' ? Number(value) : value); const result = (connection as SqliteDatabase).prepare(sql).run(...sqliteParams); return Number(result.lastInsertRowid); }, }; } export const database: QueryContext = context((pool ?? sqlite) as Connection); export async function withTransaction(work: (tx: QueryContext) => Promise): Promise { if (databaseDialect === 'postgres') { const client = await pool!.connect(); try { await client.query('BEGIN'); const result = await work(context(client)); await client.query('COMMIT'); return result; } catch (error) { await client.query('ROLLBACK'); throw error; } finally { client.release(); } } sqlite!.exec('BEGIN IMMEDIATE'); try { const result = await work(context(sqlite!)); sqlite!.exec('COMMIT'); return result; } catch (error) { sqlite!.exec('ROLLBACK'); throw error; } } export async function closeDatabase(): Promise { if (pool) await pool.end(); if (sqlite) sqlite.close(); }