bugswatch-agent/internal/dbmanager/manager.go

package dbmanager

import (
	"context"
	"crypto/md5"
	"encoding/json"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"sync"
	"time"

	"bugs.watch/agent/internal/aggregate"
	"bugs.watch/agent/internal/collectors"
	mongocol "bugs.watch/agent/internal/collectors/mongodb"
	mysqlcol "bugs.watch/agent/internal/collectors/mysql"
	pgcol "bugs.watch/agent/internal/collectors/postgres"
	rediscol "bugs.watch/agent/internal/collectors/redis"
	sqlitecol "bugs.watch/agent/internal/collectors/sqlite"
)

// Driver type identifiers. used as the db_driver field in the health payload
// so the dashboard can render the correct icon without parsing strings.
const (
	DriverNone      uint8 = 0
	DriverMySQL     uint8 = 1 // also MariaDB
	DriverPostgres  uint8 = 2
	DriverRedis     uint8 = 3
	DriverMongoDB   uint8 = 4
	DriverSQLite    uint8 = 5
)

// Entry is one database config record in databases.json.
type Entry struct {
	Driver string `json:"driver"`
	URL    string `json:"url"`
}

type dbCollector interface {
	collectors.Collector
	Close()
}

// Manager watches a config file and hot-swaps DB collectors without agent restart.
// Collect() is safe to call concurrently with Watch().
type Manager struct {
	mu         sync.RWMutex
	active     []dbCollector
	configPath string
	lastHash   [16]byte
}

func New(configPath string) *Manager {
	return &Manager{configPath: configPath}
}

// Watch polls the config file every 30 seconds and reloads on change.
// Blocks until ctx is cancelled.
func (m *Manager) Watch(ctx context.Context) {
	m.reload()

	ticker := time.NewTicker(30 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			m.swap(nil)
			return
		case <-ticker.C:
			m.reload()
		}
	}
}

// Collect runs all active DB collectors against the window.
// Holds RLock for the full duration so swaps never race with in-progress queries.
func (m *Manager) Collect(window *aggregate.Window) {
	m.mu.RLock()
	defer m.mu.RUnlock()
	if len(m.active) == 0 {
		return
	}
	window.SetDBDriver(m.driverID(m.active[0].Name()))
	for _, c := range m.active {
		c.Collect(window)
	}
}

func (m *Manager) driverID(name string) uint8 {
	switch name {
	case "mysql", "mariadb":
		return DriverMySQL
	case "postgres":
		return DriverPostgres
	case "redis":
		return DriverRedis
	case "mongodb":
		return DriverMongoDB
	case "sqlite":
		return DriverSQLite
	default:
		return DriverNone
	}
}

func (m *Manager) reload() {
	data, err := os.ReadFile(m.configPath)
	if err != nil {
		if os.IsNotExist(err) {
			m.swap(nil)
		}
		return
	}

	hash := md5.Sum(data)
	m.mu.RLock()
	unchanged := hash == m.lastHash
	m.mu.RUnlock()
	if unchanged {
		return
	}

	var entries []Entry
	if err := json.Unmarshal(data, &entries); err != nil {
		log.Printf("dbmanager: invalid config: %v", err)
		return
	}

	var next []dbCollector
	for _, e := range entries {
		c, err := newCollector(e)
		if err != nil {
			log.Printf("dbmanager: %s: %v", e.Driver, err)
			continue
		}
		next = append(next, c)
		log.Printf("dbmanager: loaded %s", e.Driver)
	}

	m.swap(next)

	m.mu.Lock()
	m.lastHash = hash
	m.mu.Unlock()
}

func (m *Manager) swap(next []dbCollector) {
	m.mu.Lock()
	old := m.active
	m.active = next
	m.mu.Unlock()

	for _, c := range old {
		c.Close()
	}
}

func newCollector(e Entry) (dbCollector, error) {
	switch e.Driver {
	case "mysql", "mariadb":
		return mysqlcol.New(e.URL)
	case "postgres", "postgresql":
		return pgcol.New(e.URL)
	case "redis":
		return rediscol.New(e.URL)
	case "mongodb":
		return mongocol.New(e.URL)
	case "sqlite":
		return sqlitecol.New(e.URL), nil
	default:
		return nil, fmt.Errorf("unknown driver %q (supported: mysql, mariadb, postgres, redis, mongodb, sqlite)", e.Driver)
	}
}

// ReadEntries reads the config file, returning an empty slice if it doesn't exist.
func ReadEntries(configPath string) ([]Entry, error) {
	data, err := os.ReadFile(configPath)
	if os.IsNotExist(err) {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}
	var entries []Entry
	if err := json.Unmarshal(data, &entries); err != nil {
		return nil, err
	}
	return entries, nil
}

// WriteEntries writes entries to the config file, creating the directory if needed.
func WriteEntries(configPath string, entries []Entry) error {
	if err := os.MkdirAll(filepath.Dir(configPath), 0755); err != nil {
		return err
	}
	data, err := json.MarshalIndent(entries, "", "  ")
	if err != nil {
		return err
	}
	return os.WriteFile(configPath, data, 0644)
}