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() {
  // Only run if .env is already configured
  const envPath = path.join(__dirname, '..', '.env');
  if (!fs.existsSync(envPath)) return;

  const dbHost = process.env.MYSQL_HOST || 'localhost';
  const dbUser = process.env.MYSQL_USER || 'root';
  const dbPass = process.env.MYSQL_PASSWORD || '';
  const dbPort = parseInt(process.env.MYSQL_PORT) || 3306;
  const dbName = process.env.MYSQL_DATABASE || 'display_wall';

  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(', ')}`);
      await recreateDatabase(dbHost, dbPort, dbUser, dbPass, dbName);
    } else {
      // console.log(`[Database Integrity] ✅ All tables accounted for.`);
    }
  } catch (err) {
    if (err.code === 'ER_BAD_DB_ERROR') {
      console.log(`[Database Integrity] ⚠️ Database "${dbName}" does not exist.`);
      await recreateDatabase(dbHost, dbPort, dbUser, dbPass, dbName);
    } else {
      console.error('[Database Integrity] ❌ Error checking integrity:', err.message);
    }
  }
}

async function recreateDatabase(host, port, user, password, dbName) {
  console.log(`[Database Integrity] 🔄 Re-initializing database "${dbName}"...`);
  
  let connection;
  try {
    connection = await mysql.createConnection({ host, port, user, password });

    // Drop and create database
    await connection.query(`DROP DATABASE IF EXISTS \`${dbName}\``);
    await connection.query(`CREATE DATABASE \`${dbName}\` CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci`);
    await connection.query(`USE \`${dbName}\``);

    // Recreate all tables
    console.log('  - Creating table "users"...');
    await connection.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
    `);

    console.log('  - Creating table "prometheus_sources"...');
    await connection.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
    `);

    console.log('  - Creating table "site_settings"...');
    await connection.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',
        updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP
      ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
    `);
    await connection.query(`
      INSERT INTO site_settings (id, page_name, title, default_theme) 
      VALUES (1, '数据可视化展示大屏', '数据可视化展示大屏', 'dark')
    `);

    console.log('  - Creating table "traffic_stats"...');
    await connection.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
    `);

    console.log('  - Creating table "server_locations"...');
    await connection.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
    `);

    console.log(`[Database Integrity] ✅ Re-initialization complete.`);
    
    // Refresh db pool in the main app context
    db.initPool();

  } catch (err) {
    console.error('[Database Integrity] ❌ Critical failure during re-initialization:', err.message);
  } finally {
    if (connection) await connection.end();
  }
}

module.exports = checkAndFixDatabase;