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',
  'latency_routes'
];

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 and type 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.`);
    }

    if (!promColumnNames.includes('type')) {
      console.log(`[Database Integrity] ⚠️ Missing column 'type' in 'prometheus_sources'. Adding it...`);
      await db.query("ALTER TABLE prometheus_sources ADD COLUMN type VARCHAR(50) DEFAULT 'prometheus' AFTER is_server_source");
      console.log(`[Database Integrity] ✅ Column 'type' 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);
    const addColumn = async (columnName, sql) => {
      if (!columnNames.includes(columnName)) {
        try {
          console.log(`[Database Integrity] ⚠️ Missing column '${columnName}' in 'site_settings'. Adding it...`);
          await db.query(sql);
          console.log(`[Database Integrity] ✅ Column '${columnName}' added.`);
        } catch (err) {
          console.error(`[Database Integrity] ❌ Failed to add column '${columnName}':`, err.message);
          // Try without AFTER if it exists
          if (sql.includes('AFTER')) {
            try {
              const fallback = sql.split(' AFTER')[0];
              console.log(`[Database Integrity] 🔄 Retrying column '${columnName}' WITHOUT 'AFTER'...`);
              await db.query(fallback);
              console.log(`[Database Integrity] ✅ Column '${columnName}' added via fallback.`);
            } catch (err2) {
              console.error(`[Database Integrity] ❌ Fallback also failed:`, err2.message);
            }
          }
        }
      }
    };

    await addColumn('show_page_name', "ALTER TABLE site_settings ADD COLUMN show_page_name TINYINT(1) DEFAULT 1 AFTER page_name");
    await addColumn('show_95_bandwidth', "ALTER TABLE site_settings ADD COLUMN show_95_bandwidth TINYINT(1) DEFAULT 0 AFTER default_theme");
    await addColumn('p95_type', "ALTER TABLE site_settings ADD COLUMN p95_type VARCHAR(20) DEFAULT 'tx' AFTER show_95_bandwidth");
    await addColumn('require_login_for_server_details', "ALTER TABLE site_settings ADD COLUMN require_login_for_server_details TINYINT(1) DEFAULT 1 AFTER p95_type");
    await addColumn('blackbox_source_id', "ALTER TABLE site_settings ADD COLUMN blackbox_source_id INT AFTER p95_type");
    await addColumn('latency_source', "ALTER TABLE site_settings ADD COLUMN latency_source VARCHAR(100) AFTER blackbox_source_id");
    await addColumn('latency_dest', "ALTER TABLE site_settings ADD COLUMN latency_dest VARCHAR(100) AFTER latency_source");
    await addColumn('latency_target', "ALTER TABLE site_settings ADD COLUMN latency_target VARCHAR(255) AFTER latency_dest");
    await addColumn('icp_filing', "ALTER TABLE site_settings ADD COLUMN icp_filing VARCHAR(255) AFTER latency_target");
    await addColumn('ps_filing', "ALTER TABLE site_settings ADD COLUMN ps_filing VARCHAR(255) AFTER icp_filing");
    await addColumn('logo_url_dark', "ALTER TABLE site_settings ADD COLUMN logo_url_dark TEXT AFTER logo_url");
    await addColumn('favicon_url', "ALTER TABLE site_settings ADD COLUMN favicon_url TEXT AFTER logo_url_dark");
  } catch (err) {
    console.error('[Database Integrity] ❌ Overall site_settings check error:', 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 '数据可视化展示大屏',
          show_page_name TINYINT(1) DEFAULT 1,
          title VARCHAR(255) DEFAULT '数据可视化展示大屏',
          logo_url TEXT,
          logo_url_dark TEXT,
          favicon_url TEXT,
          default_theme VARCHAR(20) DEFAULT 'dark',
          show_95_bandwidth TINYINT(1) DEFAULT 0,
          p95_type VARCHAR(20) DEFAULT 'tx',
          require_login_for_server_details TINYINT(1) DEFAULT 1,
          blackbox_source_id INT,
          latency_source VARCHAR(100),
          latency_dest VARCHAR(100),
          latency_target VARCHAR(255),
          icp_filing VARCHAR(255),
          ps_filing 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 'latency_routes':
      await db.query(`
        CREATE TABLE IF NOT EXISTS latency_routes (
          id INT AUTO_INCREMENT PRIMARY KEY,
          source_id INT NOT NULL,
          latency_source VARCHAR(100) NOT NULL,
          latency_dest VARCHAR(100) NOT NULL,
          latency_target VARCHAR(255) NOT NULL,
          created_at TIMESTAMP DEFAULT CURRENT_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;