nano_metrics/server/db/db.go

package db

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"log"
	"time"

	_ "modernc.org/sqlite"

	"github.com/user/nanometrics/server/models"
)

type DB struct {
	conn *sql.DB
}

const schema = `
CREATE TABLE IF NOT EXISTS agents (
	id TEXT PRIMARY KEY, hostname TEXT NOT NULL,
	ip TEXT NOT NULL DEFAULT '', status TEXT NOT NULL DEFAULT 'offline',
	last_seen INTEGER NOT NULL DEFAULT 0
);
CREATE TABLE IF NOT EXISTS metrics (
	id INTEGER PRIMARY KEY AUTOINCREMENT, agent_id TEXT NOT NULL, ts INTEGER NOT NULL,
	cpu_percent REAL, memory_used INTEGER, memory_free INTEGER, memory_total INTEGER,
	hdd_used INTEGER, hdd_free INTEGER, hdd_total INTEGER,
	uptime INTEGER, network_rx INTEGER, network_tx INTEGER, temperature REAL,
	smart_passed INTEGER, smart_temp INTEGER, smart_realloc INTEGER,
	smart_hours INTEGER, smart_wear INTEGER,
	FOREIGN KEY (agent_id) REFERENCES agents(id)
);
CREATE INDEX IF NOT EXISTS idx_metrics_agent_ts ON metrics(agent_id, ts);
CREATE TABLE IF NOT EXISTS agent_configs (
	agent_id TEXT PRIMARY KEY, config_json TEXT NOT NULL DEFAULT '{}',
	FOREIGN KEY (agent_id) REFERENCES agents(id)
);
CREATE TABLE IF NOT EXISTS agent_icons (
	agent_id TEXT PRIMARY KEY, data BLOB NOT NULL, mime_type TEXT NOT NULL DEFAULT 'image/png',
	FOREIGN KEY (agent_id) REFERENCES agents(id)
);
CREATE TABLE IF NOT EXISTS server_config (key TEXT PRIMARY KEY, value TEXT NOT NULL);
`

func Open(path string) (*DB, error) {
	conn, err := sql.Open("sqlite", path)
	if err != nil {
		return nil, fmt.Errorf("open sqlite: %w", err)
	}
	conn.SetMaxOpenConns(1)
	d := &DB{conn: conn}
	if err := d.migrate(); err != nil {
		return nil, fmt.Errorf("migrate: %w", err)
	}
	return d, nil
}

func (d *DB) migrate() error {
	_, err := d.conn.Exec(schema)
	return err
}

func (d *DB) Close() { _ = d.conn.Close() }

func (d *DB) UpsertAgent(m *models.AgentMetrics) error {
	ts := time.Now().Unix()
	_, err := d.conn.Exec(`
		INSERT INTO agents (id, hostname, ip, status, last_seen)
		VALUES (?, ?, ?, ?, ?)
		ON CONFLICT(id) DO UPDATE SET
			ip=excluded.ip, status=excluded.status, last_seen=excluded.last_seen`,
		m.Hostname, m.Hostname, m.IP, m.Status, ts)
	return err
}

func (d *DB) InsertMetrics(m *models.AgentMetrics) error {
	ts := time.Now().Unix()
	var smartPassed, smartTemp, smartRealloc, smartHours, smartWear interface{}
	if m.Smart != nil {
		b := 0
		if m.Smart.Passed {
			b = 1
		}
		smartPassed = b
		smartTemp = m.Smart.Temperature
		smartRealloc = m.Smart.ReallocatedSectors
		smartHours = m.Smart.PowerOnHours
		smartWear = m.Smart.WearLevel
	}
	_, err := d.conn.Exec(`
		INSERT INTO metrics (agent_id, ts,
			cpu_percent, memory_used, memory_free, memory_total,
			hdd_used, hdd_free, hdd_total,
			uptime, network_rx, network_tx, temperature,
			smart_passed, smart_temp, smart_realloc, smart_hours, smart_wear)
		VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)`,
		m.Hostname, ts,
		m.CPUPercent, m.MemoryUsed, m.MemoryFree, m.MemoryTotal,
		m.HDDUsed, m.HDDFree, m.HDDTotal,
		m.Uptime, m.NetworkRX, m.NetworkTX, m.Temperature,
		smartPassed, smartTemp, smartRealloc, smartHours, smartWear)
	return err
}

func (d *DB) GetAgents() ([]models.Agent, error) {
	rows, err := d.conn.Query(`SELECT id, hostname, ip, status, last_seen FROM agents`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var agents []models.Agent
	for rows.Next() {
		var a models.Agent
		if err := rows.Scan(&a.ID, &a.Hostname, &a.IP, &a.Status, &a.LastSeen); err != nil {
			return nil, err
		}
		agents = append(agents, a)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	return agents, nil
}

func (d *DB) GetMetricsHistory(agentID string, from, to int64) ([]map[string]interface{}, error) {
	rows, err := d.conn.Query(`
		SELECT ts, cpu_percent, memory_used, memory_total, hdd_used, hdd_total
		FROM metrics
		WHERE agent_id = ? AND ts >= ? AND ts <= ?
		ORDER BY ts ASC`, agentID, from, to)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var result []map[string]interface{}
	for rows.Next() {
		var ts int64
		var cpu, memUsed, memTotal, hddUsed, hddTotal interface{}
		if err := rows.Scan(&ts, &cpu, &memUsed, &memTotal, &hddUsed, &hddTotal); err != nil {
			return nil, err
		}
		result = append(result, map[string]interface{}{
			"ts": ts, "cpu_percent": cpu,
			"memory_used": memUsed, "memory_total": memTotal,
			"hdd_used": hddUsed, "hdd_total": hddTotal,
		})
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	return result, nil
}

func (d *DB) GetAgentConfig(agentID string) (*models.AgentConfig, error) {
	var raw string
	err := d.conn.QueryRow(
		`SELECT config_json FROM agent_configs WHERE agent_id = ?`, agentID,
	).Scan(&raw)
	if err == sql.ErrNoRows {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}
	var cfg models.AgentConfig
	if err := json.Unmarshal([]byte(raw), &cfg); err != nil {
		return nil, err
	}
	return &cfg, nil
}

func (d *DB) UpsertAgentConfig(agentID string, cfg *models.AgentConfig) error {
	raw, err := json.Marshal(cfg)
	if err != nil {
		return err
	}
	d.conn.Exec(`INSERT OR IGNORE INTO agents (id, hostname, ip, status, last_seen) VALUES (?,?,?,?,?)`,
		agentID, agentID, "", "offline", 0)
	_, err = d.conn.Exec(`
		INSERT INTO agent_configs (agent_id, config_json)
		VALUES (?, ?)
		ON CONFLICT(agent_id) DO UPDATE SET config_json=excluded.config_json`,
		agentID, string(raw))
	return err
}

func (d *DB) GetServerConfig() (models.ServerConfig, error) {
	cfg := models.DefaultServerConfig()
	var raw string
	if err := d.conn.QueryRow(`SELECT value FROM server_config WHERE key='ui'`).Scan(&raw); err == nil {
		_ = json.Unmarshal([]byte(raw), &cfg)
	}
	return cfg, nil
}

func (d *DB) SetServerConfig(cfg models.ServerConfig) error {
	raw, err := json.Marshal(cfg)
	if err != nil {
		return err
	}
	_, err = d.conn.Exec(`
		INSERT INTO server_config (key, value) VALUES ('ui', ?)
		ON CONFLICT(key) DO UPDATE SET value=excluded.value`, string(raw))
	return err
}

func (d *DB) SaveIcon(agentID string, data []byte, mimeType string) error {
	d.conn.Exec(`INSERT OR IGNORE INTO agents (id, hostname, ip, status, last_seen) VALUES (?,?,?,?,?)`,
		agentID, agentID, "", "offline", 0)
	_, err := d.conn.Exec(`
		INSERT INTO agent_icons (agent_id, data, mime_type) VALUES (?,?,?)
		ON CONFLICT(agent_id) DO UPDATE SET data=excluded.data, mime_type=excluded.mime_type`,
		agentID, data, mimeType)
	return err
}

func (d *DB) GetIcon(agentID string) ([]byte, string, error) {
	var data []byte
	var mime string
	err := d.conn.QueryRow(
		`SELECT data, mime_type FROM agent_icons WHERE agent_id=?`, agentID,
	).Scan(&data, &mime)
	if err != nil {
		return nil, "", err
	}
	return data, mime, nil
}

func (d *DB) PruneOldMetrics(retentionDays int) error {
	cutoff := time.Now().Unix() - int64(retentionDays)*86400
	_, err := d.conn.Exec(`DELETE FROM metrics WHERE ts < ?`, cutoff)
	return err
}

func (d *DB) MarkOffline(timeoutSec int64) error {
	_, err := d.MarkOfflineAndGetIDs(timeoutSec)
	return err
}

func (d *DB) DeleteAgent(agentID string) error {
	for _, q := range []string{
		`DELETE FROM metrics WHERE agent_id = ?`,
		`DELETE FROM agent_configs WHERE agent_id = ?`,
		`DELETE FROM agent_icons WHERE agent_id = ?`,
		`DELETE FROM agents WHERE id = ?`,
	} {
		if _, err := d.conn.Exec(q, agentID); err != nil {
			return err
		}
	}
	return nil
}

// MarkOfflineAndGetIDs marque les agents inactifs et retourne leurs IDs.
func (d *DB) MarkOfflineAndGetIDs(timeoutSec int64) ([]string, error) {
	cutoff := time.Now().Unix() - timeoutSec
	rows, err := d.conn.Query(
		`SELECT id FROM agents WHERE last_seen < ? AND status != 'offline'`, cutoff)
	if err != nil {
		return nil, err
	}
	var ids []string
	for rows.Next() {
		var id string
		_ = rows.Scan(&id)
		ids = append(ids, id)
	}
	if err = rows.Err(); err != nil {
		return nil, err
	}
	rows.Close()
	if len(ids) > 0 {
		_, err = d.conn.Exec(
			`UPDATE agents SET status='offline' WHERE last_seen < ? AND status != 'offline'`, cutoff)
	}
	return ids, err
}

func init() {
	log.SetFlags(log.LstdFlags | log.Lshortfile)
}