mengyaconnect/mengyaconnect-backend/db.go

package main

import (
	"database/sql"
	"encoding/json"
	"log"
	"os"
	"path/filepath"
	"strings"

	_ "modernc.org/sqlite"
)

// DB is the global SQLite connection shared by all handlers.
var DB *sql.DB

const schema = `
CREATE TABLE IF NOT EXISTS ssh_profiles (
    id          INTEGER  PRIMARY KEY AUTOINCREMENT,
    name        TEXT     NOT NULL UNIQUE,
    alias       TEXT     NOT NULL,
    host        TEXT     NOT NULL,
    port        INTEGER  NOT NULL DEFAULT 22,
    username    TEXT     NOT NULL,
    password    TEXT     NOT NULL DEFAULT '',
    private_key TEXT     NOT NULL DEFAULT '',
    passphrase  TEXT     NOT NULL DEFAULT '',
    created_at  DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
    updated_at  DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
);

CREATE TABLE IF NOT EXISTS commands (
    id         INTEGER  PRIMARY KEY AUTOINCREMENT,
    alias      TEXT     NOT NULL,
    command    TEXT     NOT NULL,
    sort_order INTEGER  NOT NULL DEFAULT 0,
    created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
);

CREATE TABLE IF NOT EXISTS scripts (
    id         INTEGER  PRIMARY KEY AUTOINCREMENT,
    name       TEXT     NOT NULL UNIQUE,
    content    TEXT     NOT NULL DEFAULT '',
    created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
    updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
);
`

// initDB opens the SQLite database, applies PRAGMAs and creates tables.
func initDB() error {
	if err := os.MkdirAll(dataBasePath(), 0o750); err != nil {
		return err
	}
	db, err := sql.Open("sqlite", dbPath())
	if err != nil {
		return err
	}
	// Single writer; multiple readers are fine via WAL.
	db.SetMaxOpenConns(1)
	if _, err := db.Exec(`PRAGMA journal_mode=WAL; PRAGMA foreign_keys=ON; PRAGMA busy_timeout=5000;`); err != nil {
		return err
	}
	if _, err := db.Exec(schema); err != nil {
		return err
	}
	DB = db
	log.Printf("SQLite database ready: %s", dbPath())
	return nil
}

// ─── One-time migration from legacy file storage ─────────────────────────────

// migrateFromFiles reads legacy JSON/text files and inserts their data into
// SQLite (INSERT OR IGNORE so duplicates are skipped). After migration each
// source file is renamed to *.migrated so the routine is not repeated.
func migrateFromFiles() {
	migrateSSHProfiles()
	migrateCommands()
	migrateScripts()
}

func migrateSSHProfiles() {
	dir := filepath.Join(dataBasePath(), "ssh")
	entries, err := os.ReadDir(dir)
	if err != nil {
		return // directory doesn't exist — nothing to migrate
	}
	migrated := 0
	for _, e := range entries {
		if e.IsDir() || !strings.HasSuffix(e.Name(), ".json") {
			continue
		}
		path := filepath.Join(dir, e.Name())
		raw, err := os.ReadFile(path)
		if err != nil {
			log.Printf("migrate ssh: read %s: %v", path, err)
			continue
		}
		var p SSHProfile
		if err := json.Unmarshal(raw, &p); err != nil {
			log.Printf("migrate ssh: parse %s: %v", path, err)
			continue
		}
		name := strings.TrimSuffix(e.Name(), ".json")
		if p.Port == 0 {
			p.Port = 22
		}
		_, err = DB.Exec(
			`INSERT OR IGNORE INTO ssh_profiles (name,alias,host,port,username,password,private_key,passphrase)
			 VALUES (?,?,?,?,?,?,?,?)`,
			name, p.Alias, p.Host, p.Port, p.Username, p.Password, p.PrivateKey, p.Passphrase,
		)
		if err != nil {
			log.Printf("migrate ssh: insert %s: %v", name, err)
			continue
		}
		_ = os.Rename(path, path+".migrated")
		migrated++
	}
	if migrated > 0 {
		log.Printf("migrate: imported %d SSH profile(s) from %s", migrated, dir)
	}
}

func migrateCommands() {
	path := filepath.Join(dataBasePath(), "command", "command.json")
	raw, err := os.ReadFile(path)
	if err != nil {
		return
	}
	var cmds []Command
	if err := json.Unmarshal(raw, &cmds); err != nil {
		log.Printf("migrate commands: parse %s: %v", path, err)
		return
	}
	// Check if commands table is already populated to avoid duplicating on restart.
	var count int
	_ = DB.QueryRow(`SELECT COUNT(*) FROM commands`).Scan(&count)
	if count > 0 {
		log.Printf("migrate commands: table not empty, skipping file migration")
		_ = os.Rename(path, path+".migrated")
		return
	}
	for i, cmd := range cmds {
		if _, err := DB.Exec(
			`INSERT INTO commands (alias, command, sort_order) VALUES (?,?,?)`,
			cmd.Alias, cmd.Command, i,
		); err != nil {
			log.Printf("migrate commands: insert [%d]: %v", i, err)
		}
	}
	_ = os.Rename(path, path+".migrated")
	log.Printf("migrate: imported %d command(s) from %s", len(cmds), path)
}

func migrateScripts() {
	dir := filepath.Join(dataBasePath(), "script")
	entries, err := os.ReadDir(dir)
	if err != nil {
		return
	}
	migrated := 0
	for _, e := range entries {
		if e.IsDir() {
			continue
		}
		path := filepath.Join(dir, e.Name())
		content, err := os.ReadFile(path)
		if err != nil {
			log.Printf("migrate scripts: read %s: %v", path, err)
			continue
		}
		if _, err := DB.Exec(
			`INSERT OR IGNORE INTO scripts (name, content) VALUES (?,?)`,
			e.Name(), string(content),
		); err != nil {
			log.Printf("migrate scripts: insert %s: %v", e.Name(), err)
			continue
		}
		_ = os.Rename(path, path+".migrated")
		migrated++
	}
	if migrated > 0 {
		log.Printf("migrate: imported %d script(s) from %s", migrated, dir)
	}
}