码桶

发现社区成员的开源项目

奶狗 /

ng-webot

公开
main
db.js14.5 KB
/**
 * 数据库访问层 — node:sqlite (内置, 文件型, 无需 MySQL)
 * 对外接口与原 mysql2 版本完全一致: row / rows / exec / lastInsertId / getPool / testConnect / ensureDatabase
 */
const { DatabaseSync } = require('node:sqlite');
const fs = require('fs');
const path = require('path');
const config = require('../config');

let db = null;

/** 将 node:sqlite 返回的 BigInt 等归一化为普通 JS 类型 */
function normalize(value) {
  if (typeof value === 'bigint') return Number(value);
  if (Array.isArray(value)) return value.map(normalize);
  if (value && typeof value === 'object') {
    const out = {};
    for (const k of Object.keys(value)) out[k] = normalize(value[k]);
    return out;
  }
  return value;
}

/** 获取(惰性初始化)数据库句柄,并自动建表 */
function raw() {
  if (db) return db;

  const dbPath = config.db.sqlite_path;
  const dir = path.dirname(dbPath);
  if (!fs.existsSync(dir)) fs.mkdirSync(dir, { recursive: true });

  db = new DatabaseSync(dbPath);
  db.exec('PRAGMA journal_mode = WAL;');
  db.exec('PRAGMA foreign_keys = ON;');
  // 写冲突时等待而非立即抛 SQLITE_BUSY(默认 0),避免高并发(后台调度+手动操作同时写库)偶发 500
  db.exec('PRAGMA busy_timeout = 5000;');
  runSchema();
  return db;
}

/**
 * 执行数据库操作并自动重试因写锁(SQLITE_BUSY)导致的失败。
 * node:sqlite 在高并发写时偶发 SQLITE_BUSY,重试可消除绝大多数偶发 500/阻塞。
 */
async function withRetry(fn, label = 'db') {
  let lastErr;
  for (let attempt = 0; attempt < 5; attempt++) {
    try {
      return await fn();
    } catch (e) {
      lastErr = e;
      const busy = /SQLITE_BUSY|database is locked|busy/i.test(e && e.message ? e.message : String(e));
      if (!busy || attempt === 4) break;
      await new Promise(r => setTimeout(r, 50 * (attempt + 1)));
    }
  }
  throw lastErr;
}

/** 执行 schema.sql(SQLite 版) */
function runSchema() {
  const schemaPath = path.join(__dirname, '..', 'data', 'schema.sql');
  if (!fs.existsSync(schemaPath)) {
    console.error('[db] schema.sql 不存在,跳过建表: ' + schemaPath);
    return;
  }
  const sql = fs.readFileSync(schemaPath, 'utf8');
  const statements = sql
    .split(';')
    .map(s => s.trim())
    .filter(s => s.length > 0);

  // 取某表列名集合(每表仅一次 PRAGMA)
  const colsOf = (t) => {
    try { return new Set(db.prepare(`PRAGMA table_info(${t})`).all().map(c => c.name)); }
    catch { return new Set(); }
  };
  // 幂等新增列(列已存在则跳过)
  const addCol = (t, col, ddl) => {
    try { db.exec(`ALTER TABLE ${t} ADD COLUMN ${ddl}`); }
    catch (e) { console.error(`[db] 迁移 ${t}.${col} 失败:`, e.message); }
  };

  // 建表 + 全部迁移包进单事务,消除逐条语句的 WAL fsync(显著提升启动速度)
  db.exec('BEGIN');
  try {
    for (const stmt of statements) db.exec(stmt);

    // bots.name
    if (!colsOf('bots').has('name')) addCol('bots', 'name', 'name VARCHAR(128)');

    // users 表补齐字段(合并为一次列查询,替代原先 3 次 PRAGMA)
    const ucols = colsOf('users');
    const userMigrations = [
      ['points', 'points INTEGER NOT NULL DEFAULT 0'],
      ['checkin_at', 'checkin_at INTEGER NOT NULL DEFAULT 0'],
      ['checkin_streak', 'checkin_streak INTEGER NOT NULL DEFAULT 0'],
      ['email', "email VARCHAR(255) NOT NULL DEFAULT ''"],
      ['email_verified', 'email_verified INTEGER NOT NULL DEFAULT 0'],
      ['nickname', "nickname VARCHAR(64) NOT NULL DEFAULT ''"],
      ['member_type', "member_type VARCHAR(16) NOT NULL DEFAULT 'none'"],
      ['member_until', 'member_until INTEGER NOT NULL DEFAULT 0'],
      ['ai_tokens_used', 'ai_tokens_used INTEGER NOT NULL DEFAULT 0'],
      ['ai_tokens_limit', 'ai_tokens_limit INTEGER NOT NULL DEFAULT 100000'],
      ['image_gen_used', 'image_gen_used INTEGER NOT NULL DEFAULT 0'],
      ['image_gen_limit', 'image_gen_limit INTEGER NOT NULL DEFAULT 5'],
      ['profile_url', "profile_url VARCHAR(512) NOT NULL DEFAULT ''"],
      ['ai_persona', "ai_persona VARCHAR(64) NOT NULL DEFAULT ''"],
      ['user_code', "user_code VARCHAR(6) NOT NULL DEFAULT ''"],
      ['companion_enabled', 'companion_enabled INTEGER NOT NULL DEFAULT 0'],
      ['companion_last_at', 'companion_last_at INTEGER NOT NULL DEFAULT 0'],
    ];
    for (const [col, ddl] of userMigrations) {
      if (!ucols.has(col)) addCol('users', col, ddl);
    }

    // 为用户分配专属 6 位随机 ID(无规律纯数字),回填历史用户
    {
      const usedCodes = new Set(
        db.prepare("SELECT user_code FROM users WHERE user_code IS NOT NULL AND user_code <> ''")
          .all().map((r) => r.user_code)
      );
      const emptyUsers = db.prepare("SELECT id FROM users WHERE user_code IS NULL OR user_code = ''").all();
      for (const row of emptyUsers) {
        let code;
        do { code = String(Math.floor(Math.random() * 1000000)).padStart(6, '0'); } while (usedCodes.has(code));
        usedCodes.add(code);
        db.prepare('UPDATE users SET user_code = ? WHERE id = ?').run(code, row.id);
      }
    }

    // 兑换卡表
    db.exec(`CREATE TABLE IF NOT EXISTS cards (
      id INTEGER PRIMARY KEY AUTOINCREMENT,
      code VARCHAR(32) NOT NULL,
      type VARCHAR(16) NOT NULL,
      value INTEGER NOT NULL DEFAULT 0,
      used_by INTEGER NOT NULL DEFAULT 0,
      used_at INTEGER NOT NULL DEFAULT 0,
      created_by INTEGER NOT NULL DEFAULT 0,
      created_at INTEGER NOT NULL,
      UNIQUE (code)
    )`);

    // plugins.expires_at(积分兑换有效期,0=永久)
    if (!colsOf('plugins').has('expires_at')) {
      addCol('plugins', 'expires_at', 'expires_at INTEGER NOT NULL DEFAULT 0');
    }

    // plugin_market 补齐 is_free / price / access(v4+:收费方式 0免费 1付费 2会员免费 3会员付费)
    const mcols = colsOf('plugin_market');
    if (!mcols.has('is_free')) addCol('plugin_market', 'is_free', 'is_free INTEGER NOT NULL DEFAULT 1');
    if (!mcols.has('price')) addCol('plugin_market', 'price', 'price INTEGER NOT NULL DEFAULT 0');
    if (!mcols.has('access')) addCol('plugin_market', 'access', 'access INTEGER NOT NULL DEFAULT 0');
    if (!mcols.has('listed')) addCol('plugin_market', 'listed', 'listed INTEGER NOT NULL DEFAULT 1');
    // 兼容旧数据:原 is_free=0 的视为付费(access=1)
    try { db.exec('UPDATE plugin_market SET access = 1 WHERE is_free = 0 AND access = 0'); } catch (_) {}

    // messages.status / error_msg(消息状态追踪:ok/failed/pending)
    if (!colsOf('messages').has('status')) addCol('messages', 'status', "status VARCHAR(16) NOT NULL DEFAULT 'ok'");
    if (!colsOf('messages').has('error_msg')) addCol('messages', 'error_msg', "error_msg TEXT NOT NULL DEFAULT ''");

    // 消息表索引:加速 bot_messages 列表查询、worker 去重查询、resolveCtx context_token 查询
    db.exec('CREATE INDEX IF NOT EXISTS idx_messages_bot_id ON messages(bot_id)');
    db.exec('CREATE INDEX IF NOT EXISTS idx_messages_bot_dir ON messages(bot_id, direction)');
    db.exec('CREATE INDEX IF NOT EXISTS idx_messages_bot_ctx ON messages(bot_id, context_token)');
    db.exec('CREATE INDEX IF NOT EXISTS idx_messages_peer ON messages(bot_id, peer_id)');

    // 自定义技能表(用户级,跨 bot 共享)
    db.exec(`CREATE TABLE IF NOT EXISTS custom_skills (
      id INTEGER PRIMARY KEY AUTOINCREMENT,
      user_id INTEGER NOT NULL,
      name VARCHAR(128) NOT NULL,
      icon VARCHAR(8) NOT NULL DEFAULT '',
      description TEXT NOT NULL DEFAULT '',
      prompt TEXT NOT NULL,
      created_at INTEGER NOT NULL DEFAULT (strftime('%s','now')),
      updated_at INTEGER NOT NULL DEFAULT 0
    )`);

    // 用户挂售插件表(插件市场—用户上传售卖)
    db.exec(`CREATE TABLE IF NOT EXISTS marketplace_listings (
      id INTEGER PRIMARY KEY AUTOINCREMENT,
      seller_user_id INTEGER NOT NULL,
      market_id VARCHAR(64) NOT NULL UNIQUE,
      name VARCHAR(128) NOT NULL,
      description TEXT NOT NULL DEFAULT '',
      category VARCHAR(32) NOT NULL DEFAULT '其他',
      price INTEGER NOT NULL DEFAULT 0,
      version VARCHAR(16) NOT NULL DEFAULT '1.0.0',
      downloads INTEGER NOT NULL DEFAULT 0,
      status VARCHAR(16) NOT NULL DEFAULT 'pending',
      reviewed_by INTEGER DEFAULT NULL,
      review_note TEXT NOT NULL DEFAULT '',
      created_at INTEGER NOT NULL,
      updated_at INTEGER NOT NULL
    )`);
      db.exec('CREATE INDEX IF NOT EXISTS idx_ml_status ON marketplace_listings(status)');
      db.exec('CREATE INDEX IF NOT EXISTS idx_ml_seller ON marketplace_listings(seller_user_id)');

    // 区分「自用插件」(kind='self') 与「挂售插件」(kind='market')
    const mlcols = colsOf('marketplace_listings');
    if (!mlcols.has('kind')) {
      addCol('marketplace_listings', 'kind', "kind VARCHAR(8) NOT NULL DEFAULT 'market'");
      db.exec("UPDATE marketplace_listings SET kind='self' WHERE status='private'");
      db.exec("UPDATE marketplace_listings SET kind='market' WHERE status<>'private'");
      db.exec('CREATE INDEX IF NOT EXISTS idx_ml_kind ON marketplace_listings(kind)');
    }

    // 购买记录表
    db.exec(`CREATE TABLE IF NOT EXISTS purchase_records (
      id INTEGER PRIMARY KEY AUTOINCREMENT,
      listing_id INTEGER NOT NULL,
      buyer_user_id INTEGER NOT NULL,
      seller_user_id INTEGER NOT NULL,
      price INTEGER NOT NULL,
      created_at INTEGER NOT NULL
    )`);
    db.exec('CREATE INDEX IF NOT EXISTS idx_pr_listing ON purchase_records(listing_id)');
    db.exec('CREATE INDEX IF NOT EXISTS idx_pr_buyer ON purchase_records(buyer_user_id)');

    // Token 日用量表(热力图用:按天汇总每个用户的 AI Token 消耗)
    db.exec(`CREATE TABLE IF NOT EXISTS token_usage (
      user_id INTEGER NOT NULL,
      day TEXT NOT NULL,
      tokens INTEGER NOT NULL DEFAULT 0,
      PRIMARY KEY (user_id, day)
    )`);
    db.exec('CREATE INDEX IF NOT EXISTS idx_token_usage_day ON token_usage(day)');

    // 积分变动明细(会员/后台查看积分流水)
    db.exec(`CREATE TABLE IF NOT EXISTS points_log (
      id INTEGER PRIMARY KEY AUTOINCREMENT,
      user_id INTEGER NOT NULL,
      username TEXT DEFAULT '',
      change INTEGER NOT NULL,
      balance INTEGER NOT NULL DEFAULT 0,
      type TEXT DEFAULT 'earn',
      reason TEXT DEFAULT '',
      created_at INTEGER NOT NULL
    )`);
    db.exec('CREATE INDEX IF NOT EXISTS idx_points_log_user ON points_log(user_id)');
    db.exec('CREATE INDEX IF NOT EXISTS idx_points_log_created ON points_log(created_at)');

    // 图片生成明细
    db.exec(`CREATE TABLE IF NOT EXISTS image_gen_log (
      id INTEGER PRIMARY KEY AUTOINCREMENT,
      user_id INTEGER NOT NULL,
      username TEXT DEFAULT '',
      prompt TEXT DEFAULT '',
      model TEXT DEFAULT '',
      size TEXT DEFAULT '',
      count INTEGER NOT NULL DEFAULT 1,
      created_at INTEGER NOT NULL
    )`);
    db.exec('CREATE INDEX IF NOT EXISTS idx_image_gen_log_user ON image_gen_log(user_id)');
    db.exec('CREATE INDEX IF NOT EXISTS idx_image_gen_log_created ON image_gen_log(created_at)');

    db.exec('COMMIT');
  } catch (e) {
    try { db.exec('ROLLBACK'); } catch (_) {}
    console.error('[db] runSchema 事务失败,回退逐条执行:', e.message);
    for (const stmt of statements) { try { db.exec(stmt); } catch (_) {} }
  }
}

/** 获取连接(此处即句柄本身) */
async function getPool() {
  return raw();
}

/** 测试 / 打开数据库 */
async function testConnect() {
  raw();
  return true;
}

/** 确保数据库已就绪(SQLite 为文件型,打开即创建) */
async function ensureDatabase() {
  raw();
  return true;
}

/** 查询单行 */
async function row(sql, params = []) {
  return withRetry(() => {
    const database = raw();
    const r = database.prepare(sql).get(...params);
    return r ? normalize(r) : null;
  }, 'row');
}

/** 查询多行 */
async function rows(sql, params = []) {
  return withRetry(() => {
    const database = raw();
    return normalize(database.prepare(sql).all(...params));
  }, 'rows');
}

/** 执行写操作 */
async function exec(sql, params = []) {
  return withRetry(() => {
    const database = raw();
    return database.prepare(sql).run(...params);
  }, 'exec');
}

/** 获取最后插入 ID */
async function lastInsertId() {
  const database = raw();
  const r = database.prepare('SELECT last_insert_rowid() AS id').get();
  return normalize(r).id;
}

/** 将 WAL 中未提交的数据合并回主库文件(TRUNCATE 后 WAL 清空) */
function checkpoint() {
  if (!db) return;
  try { db.exec('PRAGMA wal_checkpoint(TRUNCATE)'); } catch (_) {}
}

/** 关闭并重置当前连接,使下次 raw() 重新打开数据库文件(用于导入替换) */
function close() {
  if (!db) return;
  try { db.exec('PRAGMA wal_checkpoint(TRUNCATE)'); } catch (_) {}
  try { db.close(); } catch (_) {}
  db = null;
}

/** 返回数据库文件路径 */
function getDbPath() {
  return config.db.sqlite_path;
}

/** 记录积分变动明细(非关键路径,失败忽略) */
async function logPoints(userId, change, type, reason) {
  try {
    const u = await row('SELECT username, points FROM users WHERE id=?', [userId]);
    if (!u) return;
    await exec(
      'INSERT INTO points_log (user_id, username, change, balance, type, reason, created_at) VALUES (?,?,?,?,?,?,?)',
      [userId, u.username || '', change, u.points || 0, type || 'earn', reason || '', Math.floor(Date.now() / 1000)]
    );
  } catch (e) { console.error('[db.logPoints]', e.message); }
}

/** 记录图片生成明细(非关键路径,失败忽略) */
async function logImageGen(userId, prompt, model, size, count) {
  try {
    const u = await row('SELECT username FROM users WHERE id=?', [userId]);
    if (!u) return;
    await exec(
      'INSERT INTO image_gen_log (user_id, username, prompt, model, size, count, created_at) VALUES (?,?,?,?,?,?,?)',
      [userId, u.username || '', String(prompt || '').slice(0, 500), model || '', size || '', count || 1, Math.floor(Date.now() / 1000)]
    );
  } catch (e) { console.error('[db.logImageGen]', e.message); }
}

module.exports = { getPool, testConnect, ensureDatabase, row, rows, exec, lastInsertId, checkpoint, close, getDbPath, logPoints, logImageGen };