PromdataPanel/server/db-integrity-check.js

/**
 * Database Integrity Check
 * Runs at startup to ensure all required tables exist.
 * Recreates the database if any tables are missing.
 */
require('dotenv').config();
const mysql = require('mysql2/promise');
const db = require('./db');
const path = require('path');
const fs = require('fs');

const REQUIRED_TABLES = [
  'users',
  'prometheus_sources',
  'site_settings',
  'traffic_stats',
  'server_locations'
];

async function checkAndFixDatabase() {
  const envPath = path.join(__dirname, '..', '.env');
  if (!fs.existsSync(envPath)) return;

  try {
    // Check tables
    const [rows] = await db.query("SHOW TABLES");
    const existingTables = rows.map(r => Object.values(r)[0]);
    
    const missingTables = REQUIRED_TABLES.filter(t => !existingTables.includes(t));
    
    if (missingTables.length > 0) {
      console.log(`[Database Integrity] ⚠️ Missing tables: ${missingTables.join(', ')}. Creating them...`);
      
      for (const table of missingTables) {
        await createTable(table);
      }
      console.log(`[Database Integrity] ✅ Missing tables created.`);
    }

    // Check for is_server_source in prometheus_sources
    const [promColumns] = await db.query("SHOW COLUMNS FROM prometheus_sources");
    const promColumnNames = promColumns.map(c => c.Field);
    if (!promColumnNames.includes('is_server_source')) {
      console.log(`[Database Integrity] ⚠️ Missing column 'is_server_source' in 'prometheus_sources'. Adding it...`);
      await db.query("ALTER TABLE prometheus_sources ADD COLUMN is_server_source TINYINT(1) DEFAULT 1 AFTER description");
      console.log(`[Database Integrity] ✅ Column 'is_server_source' added.`);
    }

    // Check for new columns in site_settings
    const [columns] = await db.query("SHOW COLUMNS FROM site_settings");
    const columnNames = columns.map(c => c.Field);
    if (!columnNames.includes('show_95_bandwidth')) {
      console.log(`[Database Integrity] ⚠️ Missing column 'show_95_bandwidth' in 'site_settings'. Adding it...`);
      await db.query("ALTER TABLE site_settings ADD COLUMN show_95_bandwidth TINYINT(1) DEFAULT 0 AFTER default_theme");
      console.log(`[Database Integrity] ✅ Column 'show_95_bandwidth' added.`);
    }
    if (!columnNames.includes('p95_type')) {
      console.log(`[Database Integrity] ⚠️ Missing column 'p95_type' in 'site_settings'. Adding it...`);
      await db.query("ALTER TABLE site_settings ADD COLUMN p95_type VARCHAR(20) DEFAULT 'tx' AFTER show_95_bandwidth");
      console.log(`[Database Integrity] ✅ Column 'p95_type' added.`);
    }
    if (!columnNames.includes('blackbox_source_id')) {
      console.log(`[Database Integrity] ⚠️ Missing column 'blackbox_source_id' in 'site_settings'. Adding it...`);
      await db.query("ALTER TABLE site_settings ADD COLUMN blackbox_source_id INT AFTER p95_type");
      console.log(`[Database Integrity] ✅ Column 'blackbox_source_id' added.`);
    }
    if (!columnNames.includes('latency_source')) {
      console.log(`[Database Integrity] ⚠️ Missing column 'latency_source' in 'site_settings'. Adding it...`);
      await db.query("ALTER TABLE site_settings ADD COLUMN latency_source VARCHAR(100) AFTER blackbox_source_id");
      console.log(`[Database Integrity] ✅ Column 'latency_source' added.`);
    }
    if (!columnNames.includes('latency_dest')) {
      console.log(`[Database Integrity] ⚠️ Missing column 'latency_dest' in 'site_settings'. Adding it...`);
      await db.query("ALTER TABLE site_settings ADD COLUMN latency_dest VARCHAR(100) AFTER latency_source");
      console.log(`[Database Integrity] ✅ Column 'latency_dest' added.`);
    }
    if (!columnNames.includes('latency_target')) {
      console.log(`[Database Integrity] ⚠️ Missing column 'latency_target' in 'site_settings'. Adding it...`);
      await db.query("ALTER TABLE site_settings ADD COLUMN latency_target VARCHAR(255) AFTER latency_dest");
      console.log(`[Database Integrity] ✅ Column 'latency_target' added.`);
    }
  } catch (err) {
    console.error('[Database Integrity] ❌ Error checking integrity:', err.message);
  }
}

async function createTable(tableName) {
  console.log(`  - Creating table "${tableName}"...`);
  switch (tableName) {
    case 'users':
      await db.query(`
        CREATE TABLE IF NOT EXISTS users (
          id INT AUTO_INCREMENT PRIMARY KEY,
          username VARCHAR(255) NOT NULL UNIQUE,
          password VARCHAR(255) NOT NULL,
          salt VARCHAR(255) NOT NULL,
          created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
      `);
      break;
    case 'prometheus_sources':
      await db.query(`
        CREATE TABLE IF NOT EXISTS prometheus_sources (
          id INT AUTO_INCREMENT PRIMARY KEY,
          name VARCHAR(255) NOT NULL,
          url VARCHAR(500) NOT NULL,
          description TEXT,
          created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
          updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP
        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
      `);
      break;
    case 'site_settings':
      await db.query(`
        CREATE TABLE IF NOT EXISTS site_settings (
          id INT PRIMARY KEY DEFAULT 1,
          page_name VARCHAR(255) DEFAULT '数据可视化展示大屏',
          title VARCHAR(255) DEFAULT '数据可视化展示大屏',
          logo_url TEXT,
          default_theme VARCHAR(20) DEFAULT 'dark',
          show_95_bandwidth TINYINT(1) DEFAULT 0,
          p95_type VARCHAR(20) DEFAULT 'tx',
          blackbox_source_id INT,
          latency_source VARCHAR(100),
          latency_dest VARCHAR(100),
          latency_target VARCHAR(255),
          updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP
        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
      `);
      await db.query(`
        INSERT IGNORE INTO site_settings (id, page_name, title, default_theme, show_95_bandwidth) 
        VALUES (1, '数据可视化展示大屏', '数据可视化展示大屏', 'dark', 0)
      `);
      break;
    case 'traffic_stats':
      await db.query(`
        CREATE TABLE IF NOT EXISTS traffic_stats (
          id INT AUTO_INCREMENT PRIMARY KEY,
          rx_bytes BIGINT UNSIGNED DEFAULT 0,
          tx_bytes BIGINT UNSIGNED DEFAULT 0,
          rx_bandwidth DOUBLE DEFAULT 0,
          tx_bandwidth DOUBLE DEFAULT 0,
          timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
          UNIQUE INDEX (timestamp)
        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
      `);
      break;
    case 'server_locations':
      await db.query(`
        CREATE TABLE IF NOT EXISTS server_locations (
          id INT AUTO_INCREMENT PRIMARY KEY,
          ip VARCHAR(255) NOT NULL UNIQUE,
          country CHAR(2),
          country_name VARCHAR(100),
          region VARCHAR(100),
          city VARCHAR(100),
          latitude DOUBLE,
          longitude DOUBLE,
          last_updated TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP
        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
      `);
      break;
  }
}

module.exports = checkAndFixDatabase;