lan-manager/server/db/db.go

package db

import (
	"database/sql"
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"lan-manager/server/config"

	_ "modernc.org/sqlite"
)

var DB *sql.DB

func Init(cfg *config.Config) error {
	if err := os.MkdirAll(filepath.Dir(cfg.DBPath), 0750); err != nil {
		return err
	}
	var err error
	DB, err = sql.Open("sqlite", cfg.DBPath+"?_pragma=foreign_keys(1)")
	if err != nil {
		return err
	}
	if err := DB.Ping(); err != nil {
		return err
	}
	DB.SetMaxOpenConns(1)
	return migrate()
}

func migrate() error {
	stmts := []string{
		`CREATE TABLE IF NOT EXISTS machines (
			id INTEGER PRIMARY KEY AUTOINCREMENT,
			hostname TEXT NOT NULL,
			ip TEXT NOT NULL UNIQUE,
			mac TEXT,
			os_type TEXT NOT NULL,
			os_version TEXT,
			notes TEXT,
			ssh_port INTEGER DEFAULT 22,
			is_online INTEGER DEFAULT 0,
			last_ping_at DATETIME,
			cpu_info TEXT,
			memory_info TEXT,
			disk_info TEXT,
			uptime TEXT,
			listen_ports TEXT,
			ssh_synced_at DATETIME,
			created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
			updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
		)`,
		`ALTER TABLE machines ADD COLUMN ssh_port INTEGER DEFAULT 22`,
		`ALTER TABLE machines ADD COLUMN ssh_username TEXT`,
		`ALTER TABLE machines ADD COLUMN ssh_password TEXT`,
		`CREATE TABLE IF NOT EXISTS services (
			id INTEGER PRIMARY KEY AUTOINCREMENT,
			machine_id INTEGER NOT NULL REFERENCES machines(id) ON DELETE CASCADE,
			name TEXT NOT NULL,
			port INTEGER NOT NULL,
			protocol TEXT DEFAULT 'TCP',
			notes TEXT,
			target_machine_id INTEGER REFERENCES machines(id) ON DELETE SET NULL,
			target_notes TEXT
		)`,
		`ALTER TABLE services ADD COLUMN target_machine_id INTEGER REFERENCES machines(id) ON DELETE SET NULL`,
		`ALTER TABLE services ADD COLUMN target_notes TEXT`,
		`CREATE TABLE IF NOT EXISTS relationships (
			id INTEGER PRIMARY KEY AUTOINCREMENT,
			source_machine_id INTEGER NOT NULL REFERENCES machines(id) ON DELETE CASCADE,
			target_machine_id INTEGER NOT NULL REFERENCES machines(id) ON DELETE CASCADE,
			relation_type TEXT NOT NULL,
			source_port INTEGER,
			target_port INTEGER,
			notes TEXT
		)`,
		`CREATE TABLE IF NOT EXISTS operation_logs (
			id INTEGER PRIMARY KEY AUTOINCREMENT,
			action TEXT NOT NULL,
			entity_type TEXT NOT NULL,
			entity_id INTEGER,
			entity_name TEXT,
			old_value TEXT,
			new_value TEXT,
			source_ip TEXT,
			username TEXT,
			created_at DATETIME DEFAULT CURRENT_TIMESTAMP
		)`,
		`CREATE INDEX IF NOT EXISTS idx_machines_ip ON machines(ip)`,
		`CREATE INDEX IF NOT EXISTS idx_services_machine ON services(machine_id)`,
		`CREATE INDEX IF NOT EXISTS idx_rel_src ON relationships(source_machine_id)`,
		`CREATE INDEX IF NOT EXISTS idx_rel_tgt ON relationships(target_machine_id)`,
		`CREATE INDEX IF NOT EXISTS idx_logs_created ON operation_logs(created_at)`,
			`ALTER TABLE machines ADD COLUMN offline_count INTEGER DEFAULT 0`,
			`ALTER TABLE machines ADD COLUMN total_offline_seconds INTEGER DEFAULT 0`,
			`ALTER TABLE machines ADD COLUMN last_offline_at DATETIME`,
			`ALTER TABLE machines ADD COLUMN last_offline_reason TEXT`,
			`CREATE TABLE IF NOT EXISTS offline_logs (
				id INTEGER PRIMARY KEY AUTOINCREMENT,
				machine_id INTEGER NOT NULL REFERENCES machines(id) ON DELETE CASCADE,
				reason TEXT,
				started_at DATETIME DEFAULT CURRENT_TIMESTAMP,
				ended_at DATETIME,
				duration_seconds INTEGER
			)`,
			`CREATE INDEX IF NOT EXISTS idx_offline_logs_machine ON offline_logs(machine_id)`,
			`CREATE INDEX IF NOT EXISTS idx_offline_logs_started ON offline_logs(started_at)`,
			`CREATE TABLE IF NOT EXISTS settings (
				key TEXT PRIMARY KEY,
				value TEXT
			)`,
			// PVE 相关表
			`CREATE TABLE IF NOT EXISTS pve_hosts (
				id INTEGER PRIMARY KEY AUTOINCREMENT,
				name TEXT NOT NULL,
				hostname TEXT NOT NULL,
				port INTEGER DEFAULT 8006,
				node_name TEXT DEFAULT 'pve',
				username TEXT NOT NULL,
				password_enc TEXT NOT NULL,
				verify_ssl INTEGER DEFAULT 0,
				created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
				updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
			)`,
			`ALTER TABLE machines ADD COLUMN pve_host_id INTEGER REFERENCES pve_hosts(id) ON DELETE SET NULL`,
			`ALTER TABLE machines ADD COLUMN pve_vmid TEXT`,
			`ALTER TABLE machines ADD COLUMN pve_vm_status TEXT`,
	}
	for _, s := range stmts {
		if _, err := DB.Exec(s); err != nil {
			// ignore "duplicate column" errors for ALTER TABLE compatibility
			if strings.Contains(err.Error(), "duplicate column name") {
				continue
			}
			return fmt.Errorf("migration failed: %w", err)
		}
	}
	return nil
}