delivery-desk/api/database.ts

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[^;]+;/, '')
      .replace(/-- COLLECTION_STATUS_REPAIR_START[\s\S]+?-- COLLECTION_STATUS_REPAIR_END/, '')
      .replace(/-- REVIEW_ROUND_REPAIR_START[\s\S]+?-- REVIEW_ROUND_REPAIR_END/, '')
      .replace(/-- PROJECT_REVIEW_STATUS_REPAIR_START[\s\S]+?-- PROJECT_REVIEW_STATUS_REPAIR_END/, '')
      .replace(/-- TEXT_ANNOTATION_TARGET_REPAIR_START[\s\S]+?-- TEXT_ANNOTATION_TARGET_REPAIR_END/, '')
      .replace(/-- SINGLE_SCHEME_REPAIR_START[\s\S]+?-- SINGLE_SCHEME_REPAIR_END/, '');
  }
  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<T>(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<T>(sql: string, params?: unknown[]): Promise<T[]>;
  one<T>(sql: string, params?: unknown[]): Promise<T | undefined>;
  execute(sql: string, params?: unknown[]): Promise<{ changes: number }>;
  insertId(sql: string, params?: unknown[]): Promise<number>;
}

function context(connection: Connection): QueryContext {
  return {
    async all<T>(sql, params = []) { return (await connectionQuery<T>(connection, sql, params)).rows; },
    async one<T>(sql, params = []) { return (await connectionQuery<T>(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<T>(work: (tx: QueryContext) => Promise<T>): Promise<T> {
  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<void> {
  if (pool) await pool.end();
  if (sqlite) sqlite.close();
}