feat: persist device offline lifecycle

2026-09-30 13:17:19 +08:00
parent bb6a446cad
commit 723d7345a6
9 changed files with 314 additions and 82 deletions
@@ -0,0 +1,51 @@
+package storage
+
+import (
+	"path/filepath"
+	"testing"
+)
+
+func TestDeviceLifecyclePersistence(t *testing.T) {
+	s, err := Open(filepath.Join(t.TempDir(), "commander.db"))
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer s.DB.Close()
+
+	if err := s.UpsertRemote("node1", "online", "amd64", "linux", 100, "{\"role\":\"APPLICATION\"}"); err != nil {
+		t.Fatal(err)
+	}
+	rows, err := s.ListDevices()
+	if err != nil || len(rows) != 1 {
+		t.Fatalf("rows=%d err=%v", len(rows), err)
+	}
+	if rows[0].Status != "online" || rows[0].Arch != "amd64" || rows[0].OS != "linux" {
+		t.Fatalf("unexpected row: %+v", rows[0])
+	}
+
+	if err := s.UpdateDeviceStatus("node1", "offline", 120); err != nil {
+		t.Fatal(err)
+	}
+	rows, err = s.ListDevices()
+	if err != nil || len(rows) != 1 {
+		t.Fatalf("rows=%d err=%v", len(rows), err)
+	}
+	if rows[0].Status != "offline" || rows[0].OfflineAt == nil {
+		t.Fatalf("offline state not persisted: %+v", rows[0])
+	}
+
+	if err := s.RevokeDevice("node1", "test"); err != nil {
+		t.Fatal(err)
+	}
+	rows, err = s.ListDevices()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(rows) != 0 {
+		t.Fatalf("revoked device still present: %+v", rows)
+	}
+	revoked, err := s.IsRevoked("node1")
+	if err != nil || !revoked {
+		t.Fatalf("revoked=%v err=%v", revoked, err)
+	}
+}
@@ -18,8 +18,6 @@ func Open(path string) (*Store, error) {
 	if err != nil {
 		return nil, err
 	}
-	// SQLite has a single-writer model. Keep one application connection so
-	// concurrent callers queue instead of opening competing writer connections.
 	db.SetMaxOpenConns(1)
 	db.SetMaxIdleConns(1)
 	s := &Store{db}
@@ -31,28 +29,144 @@ func Open(path string) (*Store, error) {
 }

 func (s *Store) init() error {
-	if _, err := s.DB.Exec(`PRAGMA journal_mode = WAL`); err != nil {
+	if _, err := s.DB.Exec("PRAGMA journal_mode = WAL"); err != nil {
 		return err
 	}
-	if _, err := s.DB.Exec(`PRAGMA busy_timeout = 5000`); err != nil {
+	if _, err := s.DB.Exec("PRAGMA busy_timeout = 5000"); err != nil {
 		return err
 	}
-	_, err := s.DB.Exec(`CREATE TABLE IF NOT EXISTS devices (
-		name TEXT PRIMARY KEY, status TEXT NOT NULL, is_local INTEGER NOT NULL,
-		arch TEXT, os TEXT, updated_at INTEGER NOT NULL
-	);	CREATE TABLE IF NOT EXISTS audit_logs (
-		id INTEGER PRIMARY KEY AUTOINCREMENT, timestamp INTEGER NOT NULL,
-		caller_ip TEXT, device TEXT, action TEXT, command TEXT,
-		exit_code INTEGER, duration_ms INTEGER, error_msg TEXT
-	);`)
+	_, err := s.DB.Exec("CREATE TABLE IF NOT EXISTS devices (" +
+		"name TEXT PRIMARY KEY, status TEXT NOT NULL DEFAULT 'offline', is_local INTEGER NOT NULL DEFAULT 0, " +
+		"arch TEXT, os TEXT, last_seen INTEGER NOT NULL DEFAULT 0, last_heartbeat INTEGER NOT NULL DEFAULT 0, " +
+		"offline_at INTEGER, profile_json TEXT, created_at INTEGER NOT NULL DEFAULT 0, updated_at INTEGER NOT NULL DEFAULT 0); " +
+		"CREATE TABLE IF NOT EXISTS audit_logs (" +
+		"id INTEGER PRIMARY KEY AUTOINCREMENT, timestamp INTEGER NOT NULL, caller_ip TEXT, device TEXT, action TEXT, command TEXT, " +
+		"exit_code INTEGER, duration_ms INTEGER, error_msg TEXT); " +
+		"CREATE TABLE IF NOT EXISTS revoked_devices (name TEXT PRIMARY KEY, revoked_at INTEGER NOT NULL, reason TEXT);")
+	if err != nil {
+		return err
+	}
+	return s.migrateDevices()
+}
+
+func (s *Store) migrateDevices() error {
+	cols := map[string]bool{}
+	rows, err := s.DB.Query("PRAGMA table_info(devices)")
+	if err != nil {
+		return err
+	}
+	defer rows.Close()
+	for rows.Next() {
+		var cid, notnull, pk int
+		var name, typ string
+		var dflt sql.NullString
+		if err := rows.Scan(&cid, &name, &typ, &notnull, &dflt, &pk); err != nil {
+			return err
+		}
+		cols[name] = true
+	}
+	if err := rows.Err(); err != nil {
+		return err
+	}
+	migrations := []struct{ name, sql string }{
+		{"last_seen", "ALTER TABLE devices ADD COLUMN last_seen INTEGER NOT NULL DEFAULT 0"},
+		{"last_heartbeat", "ALTER TABLE devices ADD COLUMN last_heartbeat INTEGER NOT NULL DEFAULT 0"},
+		{"offline_at", "ALTER TABLE devices ADD COLUMN offline_at INTEGER"},
+		{"profile_json", "ALTER TABLE devices ADD COLUMN profile_json TEXT"},
+		{"created_at", "ALTER TABLE devices ADD COLUMN created_at INTEGER NOT NULL DEFAULT 0"},
+	}
+	for _, m := range migrations {
+		if !cols[m.name] {
+			if _, err := s.DB.Exec(m.sql); err != nil {
+				return err
+			}
+		}
+	}
+	_, err = s.DB.Exec("UPDATE devices SET created_at=CASE WHEN created_at=0 THEN updated_at ELSE created_at END, " +
+		"last_seen=CASE WHEN last_seen=0 THEN updated_at ELSE last_seen END, " +
+		"last_heartbeat=CASE WHEN last_heartbeat=0 THEN updated_at ELSE last_heartbeat END " +
+		"WHERE created_at=0 OR last_seen=0 OR last_heartbeat=0")
 	return err
 }

 func (s *Store) UpsertLocal(name, arch, osName string) error {
-	_, err := s.DB.Exec(`INSERT INTO devices(name,status,is_local,arch,os,updated_at)
-		VALUES(?,?,?,?,?,?) ON CONFLICT(name) DO UPDATE SET status=excluded.status,
-		arch=excluded.arch, os=excluded.os, updated_at=excluded.updated_at`,
-		name, "online", 1, arch, osName, time.Now().Unix())
+	now := time.Now().Unix()
+	_, err := s.DB.Exec("INSERT INTO devices(name,status,is_local,arch,os,last_seen,last_heartbeat,offline_at,profile_json,created_at,updated_at) "+
+		"VALUES(?,?,?,?,?,?,?,?,?,?,?) ON CONFLICT(name) DO UPDATE SET status='online', is_local=1, arch=excluded.arch, os=excluded.os, offline_at=NULL, updated_at=excluded.updated_at",
+		name, "online", 1, arch, osName, now, now, nil, nil, now, now)
+	return err
+}
+
+func (s *Store) UpsertRemote(name, status, arch, osName string, lastHeartbeat int64, profileJSON string) error {
+	now := time.Now().Unix()
+	if lastHeartbeat <= 0 {
+		lastHeartbeat = now
+	}
+	_, err := s.DB.Exec("INSERT INTO devices(name,status,is_local,arch,os,last_seen,last_heartbeat,offline_at,profile_json,created_at,updated_at) "+
+		"VALUES(?,?,?,?,?,?,?,?,?,?,?) ON CONFLICT(name) DO UPDATE SET status=excluded.status, is_local=0, arch=excluded.arch, os=excluded.os, "+
+		"last_seen=excluded.last_seen, last_heartbeat=excluded.last_heartbeat, offline_at=NULL, profile_json=excluded.profile_json, updated_at=excluded.updated_at",
+		name, status, 0, arch, osName, lastHeartbeat, lastHeartbeat, nil, profileJSON, now, now)
+	return err
+}
+
+func (s *Store) UpdateRemoteProfile(name string, lastHeartbeat int64, arch, osName, profileJSON string) error {
+	now := time.Now().Unix()
+	if lastHeartbeat <= 0 {
+		lastHeartbeat = now
+	}
+	_, err := s.DB.Exec("UPDATE devices SET status='online', arch=?, os=?, last_seen=?, last_heartbeat=?, offline_at=NULL, profile_json=?, updated_at=? WHERE name=? AND is_local=0",
+		arch, osName, lastHeartbeat, lastHeartbeat, profileJSON, now, name)
+	return err
+}
+
+func (s *Store) UpdateDeviceStatus(name, status string, lastHeartbeat int64) error {
+	now := time.Now().Unix()
+	if lastHeartbeat <= 0 {
+		lastHeartbeat = now
+	}
+	var offlineAt any
+	if status == "offline" {
+		offlineAt = now
+	}
+	_, err := s.DB.Exec("UPDATE devices SET status=?, last_seen=?, last_heartbeat=?, offline_at=?, updated_at=? WHERE name=? AND is_local=0",
+		status, lastHeartbeat, lastHeartbeat, offlineAt, now, name)
+	return err
+}
+
+type DeviceRecord struct {
+	Name, Status, Arch, OS, ProfileJSON string
+	IsLocal                             bool
+	LastSeen, LastHeartbeat             int64
+	OfflineAt                           *int64
+	UpdatedAt                           int64
+}
+
+func (s *Store) ListDevices() ([]DeviceRecord, error) {
+	rows, err := s.DB.Query("SELECT name,status,is_local,COALESCE(arch,''),COALESCE(os,''),last_seen,last_heartbeat,offline_at,COALESCE(profile_json,''),updated_at FROM devices ORDER BY name")
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+	var out []DeviceRecord
+	for rows.Next() {
+		var d DeviceRecord
+		var local int
+		var offline sql.NullInt64
+		if err := rows.Scan(&d.Name, &d.Status, &local, &d.Arch, &d.OS, &d.LastSeen, &d.LastHeartbeat, &offline, &d.ProfileJSON, &d.UpdatedAt); err != nil {
+			return nil, err
+		}
+		d.IsLocal = local != 0
+		if offline.Valid {
+			v := offline.Int64
+			d.OfflineAt = &v
+		}
+		out = append(out, d)
+	}
+	return out, rows.Err()
+}
+
+func (s *Store) DeleteDevice(name string) error {
+	_, err := s.DB.Exec("DELETE FROM devices WHERE name=? AND is_local=0", name)
 	return err
 }

@@ -69,8 +183,8 @@ type AuditLog struct {
 }

 func (s *Store) Audit(ip, device, action, command string, code int, ms int64, msg string) error {
-	_, err := s.DB.Exec(`INSERT INTO audit_logs(timestamp,caller_ip,device,action,command,exit_code,duration_ms,error_msg)
-		VALUES(?,?,?,?,?,?,?,?)`, time.Now().Unix(), ip, device, action, command, code, ms, msg)
+	_, err := s.DB.Exec("INSERT INTO audit_logs(timestamp,caller_ip,device,action,command,exit_code,duration_ms,error_msg) VALUES(?,?,?,?,?,?,?,?)",
+		time.Now().Unix(), ip, device, action, command, code, ms, msg)
 	return err
 }

@@ -81,8 +195,7 @@ func (s *Store) RecentAudits(limit int) ([]AuditLog, error) {
 	if limit > 500 {
 		limit = 500
 	}
-	rows, err := s.DB.Query(`SELECT id,timestamp,caller_ip,device,action,command,exit_code,duration_ms,error_msg
-		FROM audit_logs ORDER BY id DESC LIMIT ?`, limit)
+	rows, err := s.DB.Query("SELECT id,timestamp,caller_ip,device,action,command,exit_code,duration_ms,error_msg FROM audit_logs ORDER BY id DESC LIMIT ?", limit)
 	if err != nil {
 		return nil, err
 	}
@@ -98,29 +211,29 @@ func (s *Store) RecentAudits(limit int) ([]AuditLog, error) {
 	return out, rows.Err()
 }

-func (s *Store) InitRevocationTable() error {
-	_, err := s.DB.Exec(`CREATE TABLE IF NOT EXISTS revoked_devices (
-		name TEXT PRIMARY KEY,
-		revoked_at INTEGER NOT NULL,
-		reason TEXT
-	);`)
-	return err
-}
-
 func (s *Store) RevokeDevice(name, reason string) error {
-	_, err := s.DB.Exec(`INSERT INTO revoked_devices(name, revoked_at, reason)
-		VALUES(?, ?, ?) ON CONFLICT(name) DO UPDATE SET revoked_at=excluded.revoked_at, reason=excluded.reason`,
-		name, time.Now().Unix(), reason)
-	return err
+	tx, err := s.DB.Begin()
+	if err != nil {
+		return err
+	}
+	defer tx.Rollback()
+	if _, err := tx.Exec("INSERT INTO revoked_devices(name,revoked_at,reason) VALUES(?,?,?) ON CONFLICT(name) DO UPDATE SET revoked_at=excluded.revoked_at,reason=excluded.reason",
+		name, time.Now().Unix(), reason); err != nil {
+		return err
+	}
+	if _, err := tx.Exec("DELETE FROM devices WHERE name=? AND is_local=0", name); err != nil {
+		return err
+	}
+	return tx.Commit()
 }

 func (s *Store) UnrevokeDevice(name string) error {
-	_, err := s.DB.Exec(`DELETE FROM revoked_devices WHERE name = ?`, name)
+	_, err := s.DB.Exec("DELETE FROM revoked_devices WHERE name = ?", name)
 	return err
 }

 func (s *Store) IsRevoked(name string) (bool, error) {
 	var count int
-	err := s.DB.QueryRow(`SELECT COUNT(*) FROM revoked_devices WHERE name = ?`, name).Scan(&count)
+	err := s.DB.QueryRow("SELECT COUNT(*) FROM revoked_devices WHERE name = ?", name).Scan(&count)
 	return count > 0, err
 }